Fixed a bug in save / load of a FireWord model.

Xeadriel · Mar 7, 2023 · 7161799 · 7161799
1 parent ef84a0c
commit 7161799
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/firelang/models/_fireword.py b/firelang/models/_fireword.py
@@ -248,7 +248,7 @@ def loss_skipgram(
         return loss
 
     @staticmethod
-    def from_pretrained(dirpath) -> FireWord:
+    def from_pretrained(dirpath, strict: bool = True) -> FireWord:
         dirpath = os.path.abspath(dirpath)
         if not os.path.exists(dirpath):
             raise FileNotFoundError(f"Directory not found at {dirpath}")
@@ -263,7 +263,7 @@ def from_pretrained(dirpath) -> FireWord:
         # state_dict
         word = FireWord(config=config, vocab=vocab)
         state_dict = torch.load(f"{dirpath}/pytorch_model.bin")
-        word.load_state_dict(state_dict)
+        word.load_state_dict(state_dict, strict=strict)
         return word
 
     def save(self, dirpath):
@@ -275,7 +275,7 @@ def save(self, dirpath):
 
         # config
         with open(f"{dirpath}/config.json", "wt") as f:
-            json.dump(self.config, f)
+            json.dump(self.config.__dict__, f)
 
         # vocab
         self.vocab.to_json(f"{dirpath}/vocab.json")