chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,40 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch.nn.functional as F
+
+from swift.utils import get_logger
+from .trainer import Trainer
+from .utils import gather_for_unpadded_tensors
+
+logger = get_logger()
+
+
+class EmbeddingTrainer(Trainer):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gather_function = gather_for_unpadded_tensors
+        mrl_dims = self.args.mrl_dims
+        if mrl_dims and self.compute_loss_func is not None:
+            origin_loss_func = self.compute_loss_func
+
+            def mrl_loss_func(outputs, labels, **kwargs):
+                # Matryoshka Representation Learning: compute loss on each truncated dimension
+                # and aggregate with the corresponding weights.
+                last_hidden_state = outputs['last_hidden_state']
+                loss = None
+                for dim, weight in mrl_dims.items():
+                    if dim > last_hidden_state.shape[-1]:
+                        logger.warning_once(f'MRL: skipping dimension {dim} because it exceeds the model hidden size '
+                                            f'({last_hidden_state.shape[-1]}).')
+                        continue
+                    sliced = F.normalize(last_hidden_state[..., :dim], p=2, dim=-1)
+                    cur_loss = weight * origin_loss_func({'last_hidden_state': sliced}, labels, **kwargs)
+                    loss = cur_loss if loss is None else loss + cur_loss
+                return loss
+
+            self.compute_loss_func = mrl_loss_func
+
+    def evaluation_loop(self, *args, **kwargs):
+        output = super().evaluation_loop(*args, **kwargs)
+        self.gather_function = gather_for_unpadded_tensors
+        return output