print test error

YoshikaiY · Apr 23, 2020 · e8fc16c · e8fc16c
1 parent e2e4d3f
commit e8fc16c
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/moses/baselines/combinatorial.py b/moses/baselines/combinatorial.py
@@ -177,8 +177,7 @@ def generate_one(self, seed=None):
                 mol = self.sample_fragment(counts_masked)
             else:
                 if self.mode == 1:  # Choose connection atom first
-                    atom_mol = np.random.choice(connections_mol)
-                    connections_mol = [atom_mol]
+                    connections_mol = [np.random.choice(connections_mol)]
 
                 con_filter = self.get_connection_filter(connections_mol)
                 # Mask fragments with possible reactions
@@ -232,9 +231,13 @@ def get_connection_rule(self, fragment):
 
     @staticmethod
     def sample_fragment(counts):
-        new_fragment = counts.sample(
-            weights=counts['frequency']
-        )
+        try:
+            new_fragment = counts.sample(
+                weights=counts['frequency']
+            )
+        except:
+            print(counts)
+            raise
         new_fragment = dict(new_fragment.iloc[0])
         fragment = Chem.MolFromSmiles(new_fragment['fragment'])
         return fragment