fix-tuncate

modeling_xlm_roberta.py

@@ -600,7 +600,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
 
         truncate_dim = truncate_dim or self.config.truncate_dim
         if truncate_dim:
-            all_embeddings = self.truncate_embeddings(all_embeddings, truncate_dim)
+            all_embeddings = self.truncate_embeddings(all_embeddings, truncate_dim, normalize_embeddings)
 
         if convert_to_tensor:
             all_embeddings = torch.stack(all_embeddings)
@@ -613,19 +613,24 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         self.train(is_training)
         return all_embeddings
 
-    def truncate_embeddings(self, embeddings, truncate_dim):
-        if not self.config.matryoshka_dimensions:
-            logger.warning(
-                "Matryoshka embeddings are not supported, so dimension truncation will not be performed."
-            )
-            return embeddings
-        elif truncate_dim in self.config.matryoshka_dimensions:
-            return [tensor[:truncate_dim] for tensor in embeddings]
-        else:
-            raise ValueError(
-                f"The provided `truncate_dim` value of {truncate_dim} is not supported. "
-                f"Supported dimensions are {self.config.matryoshka_dimensions}."
-            )
+def truncate_embeddings(self, embeddings, truncate_dim, normalize_embeddings):
+    if not self.config.matryoshka_dimensions:
+        logger.warning(
+            "Matryoshka embeddings are not supported, so dimension truncation will not be performed."
+        )
+        return embeddings
+    elif truncate_dim in self.config.matryoshka_dimensions:
+        truncated_embeddings = [tensor[:truncate_dim] for tensor in embeddings]
+        if normalize_embeddings:
+            truncated_embeddings = [
+                torch.nn.functional.normalize(tensor, p=2, dim=0) for tensor in truncated_embeddings
+            ]
+        return truncated_embeddings
+    else:
+        raise ValueError(
+            f"The provided `truncate_dim` value of {truncate_dim} is not supported. "
+            f"Supported dimensions are {self.config.matryoshka_dimensions}."
+        )
 
     def mean_pooling(
         self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor