slightly safer fix

Noble-Lab · Mar 13, 2024 · b10186f · b10186f
1 parent 6525292
commit b10186f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 18 deletions.
diff --git a/casanovo/casanovo.py b/casanovo/casanovo.py
@@ -403,7 +403,7 @@ def create_mgf_from_tide(
             # No need to do anything if the scan is not found in the scan map
             pass
     try:
-        mgf.write(all_spec, output_file)
+        mgf.write(all_spec, output_file, file_mode="w")
         logger.info("Annotated .mgf file written to %s.", output_file)
     except Exception as e:
         print(

diff --git a/casanovo/denovo/model.py b/casanovo/denovo/model.py
@@ -1043,24 +1043,31 @@ def on_predict_epoch_end(self, results) -> None:
                     )
                 )
             # Write rows
-            for group in results[0]:
-                for batch in group:
-                    for index, t_or_d, peptide, score, per_aa_scores in list(
-                        zip(*batch)
-                    ):
-                        # Remove scores of 0 (padding)
-                        per_aa_scores = per_aa_scores.numpy()
-                        per_aa_scores = list(per_aa_scores[per_aa_scores != 0])
-                        score = score.numpy()
-                        csv_writer.writerow(
-                            (
-                                index,
-                                peptide,
-                                bool(t_or_d),
-                                score,
-                                per_aa_scores,
+            for g1 in results:
+                for batch in g1:
+                    for pairs in batch:
+                        for (
+                            index,
+                            t_or_d,
+                            peptide,
+                            score,
+                            per_aa_scores,
+                        ) in list(zip(*pairs)):
+                            # Remove scores of 0 (padding)
+                            per_aa_scores = per_aa_scores.numpy()
+                            per_aa_scores = list(
+                                per_aa_scores[per_aa_scores != 0]
+                            )
+                            score = score.numpy()
+                            csv_writer.writerow(
+                                (
+                                    index,
+                                    peptide,
+                                    bool(t_or_d),
+                                    score,
+                                    per_aa_scores,
+                                )
                             )
-                        )
         out_f.close()