save vocabulary file after training

Ard0rAI · Feb 18, 2020 · bfa93b9 · bfa93b9
1 parent 8585782
commit bfa93b9
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/spert/spert_trainer.py b/spert/spert_trainer.py
@@ -116,7 +116,7 @@ def train(self, train_path: str, valid_path: str, types_path: str, input_reader_
         # save final model
         extra = dict(epoch=args.epochs, updates_epoch=updates_epoch, epoch_iteration=0)
         global_iteration = args.epochs * updates_epoch
-        self._save_model(self._save_path, model, global_iteration,
+        self._save_model(self._save_path, model, self._tokenizer, global_iteration,
                          optimizer=optimizer if self.args.save_optimizer else None, extra=extra,
                          include_iteration=False, name='final_model')
 

diff --git a/spert/trainer.py b/spert/trainer.py
@@ -9,6 +9,7 @@
 from torch.nn import DataParallel
 from torch.optim import Optimizer
 from transformers import PreTrainedModel
+from transformers import PreTrainedTokenizer
 
 from spert import util
 from spert.opt import tensorboardX
@@ -91,16 +92,18 @@ def _log_csv(self, dataset_label: str, data_label: str, *data: Tuple[object]):
         logs = self._log_paths[dataset_label]
         util.append_csv(logs[data_label], *data)
 
-    def _save_best(self, model, optimizer, accuracy, iteration, label, extra=None):
+    def _save_best(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, optimizer: Optimizer,
+                   accuracy: float, iteration: int, label: str, extra=None):
         if accuracy > self._best_results[label]:
             self._logger.info("[%s] Best model in iteration %s: %s%% accuracy" % (label, iteration, accuracy))
-            self._save_model(self._save_path, model, iteration,
+            self._save_model(self._save_path, model, tokenizer, iteration,
                              optimizer=optimizer if self.args.save_optimizer else None,
                              save_as_best=True, name='model_%s' % label, extra=extra)
             self._best_results[label] = accuracy
 
-    def _save_model(self, save_path: str, model: PreTrainedModel, iteration: int, optimizer: Optimizer = None,
-                    save_as_best: bool = False, extra: dict = None, include_iteration: int = True, name: str = 'model'):
+    def _save_model(self, save_path: str, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
+                    iteration: int, optimizer: Optimizer = None, save_as_best: bool = False,
+                    extra: dict = None, include_iteration: int = True, name: str = 'model'):
         extra_state = dict(iteration=iteration)
 
         if optimizer:
@@ -117,10 +120,16 @@ def _save_model(self, save_path: str, model: PreTrainedModel, iteration: int, op
 
         util.create_directories_dir(dir_path)
 
+        # save model
         if isinstance(model, DataParallel):
             model.module.save_pretrained(dir_path)
         else:
             model.save_pretrained(dir_path)
+
+        # save vocabulary
+        tokenizer.save_pretrained(dir_path)
+
+        # save extra
         state_path = os.path.join(dir_path, 'extra.state')
         torch.save(extra_state, state_path)