update trainer._save

rgdyy · May 22, 2024 · 005bbed · 005bbed
1 parent 53944b0
commit 005bbed
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/tevatron/retriever/trainer.py b/src/tevatron/retriever/trainer.py
@@ -30,6 +30,9 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
         if not isinstance(self.model, supported_classes):
             raise ValueError(f"Unsupported model class {self.model}")
         else:
+            prefix = 'encoder.'
+            assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
+            state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
             self.model.encoder.save_pretrained(
                 output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
             )