add save_model_eights method

CHANGELOG.md

@@ -252,3 +252,5 @@
 - Improve error handling when not a DatasetContainer is use in retrain and test API
 
 ## dev
+
+- Add `save_model_weights` method to `AddressParser` to save model weights (PyTorch state dictionary)

deepparse/parser/address_parser.py

@@ -13,6 +13,7 @@
 import platform
 import re
 import warnings
+from pathlib import Path
 from typing import List, Union, Dict, Tuple
 
 import torch
@@ -851,6 +852,35 @@ def test(
 
         return test_res
 
+    def save_model_weights(self, file_path: Union[str, Path]) -> None:
+        """
+        Method to save, in a Pickle format, the address parser model weights (PyTorch state dictionary).
+
+        file_path (Union[str, Path]): A complete file path with a pickle extension to save the model weights.
+            It can either be a string (e.g. 'path/to/save.p') or a path like path (e.g. Path('path/to/save.p').
+
+        Examples:
+
+            .. code-block:: python
+
+                address_parser = AddressParser(device=0)
+
+                a_path = Path('some/path/to/save.p')
+                address_parser.save_address_parser_weights(a_path)
+
+
+            .. code-block:: python
+
+                address_parser = AddressParser(device=0)
+
+                a_path = 'some/path/to/save.p'
+                address_parser.save_address_parser_weights(a_path)
+
+        """
+        self.model.state_dict()
+
+        torch.save(self.model.state_dict(), file_path)
+
     def _fill_tagged_addresses_components(
         self,
         tags_predictions: List,

major_release_todo.md

@@ -1 +1,2 @@
-- Remove deprecated `download_from_url` function
+- Remove deprecated `download_from_url` function
+- https://zenodo.org/account/settings/github/repository/GRAAL-Research/deepparse

tests/parser/integration/test_integration_address_parser.py

@@ -0,0 +1,80 @@
+# Bug with PyTorch source code makes torch.tensor as not callable for pylint.
+# pylint: disable=not-callable
+
+# Pylint error for TemporaryDirectory ask for with statement
+# pylint: disable=consider-using-with
+
+import os
+from collections import OrderedDict
+from os.path import exists
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest import skipIf
+
+import torch
+
+from tests.parser.integration.base_predict import (
+    AddressParserBase,
+)
+
+
+@skipIf(
+    not os.path.exists(os.path.join(os.path.expanduser("~"), ".cache", "deepparse", "cc.fr.300.bin")),
+    "download of model too long for test in runner",
+)
+class AddressParserTest(AddressParserBase):
+    @classmethod
+    def setUpClass(cls):
+        super(AddressParserTest, cls).setUpClass()
+
+        cls.temp_dir_obj = TemporaryDirectory()
+        cls.a_saving_dir_path = cls.temp_dir_obj.name
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.temp_dir_obj.cleanup()
+
+    def assert_file_exist(self, file_path):
+        file_exists = exists(file_path)
+        self.assertTrue(file_exists)
+
+    def setUp(self) -> None:
+        a_config = {"model_type": "fasttext", "device": "cpu", "verbose": False}
+        self.setup_model_with_config(a_config)
+
+    def test_givenAModelToExportDictStr_thenExportIt(self):
+        a_file_path = os.path.join(self.a_saving_dir_path, "exported_model.p")
+
+        self.a_model.save_model_weights(file_path=a_file_path)
+
+        self.assert_file_exist(a_file_path)
+
+    def test_givenAModelToExportDictPathALike_thenExportIt(self):
+        a_file_path = Path(os.path.join(self.a_saving_dir_path, "exported_model.p"))
+
+        self.a_model.save_model_weights(file_path=a_file_path)
+
+        self.assert_file_exist(a_file_path)
+
+    def test_givenAnExportedModelUsingTheMethod_whenReloadIt_thenReload(self):
+        a_file_path = Path(os.path.join(self.a_saving_dir_path, "exported_model.p"))
+
+        self.a_model.save_model_weights(file_path=a_file_path)
+
+        weights = torch.load(a_file_path)
+
+        self.assertIsInstance(weights, OrderedDict)
+
+        model_layer_keys = [
+            'encoder.lstm.weight_ih_l0',
+            'encoder.lstm.weight_hh_l0',
+            'encoder.lstm.bias_ih_l0',
+            'encoder.lstm.bias_hh_l0',
+            'decoder.lstm.weight_ih_l0',
+            'decoder.lstm.weight_hh_l0',
+            'decoder.lstm.bias_ih_l0',
+            'decoder.lstm.bias_hh_l0',
+            'decoder.linear.weight',
+            'decoder.linear.bias',
+        ]
+        self.assertEqual(model_layer_keys, list(weights.keys()))

tests/parser/test_address_parser.py

@@ -1,5 +1,5 @@
 # Since we use a patch as model mock we skip the unused argument error
-# pylint: disable=unused-argument, no-member, too-many-public-methods, too-many-lines
+# pylint: disable=unused-argument, no-member, too-many-public-methods, too-many-lines, too-many-arguments
 
 # Pylint error for TemporaryDirectory ask for with statement
 # pylint: disable=consider-using-with
@@ -76,9 +76,13 @@ def setUpClass(cls):
             "EOS",
         ]
 
+        cls.export_temp_dir_obj = TemporaryDirectory()
+        cls.a_saving_dir_path = cls.export_temp_dir_obj.name
+
     @classmethod
     def tearDownClass(cls) -> None:
         cls.temp_dir_obj.cleanup()
+        cls.export_temp_dir_obj.cleanup()
 
     def setUp(self):
         super().setUp()
@@ -1617,6 +1621,33 @@ def test_givenANewCacheDirFastText_thenInitWeightsInNewCacheDir(self, embeddings
             )
             download_weights_mock.assert_called_with(verbose=self.verbose, cache_dir=self.a_cache_dir)
 
+    @patch("deepparse.parser.address_parser.torch.save")
+    @patch("deepparse.parser.address_parser.FastTextSeq2SeqModel")
+    @patch("deepparse.parser.address_parser.fasttext_data_padding")
+    @patch("deepparse.parser.address_parser.FastTextVectorizer")
+    @patch("deepparse.parser.address_parser.FastTextEmbeddingsModel")
+    @patch("deepparse.parser.address_parser.download_fasttext_embeddings")
+    def test_givenAModelToExportDict_thenCallTorchSaveWithProperArgs(
+        self,
+        download_weights_mock,
+        embeddings_model_mock,
+        vectorizer_model_mock,
+        data_padding_mock,
+        model_mock,
+        torch_save_mock,
+    ):
+        address_parser = AddressParser(
+            model_type=self.a_fasttext_model_type,
+            device=self.a_cpu_device,
+            verbose=self.verbose,
+        )
+
+        a_file_path = os.path.join(self.a_saving_dir_path, "exported_model.p")
+        address_parser.save_model_weights(file_path=a_file_path)
+
+        torch_save_mock.assert_called()
+        torch_save_mock.assert_called_with(model_mock().state_dict(), a_file_path)
+
 
 if __name__ == "__main__":
     unittest.main()