diff --git a/folding/registries/miner_registry.py b/folding/registries/miner_registry.py index 876807ba..b766cdb5 100644 --- a/folding/registries/miner_registry.py +++ b/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] = {} diff --git a/folding/utils/ops.py b/folding/utils/ops.py index 9ec732ff..4b9be51c 100644 --- a/folding/utils/ops.py +++ b/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) diff --git a/folding/validators/reward.py b/folding/validators/reward.py index d483c742..c8675498 100644 --- a/folding/validators/reward.py +++ b/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