save tokenizer to best checkpoint folder

SeanLee97 · Jun 29, 2024 · db997f2 · db997f2
1 parent e44442c
commit db997f2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/angle_emb/angle.py b/angle_emb/angle.py
@@ -1493,6 +1493,7 @@ def fit(self,
             best_ckpt_dir = None
             if output_dir is not None:
                 best_ckpt_dir = os.path.join(output_dir, 'best-checkpoint')
+                self.tokenizer.save_pretrained(best_ckpt_dir)
             evaluate_callback = EvaluateCallback(self.backbone, valid_ds,
                                                  partial(self.evaluate, batch_size=batch_size, device=self.device),
                                                  save_dir=best_ckpt_dir)