add fit test

5ky9uy · Apr 23, 2020 · bcd5247 · bcd5247
1 parent e8fc16c
commit bcd5247
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/moses/baselines/combinatorial.py b/moses/baselines/combinatorial.py
@@ -231,13 +231,9 @@ def get_connection_rule(self, fragment):
 
     @staticmethod
     def sample_fragment(counts):
-        try:
-            new_fragment = counts.sample(
-                weights=counts['frequency']
-            )
-        except:
-            print(counts)
-            raise
+        new_fragment = counts.sample(
+            weights=counts['frequency']
+        )
         new_fragment = dict(new_fragment.iloc[0])
         fragment = Chem.MolFromSmiles(new_fragment['fragment'])
         return fragment

diff --git a/tests/test_baselines.py b/tests/test_baselines.py
@@ -30,6 +30,11 @@ def test_hmm(self):
     def test_combinatorial(self):
         model = CombinatorialGenerator()
         model.fit(self.train[:100])
+        self.assertEqual(
+            model.fragment_counts.shape,
+            (156, 5),
+            "Model was not fitted properly"
+        )
         sample_original = model.generate_one(1)
         with tempfile.NamedTemporaryFile() as f:
             model.save(f.name)