Support torch_dtype and CLS pooling (#6)

- support cls pooling and torch_dtype (dfe03e244d4edb62b846a3ed83842cc6c89e9129)

configuration_xlm_roberta.py

@@ -1,4 +1,5 @@
 from transformers import PretrainedConfig
+import torch
 
 class XLMRobertaFlashConfig(PretrainedConfig):
     def __init__(
@@ -22,6 +23,8 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             use_cache=True,
             classifier_dropout=None,
             use_flash_attn=True,
+            torch_dtype=None,
+            emb_pooler=None,
             **kwargs,
     ):
         super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
@@ -42,4 +45,9 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.use_cache = use_cache
         self.classifier_dropout = classifier_dropout
         self.use_flash_attn = use_flash_attn
-        
+        self.emb_pooler = emb_pooler
+        if torch_dtype and hasattr(torch, torch_dtype) and type(getattr(torch, torch_dtype)) is torch.dtype:
+            self.torch_dtype = getattr(torch, torch_dtype)
+        else:
+            self.torch_dtype = torch_dtype
+

modeling_xlm_roberta.py

@@ -395,6 +395,17 @@ class XLMRobertaPreTrainedModel(PreTrainedModel):
         if isinstance(module, XLMRobertaEncoder):
             module.gradient_checkpointing = value
 
+    @classmethod
+    def from_pretrained(
+        cls,
+        *args,
+        **kwargs,
+    ):
+        if not 'torch_dtype' in kwargs:
+            kwargs['torch_dtype'] = 'auto'
+        return super().from_pretrained(*args, **kwargs)
+
+
 
 class XLMRobertaModel(XLMRobertaPreTrainedModel):
     def __init__(self, config: XLMRobertaFlashConfig, add_pooling_layer=True):
@@ -545,9 +556,14 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             elif output_value is None:
                 raise NotImplementedError
             else:
-                embeddings = self.mean_pooling(
-                    token_embs, encoded_input['attention_mask']
-                )
+                if self.config.emb_pooler == 'cls':
+                    embeddings = self.cls_pooling(
+                        token_embs, encoded_input['attention_mask']
+                    )
+                else:
+                    embeddings = self.mean_pooling(
+                        token_embs, encoded_input['attention_mask']
+                    )
 
                 if normalize_embeddings:
                     embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
@@ -580,6 +596,12 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         )
 
 
+    def cls_pooling(
+        self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor
+    ):
+        return token_embeddings[:,0]
+
+
     def forward(
         self,
         input_ids,