fixed minor issues

main.py

@@ -8,7 +8,7 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
 
   parser.add_argument("-d", "--device", type=str, default="cuda:0", help="Set the device (cpu or cuda:0)")
-  parser.add_argument("-o", "--output", type=str, default="./results/mutants", help="Set the path for the output directory")
+  parser.add_argument("-o", "--output", type=str, default="./results", help="Set the path for the output directory")
   parser.add_argument("-v", "--verbose", type=int, choices=[0, 1, 2], default=1, help="Set the verbosity between 0 and 2")
   parser.add_argument("--receptor", type=str, required=True, help="Set the receptor filepath")
   parser.add_argument("--ligand", type=str, required=True, help="Set the ligand filepath")
@@ -30,7 +30,7 @@ if __name__ == "__main__":
   flint.input(args.receptor, args.ligand)
 
   # begin the inference / generate mutants
-  flint.generate()
+  #flint.generate()
 
   # output the results and write the summary file
   flint.results()
\ No newline at end of file

model/Model.py

@@ -150,7 +150,7 @@ class Model:
         # well-predicted AA on total mask redisue
         # root mean squared deviation (RMSD)
         aa_ratio, rmsd, attend_logits = self.model.generate(
-          batch, output_folder=os.path.join(self.outputdir, f"batch_{b}")
+          batch, target_path=os.path.join(self.outputdir, f"batch_{b}")
         )
         
         shutil.copyfile(self.sources[0], os.path.join(self.outputdir, f"batch_{b}", f"{b}_orig_whole.pdb")) 
@@ -166,7 +166,7 @@ class Model:
     write results in a summary file, along with all generated PDBs.
     @return (Model): the instance of Model, for chainability purposes.
     """
-
+    
     # initialize the resulting summary TSV
     summary = "ID\tdelta_G\tKd\tmutations (AA)\n"
 
@@ -175,7 +175,7 @@ class Model:
 
     for b in range(self._nbatch()):
       for i in range(self.size):
-        receptor_path = os.path.join(self.outputdir, f"batch_{b}", f"{i}.pdb")
+        receptor_path = os.path.join(self.outputdir, f"batch_{b}", f"{i}_whole.pdb")
         ligand_path = os.path.join(self.outputdir, f"batch_{b}", f"{i}.sdf")
 
         # compute the docking window around ligand
@@ -224,4 +224,4 @@ class Model:
     """
 
     os.makedirs(self.outputdir, exist_ok=True)
-    return len(os.listdir(self.outputdir))
\ No newline at end of file
+    return len([f for f in os.listdir(self.outputdir) if os.path.isdir(os.path.join(self.outputdir, f))])
\ No newline at end of file