fix(wrapper): Allow arbitrary model file names

microsoft · Jun 9, 2022 · 0373a19 · 0373a19
1 parent d240e8f
commit 0373a19
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/molecule_generation/wrapper.py b/molecule_generation/wrapper.py
@@ -45,11 +45,11 @@ def _get_model_file(cls, dir: Pathlike) -> pathlib.Path:
         Raises:
             ValueError, if the model pickle is not found or is not unique.
         """
-        # All candidate files must end with "_best.pkl"
-        candidates = list(pathlib.Path(dir).glob("*_best.pkl"))
+        # Candidate files must end with ".pkl"
+        candidates = list(pathlib.Path(dir).glob("*.pkl"))
         if len(candidates) != 1:
             raise ValueError(
-                f"There must be exactly one file matching the pattern. Found the following: {candidates}."
+                f"There must be exactly one *.pkl file. Found the following: {candidates}."
             )
         else:
             return candidates[0]