Update utils.py

utils.py

@@ -59,10 +59,8 @@ class DataCollatorForCausalLMEval(object):
 
     def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
         
-        print(instances)
         srcs = instances[0]['src']
         task_type = instances[0]['task_type']
-        print(task_type)
 
         if task_type == 'retrosynthesis':
             src_start_str = self.product_start_str
@@ -78,7 +76,6 @@ class DataCollatorForCausalLMEval(object):
         data_dict = {
             'generation_prompts': generation_prompts
         }
-        print(data_dict)
         return data_dict
 
 def smart_tokenizer_and_embedding_resize(
@@ -131,7 +128,6 @@ class ReactionPredictionModel():
                     )
                 self.load_forward_model(candidate_models[model])
 
-        print(self.forward_model.device, self.retro_model.device)
         string_template_path = hf_hub_download(candidate_models[list(candidate_models.keys())[0]], filename="string_template.json", token = os.environ.get("TOKEN"))
         string_template = json.load(open(string_template_path, 'r'))
         reactant_start_str = string_template['REACTANTS_START_STRING']
@@ -220,8 +216,6 @@ class ReactionPredictionModel():
 
             if task_type == "retrosynthesis":
                 inputs = {k: v.to(self.retro_model.device) for k, v in inputs.items()}
-                print(inputs)
-                print(self.retro_model.device)
                 with torch.no_grad():
                     outputs = self.retro_model.generate(**inputs, max_length=512, num_return_sequences=10,
                                    do_sample=False, num_beams=10,
@@ -232,8 +226,6 @@ class ReactionPredictionModel():
                                    )
             else:
                 inputs = {k: v.to(self.forward_model.device) for k, v in inputs.items()}
-                print(inputs)
-                print(self.forward_model.device)
                 with torch.no_grad():
                     outputs = self.forward_model.generate(**inputs, max_length=512, num_return_sequences=10,
                                    do_sample=False, num_beams=10,
@@ -243,11 +235,9 @@ class ReactionPredictionModel():
                                    length_penalty=0.0,
                                    )
 
-            print(outputs)
             original_smiles_list = self.tokenizer.batch_decode(outputs.detach().cpu().numpy()[:, len(inputs['input_ids'][0]):],
                                                           skip_special_tokens=True) 
             original_smiles_list = map(lambda x: x.replace(" ", ""), original_smiles_list)
-            print(original_smiles_list)
             # canonize the SMILES
             canonized_smiles_list = []
             temp = []
@@ -262,7 +252,6 @@ class ReactionPredictionModel():
             predictions.append(canonized_smiles_list)
         
         rank, invalid_rate = compute_rank(predictions)
-        print(predictions, rank)
         return rank
     
     def predict_single_smiles(self, smiles, task_type):