Update script.py

script.py

@@ -111,7 +111,7 @@ def save_submission(submission, path):
     submission (List[Dict[]]): The submission to save.
     path (str): The path to save the submission to.
     """
-    sub = pd.DataFrame(submission, columns=["__key__", "wf_vertices", "wf_edges", "edge_semantics"])
+    sub = pd.DataFrame(submission, columns=["__key__", "wf_vertices", "wf_edges"])
     sub.to_parquet(path)
     print(f"Submission saved to {path}")
 
@@ -130,7 +130,7 @@ if __name__ == "__main__":
             results.append(pool.submit(predict, sample, visualize=False))
         
         for i, result in enumerate(tqdm(results)):
-            key, pred_vertices, pred_edges, semantics = result.result()
+            key, pred_vertices, pred_edges = result.result()
             solution.append({
                             '__key__': key,
                             'wf_vertices': pred_vertices.tolist(),