Fix dlt.TranslationModel.load_obj (#34)

* Create test for saving/loading models

* Improve tests

* TranslationModel.save_obj now saves dlt config

* Bump version

dl_translate/_translation_model.py

@@ -1,4 +1,5 @@
 import os
+import json
 from typing import Union, List, Dict
 
 import transformers
@@ -252,6 +253,9 @@ def save_obj(self, path: str = "saved_model") -> None:
         torch.save(self._transformers_model, os.path.join(path, "weights.pt"))
         self._tokenizer.save_pretrained(path)
 
+        dlt_config = dict(model_family=self.model_family)
+        json.dump(dlt_config, open(os.path.join(path, "dlt_config.json"), "w"))
+
     @classmethod
     def load_obj(cls, path: str = "saved_model", **kwargs):
         """
@@ -261,5 +265,10 @@ def load_obj(cls, path: str = "saved_model", **kwargs):
         {{params}}
         {{path}} The directory where your torch model and tokenizer are stored
         """
-        load_dir = os.path.join(path, "weights.pt")
-        return cls(model_or_path=load_dir, tokenizer_path=path, **kwargs)
+        config_prev = json.load(open(os.path.join(path, "dlt_config.json"), "rb"))
+        config_prev.update(kwargs)
+        return cls(
+            model_or_path=os.path.join(path, "weights.pt"),
+            tokenizer_path=path,
+            **config_prev,
+        )

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="dl-translate",
-    version="0.2.0",
+    version="0.2.1",
     author="Xing Han Lu",
     author_email="github@xinghanlu.com",
     description="A deep learning-based translation library built on Huggingface transformers",

tests/long/test_save_load.py

@@ -0,0 +1,15 @@
+import os
+
+import dl_translate as dlt
+
+
+def test_save():
+    mt = dlt.TranslationModel()
+    mt.save_obj("saved_model")
+    assert os.path.exists("saved_model/weights.pt")
+    assert os.path.exists("saved_model/tokenizer_config.json")
+
+
+def test_load():
+    mt = dlt.TranslationModel.load_obj("saved_model")
+    assert isinstance(mt, dlt.TranslationModel)