ct2rs 0.9.19

Rust bindings for OpenNMT/CTranslate2
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

#include "ctranslate2/models/model_factory.h"

#include <mutex>

#include "ctranslate2/models/whisper.h"
#include "ctranslate2/models/wav2vec2.h"
#include "ctranslate2/models/wav2vec2bert.h"
#include "ctranslate2/models/transformer.h"

namespace ctranslate2 {
  namespace models {

    static void register_supported_models() {
      // Empty spec name, TransformerBase, and TransformerBig are there for backward compatibility.
      register_model<TransformerModel>("", /*num_heads=*/8);
      register_model<TransformerModel>("TransformerBase", /*num_heads=*/8);
      register_model<TransformerModel>("TransformerBig", /*num_heads=*/16);
      register_model<TransformerModel>("TransformerSpec");

      register_model<TransformerDecoderModel>("TransformerDecoderSpec");

      register_model<TransformerEncoderModel>("TransformerEncoderSpec");

      register_model<WhisperModel>("WhisperSpec");

      register_model<Wav2Vec2Model>("Wav2Vec2Spec");

      register_model<Wav2Vec2BertModel>("Wav2Vec2BertSpec");
    }

    std::shared_ptr<Model> create_model(const std::string& name) {
      static std::once_flag init_flag;
      std::call_once(init_flag, register_supported_models);

      return ModelFactory::get_instance().create_model(name);
    }

  }
}