Misc/code cleanup

folding/registries/miner_registry.py

@@ -1,5 +1,5 @@
 from itertools import chain
-from typing import List, Callable
+from typing import List, Callable, Dict, Union
 
 import folding.utils.constants as c
 from folding.utils.ops import write_pkl, load_pkl
@@ -13,7 +13,7 @@ class MinerRegistry:
 
     def __init__(self, miner_uids: List[int]):
         self.tasks: List[str] = list(EVALUATION_REGISTRY.keys())
-        self.registry = dict.fromkeys(miner_uids)
+        self.registry: Dict[int, Dict[str: Union[int, float, List]]] = dict.fromkeys(miner_uids)
 
         for miner_uid in miner_uids:
             self.registry[miner_uid] = {}

folding/utils/ops.py

@@ -374,15 +374,13 @@ def save_files(files: Dict, output_directory: str, write_mode: str = "wb") -> Di
     logger.info(f"⏰ Saving files to {output_directory}...")
     check_if_directory_exists(output_directory=output_directory)
 
+    # loop over all of the output files and save to local disk
     filetypes = {}
     for filename, content in files.items():
         filetypes[filename.split(".")[-1]] = filename
 
         logger.info(f"Saving file {filename} to {output_directory}")
-        if "em.cpt" in filename:
-            filename = "em_binary.cpt"
 
-        # loop over all of the output files and save to local disk
         with open(os.path.join(output_directory, filename), write_mode) as f:
             f.write(content)
 

folding/validators/reward.py

@@ -112,6 +112,7 @@ def get_energies(
         protein (Protein): instance of the Protein class
         responses (List[JobSubmissionSynapse]): list of JobSubmissionSynapse objects
         uids (List[int]): list of uids
+        job_type (str): type of job ("SyntheticMD", "OrganicMD", "SyntheticML", "OrganicML")
 
     Returns:
         Tuple: Tuple containing the energies and the event dictionary