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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

#pragma once

#include <optional>
#include "ctranslate2/layers/attention.h"
#include "ctranslate2/layers/flash_attention.h"
#include "ctranslate2/layers/common.h"
#include "ctranslate2/layers/transformer.h"
#include "ctranslate2/padder.h"

namespace ctranslate2 {
  namespace layers {

    class EncoderLayer : public Layer {
    public:
      EncoderLayer(const models::Model& model,
                   const std::string& scope,
                   const bool pre_norm = true,
                   const ops::ActivationType activation_type = ops::ActivationType::ReLU,
                   const bool use_flash_attention = false);

      void operator()(const StorageView& input, StorageView& output) const;

      DataType output_type() const override {
        return _final_layer_norm.output_type();
      }

      dim_t output_size() const override {
        return _final_layer_norm.output_size();
      }

      const AttentionLayer& get_self_attention() const {
        return *_self_attention;
      }

    private:
      const dim_t _num_heads;
      const LayerNorm _ffn1_layer_norm;
      const FeedForwardNetwork _ff1;
      const LayerNorm _self_attn_layer_norm;
      std::unique_ptr<AttentionLayer> _self_attention;
      const ops::Transpose _transpose;
      const LayerNorm _layer_norm;
      const Conv1D _pconv1;
      const ops::Sigmoid _sigmoid;
      const Conv1D _dconv;
      const LayerNorm _dlayer_norm;
      const ops::Swish _swish;
      const Conv1D _pconv2;
      const LayerNorm _ffn2_layer_norm;
      const FeedForwardNetwork _ff2;
      const LayerNorm _final_layer_norm;
    };

    class AdapterLayer : public Layer {
    public:
      AdapterLayer(const models::Model& model,
                   const std::string& scope,
                   const bool pre_norm = true,
                   const ops::ActivationType activation_type = ops::ActivationType::ReLU,
                   const bool use_flash_attention = false);

      void operator()(const StorageView& input, StorageView& output) const;

      DataType output_type() const override {
        return _ffn.output_type();
      }

      dim_t output_size() const override {
        return _ffn.output_size();
      }

      const AttentionLayer& get_self_attention() const {
        return *_self_attention;
      }

    private:
      const dim_t _num_heads;
      const LayerNorm _residual_layer_norm;
      const ops::Transpose _transpose;
      const Conv1D _residual_conv;
      const ops::Sigmoid _sigmoid;
      const LayerNorm _attn_layer_norm;
      const Conv1D _attn_conv;
      std::unique_ptr<AttentionLayer> _self_attention;
      const LayerNorm _ffn_layer_norm;
      const FeedForwardNetwork _ffn;
    };

    class Wav2Vec2BertEncoder : public Layer {
    public:
      Wav2Vec2BertEncoder(const models::Model& model, const std::string& scope);

      void operator()(const StorageView& features, StorageView& output);

      DataType output_type() const override {
        if (_lm_head) {
          return (*_lm_head).output_type();
        }
        else {
          return DataType::FLOAT32;
        }
      }

      dim_t output_size() const override {
        if (_lm_head) {
          return (*_lm_head).output_size();
        }
        else {
          return 1024;
        }
      }

      dim_t input_size() const {
        return 1024;
      }

      bool is_encoded(const StorageView& features) const {
        // Input features shape: [batch_size, input_size, input_time]
        // Encoder output shape: [batch_size, input_time // 2, output_size]
        //
        // input_time is variable so we check that dimension 1 is different than its original value.

        return (features.rank() == 3
                && features.dim(2) == output_size()
                && features.dim(1) != input_size());
      }

    private:
      const StorageView* _return_logits;
      const LayerNorm _fp_layer_norm;
      const Dense _fp_projection;
      const std::vector<std::unique_ptr<const EncoderLayer>> _encoder_layers;
      const std::vector<std::unique_ptr<const AdapterLayer>> _adapt_layers;
      std::optional<Dense> _lm_head;
    };

  }
}