ct2rs 0.9.21

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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

#pragma once

#include "ctranslate2/layers/wav2vec2bert.h"
#include "ctranslate2/models/model.h"
#include "ctranslate2/replica_pool.h"

namespace ctranslate2 {
  namespace models {

    struct Wav2Vec2BertOptions {
      // Maximum generation length.
      size_t max_length = 448;

      // Randomly sample from the top K candidates (set 0 to sample from the full distribution).
      size_t sampling_topk = 1;

      // Maximum index of the first predicted timestamp.
      size_t max_initial_timestamp_index = 50;

      // Suppress blank outputs at the beginning of the sampling.
      bool suppress_blank = true;

      // List of token IDs to suppress.
      // -1 will suppress a default set of symbols as defined in the model config.json file.
      std::vector<int> suppress_tokens = {-1};
    };


    class Wav2Vec2BertModel : public Model {
    public:
      const Vocabulary& get_vocabulary() const;
      size_t current_spec_revision() const override;
      bool is_quantizable(const std::string& variable_name) const override;
      bool is_linear_weight(const std::string& variable_name) const override;
      std::unique_ptr<Model> clone() const override;

      bool use_global_int16_scale() const override {
        return false;
      }

    protected:
      void initialize(ModelReader& model_reader) override;
    private:
      std::shared_ptr<const Vocabulary> _vocabulary;
    };

    class Wav2Vec2BertReplica : public ModelReplica {
    public:
      static std::unique_ptr<Wav2Vec2BertReplica> create_from_model(const Model& model);

      Wav2Vec2BertReplica(const std::shared_ptr<const Wav2Vec2BertModel>& model);

      StorageView encode(StorageView features, const bool to_cpu);

    private:
      const std::shared_ptr<const Wav2Vec2BertModel> _model;
      const std::unique_ptr<layers::Wav2Vec2BertEncoder> _encoder;

      StorageView maybe_encode(StorageView features);
    };

    class Wav2Vec2Bert : public ReplicaPool<Wav2Vec2BertReplica> {
    public:
      using ReplicaPool::ReplicaPool;

      std::future<StorageView> encode(const StorageView& features, const bool to_cpu);

    };

  }
}