Upload folder using huggingface_hub

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d1f3a5b7e501b52e0d9d4ad11a9396cb2956e01e2972e4062c2dbb844c1419b7
-size 496547472
+oid sha256:273806a81277195257bb969cda50cc595f645d2528869b869f1a280dfc6f8368
+size 498910032

modeling_data2vec2.py

@@ -26,6 +26,7 @@ import math
 import warnings
 from typing import Optional, Tuple, Dict, List, Callable, Any
 from functools import partial
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -35,9 +36,7 @@ from torch import nn
 from torch import Tensor
 
 from transformers import PreTrainedModel
-from transformers.modeling_outputs import (
-    Wav2Vec2BaseModelOutput,
-)
+from transformers.utils import ModelOutput
 from .configuration_data2vec2 import (
     Data2Vec2MultiConfig,
     D2v2ModalityConfig,
@@ -59,6 +58,15 @@ from .utils_data2vec2 import (
 )
 
 
+@dataclass
+class Data2vec2BaseModelOutput(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None # output of the encoder-only model
+    pooler_output: Optional[torch.FloatTensor] = None # pooled output for text tasks, which is the first token representation followed by a dense layer and activation function
+    local_features: Optional[torch.FloatTensor] = None # features before the Transformer encoder
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None # TODO: only support manual implementation with fast=False in the forward pass of AltAttention as pytorch's dspa does not output attention weights
+
+
 #################################################
 ### modeling_data2vec2_base.py
 # copied from fairseq.modules.grad_multiply
@@ -1221,6 +1229,22 @@ class TextEncoder(ModalitySpecificEncoder):
 #################################################
 
 
+# copied from transformers.models.data2vec.modeling_data2vec.Data2VecTextPooler
+class Data2VecTextPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.embed_dim, config.embed_dim)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
 class Data2Vec2MultiPreTrainedModel(PreTrainedModel):
     # use init_bert_params from fairseq
     # copied from fairseq.modules.transformer_sentence_encoder.py
@@ -1281,7 +1305,9 @@ class Data2Vec2MultiModel(Data2Vec2MultiPreTrainedModel):
     config_class = Data2Vec2MultiConfig
     base_model_prefix = "data2vec2"
 
-    def __init__(self, config: Data2Vec2MultiConfig):
+    def __init__(
+        self, config: Data2Vec2MultiConfig, add_pooling_layer: bool = True
+    ):
         super().__init__(config)
         self.config = config
         modalities_cfg = config.modalities
@@ -1327,6 +1353,10 @@ class Data2Vec2MultiModel(Data2Vec2MultiPreTrainedModel):
 
         self.blocks = nn.ModuleList([make_block(dpr[i]) for i in range(config.depth)])
 
+        self.text_pooler = None
+        if add_pooling_layer and config.supported_modality == "TEXT":
+            self.text_pooler = Data2VecTextPooler(config)
+
         self.norm = None
         if config.layer_norm_first:
             self.norm = make_layer_norm(config.embed_dim)
@@ -1355,6 +1385,9 @@ class Data2Vec2MultiModel(Data2Vec2MultiPreTrainedModel):
         """
         for mod in self.modalities:
             self.modality_encoders[mod]._freeze_parameters()
+        for block in self.blocks:
+            for p in block.parameters():
+                p.requires_grad = False
 
     def make_modality_encoder(
         self,
@@ -1405,7 +1438,7 @@ class Data2Vec2MultiModel(Data2Vec2MultiPreTrainedModel):
             precomputed_mask=None,
         )
         x = extractor_out["x"]
-        extract_features = x
+        local_features = x
 
         # encoder_mask = extractor_out["encoder_mask"]
         masked_padding_mask = extractor_out["padding_mask"]
@@ -1447,20 +1480,26 @@ class Data2Vec2MultiModel(Data2Vec2MultiPreTrainedModel):
                 :, feature_extractor.modality_cfg.num_extra_tokens :
             ]
 
+        txt_pooled_output = (
+            self.text_pooler(x) if self.text_pooler is not None else None
+        )
+
         if not return_dict:
             return tuple(
                 v
                 for v in [
                     x,
-                    extract_features,
+                    txt_pooled_output,
+                    local_features,
                     layer_results,
                 ]
                 if v is not None
             )
         
-        return Wav2Vec2BaseModelOutput(
+        return Data2vec2BaseModelOutput(
             last_hidden_state=x,
-            extract_features=extract_features,
+            pooler_output=txt_pooled_output,
+            local_features=local_features,
             hidden_states=layer_results if output_hidden_states else None,
             attentions=None, # switch to manual implementation with fast=False in forward pass of AltAttention as pytorch's dspa does not output attention weights
         )