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

#pragma once

#include "ctranslate2/layers/transformer.h"

namespace ctranslate2 {
  namespace layers {

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

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

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

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

      dim_t max_output_time() const {
        return _position_embedding.num_positions();
      }

      dim_t input_size() const {
        return _conv1.input_size();
      }

      dim_t max_input_time() const {
        return max_output_time() * 2;
      }

      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 Conv1D _conv1;
      const Conv1D _conv2;
      const ops::Transpose _transpose;
      PositionEmbedding _position_embedding;
      const dim_t _num_heads;
      const std::vector<std::unique_ptr<const TransformerEncoderLayer>> _layers;
      const LayerNorm _output_norm;
    };

    class WhisperDecoder : public TransformerDecoder {
    public:
      using TransformerDecoder::TransformerDecoder;

      bool return_normalized_attention() const override {
        return false;
      }

      void forward_prompt(const StorageView& prompt,
                          DecoderState& state,
                          StorageView* outputs = nullptr);

      void compute_logits_for_steps(const StorageView& outputs,
                                    const StorageView& steps,
                                    StorageView& logits);
    };

  }
}