Support for returning the hidden vector in the Wav2Vec2 and Wav2Vec2Bert models (#1867)

* adding the hidden vector return feature

* format update

* format update

* dummy

* dummy

---------

Co-authored-by: hkwon <[email protected]>

Author: homink

include/ctranslate2/layers/wav2vec2.h

@@ -59,11 +59,21 @@ namespace ctranslate2 {
       void operator()(const StorageView& features, StorageView& output);
 
       DataType output_type() const override {
-        return _output_norm.output_type();
+        if (_lm_head) {
+          return (*_lm_head).output_type();
+        }
+        else {
+          return _output_norm.output_type();
+        }
       }
 
       dim_t output_size() const override {
-        return _output_norm.output_size();
+        if (_lm_head) {
+          return (*_lm_head).output_size();
+        }
+        else {
+          return _output_norm.output_size();
+        }
       }
 
       dim_t input_size() const {
@@ -81,8 +91,10 @@ namespace ctranslate2 {
                 && features.dim(1) != input_size());
       }
 
-    private:
       const StorageView* _upgraded_model;
+
+    private:
+      const StorageView* _return_logits;
       std::optional<Wav2Vec2LayerNormConvLayer> _feat_layer0;
       std::optional<std::vector<std::unique_ptr<const Wav2Vec2LayerNormConvLayer>>> _feat_layers;
       std::optional<LayerNorm> _fp_norm;

include/ctranslate2/layers/wav2vec2bert.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <optional>
 #include "ctranslate2/layers/attention.h"
 #include "ctranslate2/layers/flash_attention.h"
 #include "ctranslate2/layers/common.h"
@@ -92,11 +93,21 @@ namespace ctranslate2 {
       void operator()(const StorageView& features, StorageView& output);
 
       DataType output_type() const override {
-        return _lm_head.output_type();
+        if (_lm_head) {
+          return (*_lm_head).output_type();
+        }
+        else {
+          return DataType::FLOAT32;
+        }
       }
 
       dim_t output_size() const override {
-        return _lm_head.output_size();
+        if (_lm_head) {
+          return (*_lm_head).output_size();
+        }
+        else {
+          return 1024;
+        }
       }
 
       dim_t input_size() const {
@@ -115,11 +126,12 @@ namespace ctranslate2 {
       }
 
     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;
-      const Dense _lm_head;
+      std::optional<Dense> _lm_head;
     };
 
   }

python/ctranslate2/converters/transformers.py

@@ -1004,10 +1004,13 @@ def architecture_name(self):
         return "Wav2Vec2ForCTC"
 
     def get_model_spec(self, model):
+        return_hidden = getattr(model.wav2vec2.config, "return_hidden", False)
         spec = wav2vec2_spec.Wav2Vec2Spec(
             model.wav2vec2.config.num_feat_extract_layers,
             model.wav2vec2.encoder.config.num_hidden_layers,
             model.wav2vec2.encoder.config.num_attention_heads,
+            model.lm_head.weight.shape[0],
+            return_hidden,
         )
 
         # layer component name matching (no duplications saving)
@@ -1065,7 +1068,9 @@ def set_encoder(self, spec, model, config):
         self.set_feature_projection(spec, model.wav2vec2.feature_projection)
         self.set_pos_conv_embed(spec, model.wav2vec2.encoder, config)
         super().set_encoder(spec, model.wav2vec2.encoder)
-        self.set_linear(spec.lm_head, model.lm_head)
+        return_hidden = getattr(model.wav2vec2.config, "return_hidden", False)
+        if not return_hidden:
+            self.set_linear(spec.lm_head, model.lm_head)
 
     def set_common_layers(self, spec, module):
         self.set_layer_norm(spec.layer_norm, module.layer_norm)
@@ -1078,9 +1083,12 @@ def architecture_name(self):
         return "Wav2Vec2BertForCTC"
 
     def get_model_spec(self, model):
+        return_hidden = getattr(model.wav2vec2_bert.config, "return_hidden", False)
         spec = wav2vec2bert_spec.Wav2Vec2BertSpec(
             model.wav2vec2_bert.config.num_adapter_layers,
             model.wav2vec2_bert.config.num_hidden_layers,
+            model.lm_head.weight.shape[0],
+            return_hidden,
         )
         self.set_encoder(spec.encoder, model)
         return spec
@@ -1170,7 +1178,9 @@ def set_encoder(self, spec, model):
         self.set_wav2vec2bert_adapter(
             spec.adapter_layers, model.wav2vec2_bert.adapter.layers
         )
-        self.set_linear(spec.lm_head, model.lm_head)
+        return_hidden = getattr(model.wav2vec2_bert.config, "return_hidden", False)
+        if not return_hidden:
+            self.set_linear(spec.lm_head, model.lm_head)
 
     def set_conv1d(self, spec, module):
         spec.weight = module.weight

python/ctranslate2/specs/wav2vec2_spec.py

@@ -13,9 +13,22 @@ def __init__(self):
 
 
 class Wav2Vec2Spec(model_spec.LanguageModelSpec):
-    def __init__(self, feat_layers, num_layers, num_heads):
+    def __init__(
+        self,
+        feat_layers,
+        num_layers,
+        num_heads,
+        vocab_size,
+        return_hidden,
+    ):
         super().__init__()
-        self.encoder = Wav2Vec2EncoderSpec(feat_layers, num_layers, num_heads)
+        self.vocab_size = np.dtype("int16").type(vocab_size)
+        self.encoder = Wav2Vec2EncoderSpec(
+            feat_layers,
+            num_layers,
+            num_heads,
+            return_hidden,
+        )
 
     @property
     def name(self):
@@ -29,7 +42,7 @@ def get_default_config(self):
         return Wav2Vec2Config()
 
     def get_vocabulary_size(self):
-        return self.encoder.lm_head.weight.shape[0]
+        return int(self.vocab_size.numpy())
 
 
 class Wav2Vec2LayerNormConvLayer(model_spec.LayerSpec):
@@ -44,7 +57,7 @@ def __init__(self):
 
 
 class Wav2Vec2EncoderSpec(model_spec.LayerSpec):
-    def __init__(self, feat_layers, num_layers, num_heads):
+    def __init__(self, feat_layers, num_layers, num_heads, return_hidden):
         self.num_heads = np.dtype("int16").type(num_heads)
         self.feat_layer0 = Wav2Vec2LayerNormConvLayer()
         self.feat_layer = [Wav2Vec2LayerNormConvLayer() for i in range(feat_layers - 1)]
@@ -55,4 +68,5 @@ def __init__(self, feat_layers, num_layers, num_heads):
         self.layer = [
             transformer_spec.TransformerEncoderLayerSpec() for _ in range(num_layers)
         ]
-        self.lm_head = common_spec.LinearSpec()
+        if not return_hidden:
+            self.lm_head = common_spec.LinearSpec()

python/ctranslate2/specs/wav2vec2bert_spec.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from ctranslate2.specs import attention_spec, common_spec, model_spec
 
 
@@ -9,11 +11,19 @@ def __init__(self):
 
 
 class Wav2Vec2BertSpec(model_spec.LanguageModelSpec):
-    def __init__(self, num_hidden_layers, num_adapter_layers):
+    def __init__(
+        self,
+        num_hidden_layers,
+        num_adapter_layers,
+        vocab_size,
+        return_hidden,
+    ):
         super().__init__()
+        self.vocab_size = np.dtype("int16").type(vocab_size)
         self.encoder = Wav2Vec2BertEncoderSpec(
             num_adapter_layers,
             num_hidden_layers,
+            return_hidden,
         )
 
     @property
@@ -28,7 +38,7 @@ def get_default_config(self):
         return Wav2Vec2BertConfig()
 
     def get_vocabulary_size(self):
-        return self.encoder.lm_head.weight.shape[0]
+        return int(self.vocab_size.numpy())
 
 
 class Wav2Vec2BertFeedForwardSpec(model_spec.LayerSpec):
@@ -78,9 +88,10 @@ def __init__(self):
 
 
 class Wav2Vec2BertEncoderSpec(model_spec.LayerSpec):
-    def __init__(self, num_hidden_layers, num_adapter_layers):
+    def __init__(self, num_hidden_layers, num_adapter_layers, return_hidden):
         self.fp_layer_norm = common_spec.LayerNormSpec()
         self.fp_projection = common_spec.LinearSpec()
         self.encoder_layers = [EncoderSpec() for _ in range(num_hidden_layers)]
         self.adapter_layers = [AdapterSpec() for _ in range(num_adapter_layers)]
-        self.lm_head = common_spec.LinearSpec()
+        if not return_hidden:
+            self.lm_head = common_spec.LinearSpec()

src/layers/wav2vec2.cc

@@ -46,7 +46,8 @@ namespace ctranslate2 {
     }
 
     Wav2Vec2Encoder::Wav2Vec2Encoder(const models::Model& model, const std::string& scope)
-      : _upgraded_model(model.get_variable_if_exists(scope + "/lm_head/weight"))
+      : _return_logits(model.get_variable_if_exists(scope + "/lm_head/weight"))
+      , _upgraded_model(model.get_variable_if_exists(scope + "/fp_projection/weight"))
       , _num_heads(model.get_attribute_with_default<int32_t>(scope + "/num_heads", 8))
       , _transpose({0, 2, 1})
       , _layers(build_layers_list<const TransformerEncoderLayer>(model,
@@ -65,7 +66,9 @@ namespace ctranslate2 {
         _fp_norm.emplace(model, scope + "/fp_layer_norm");
         _fp_ff.emplace(model, scope + "/fp_projection", nullptr, true);
         _pos_conv_embed.emplace(model, scope + "/pos_conv_embed");
-        _lm_head.emplace(model, scope + "/lm_head", nullptr, true);
+        if (_return_logits) {
+          _lm_head.emplace(model, scope + "/lm_head", nullptr, true);
+        }
       }
     }
 
@@ -101,12 +104,16 @@ namespace ctranslate2 {
           (*layer)(feat_buffer2, nullptr, feat_buffer);
           feat_buffer2 = std::move(feat_buffer);
         }
-        _output_norm(feat_buffer2, feat_buffer);
-
-        (*_lm_head)(feat_buffer, output); //_lm_head(feat_buffer, output);
+        if (_return_logits) {
+          _output_norm(feat_buffer2, feat_buffer);
+          (*_lm_head)(feat_buffer, output);
+        }
+        else {
+          _output_norm(feat_buffer2, output);
+        }
       }
       else { // backward compatibility for the previous converted model
-        StorageView input(output_type(), features.device());
+        StorageView input(features.dtype(), features.device());
         input = features;
         for (const auto& layer : _layers) {
           (*layer)(input, nullptr, output);

src/layers/wav2vec2bert.cc

@@ -164,7 +164,8 @@ namespace ctranslate2 {
     }
 
     Wav2Vec2BertEncoder::Wav2Vec2BertEncoder(const models::Model& model, const std::string& scope)
-      : _fp_layer_norm(model, scope + "/fp_layer_norm")
+      : _return_logits(model.get_variable_if_exists(scope + "/lm_head/weight"))
+      , _fp_layer_norm(model, scope + "/fp_layer_norm")
       , _fp_projection(model, scope + "/fp_projection", nullptr, true)
       , _encoder_layers(build_layers_list<const EncoderLayer>(model,
                                                               scope + "/encoder_layers",
@@ -175,8 +176,10 @@ namespace ctranslate2 {
                                                             scope + "/adapter_layers",
                                                             /*pre_norm=*/true,
                                                             ops::ActivationType::ReLU,
-                                                            /*use_flash_attention=*/false))
-      , _lm_head(model, scope + "/lm_head", nullptr, true) {
+                                                            /*use_flash_attention=*/false)) {
+      if (_return_logits) {
+        _lm_head.emplace(model, scope + "/lm_head", nullptr, true);
+      }
     }
 
     void Wav2Vec2BertEncoder::operator()(const StorageView& features, StorageView& output) {
@@ -203,7 +206,12 @@ namespace ctranslate2 {
         buffer2 = std::move(buffer1);
       }
 
-      _lm_head(buffer2, output);
+      if (_return_logits) {
+        (*_lm_head)(buffer2, output);
+      }
+      else {
+        output = std::move(buffer2);
+      }
     }
 
   }

src/models/wav2vec2.cc

@@ -78,7 +78,12 @@ namespace ctranslate2 {
       features.move_to(device, dtype);
 
       StorageView encoder_output(dtype, device);
-      (*_encoder)(features, encoder_output);
+      if (_encoder->_upgraded_model) {
+        encoder_output = maybe_encode(std::move(features));
+      }
+      else {
+        (*_encoder)(features, encoder_output);
+      }
 
       if (to_cpu) {
         if (device != Device::CPU)

src/models/wav2vec2bert.cc

@@ -77,8 +77,7 @@ namespace ctranslate2 {
       const DataType dtype = _encoder->output_type();
       features.move_to(device, dtype);
 
-      StorageView encoder_output(dtype, device);
-      (*_encoder)(features, encoder_output);
+      StorageView encoder_output = maybe_encode(std::move(features));
 
       if (to_cpu) {
         if (device != Device::CPU)