Update script.py

script.py

@@ -2,25 +2,30 @@
 
 ### You can change the rest of the code to define and test your solution.
 ### However, you should not change the signature of the provided function.
-### The script saves "submission.parquet" file in the current directory.
+### The script saves "submission.json" file in the current directory.
 ### You can use any additional files and subdirectories to organize your code.
 
 from pathlib import Path
 from tqdm import tqdm
-import pandas as pd
+import json
 import numpy as np
 from datasets import load_dataset
 from typing import Dict
-from joblib import Parallel, delayed
-import os
-import json
-import gc
 from hoho2025.example_solutions import predict_wireframe
-# check the https://github.com/s23dr/hoho2025/blob/main/hoho2025/example_solutions.py for the example solution
+from joblib import Parallel, delayed
 
-def empty_solution():
+def empty_solution(sample):
     '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
-    return np.zeros((2,3)), [(0, 1)]
+    return np.zeros((2,3)), [(0, 1)], sample['order_id']
+
+def predict_wireframe_safely(sample):
+    try:
+        pred_vertices, pred_edges = predict_wireframe(sample)
+    except Exception as e:
+        print (f"Failed due to {e},  returning empty solution")
+        pred_vertices, pred_edges = empty_solution(sample)
+    pred_edges  = [(int(a), int(b)) for a, b in pred_edges] # to remove possible np.int64
+    return pred_vertices, pred_edges, sample['order_id']
 
 class Sample(Dict):
     def pick_repr_data(self, x):
@@ -35,8 +40,10 @@ class Sample(Dict):
     def __repr__(self):
         # return str({k: v.shape if hasattr(v, 'shape') else [type(v[0])] if isinstance(v, list) else type(v) for k,v in self.items()})
         return str({k: self.pick_repr_data(v) for k,v in self.items()})
-    
 
+
+    
+import json
 if __name__ == "__main__":
     print ("------------ Loading dataset------------ ")
     param_path = Path('params.json')
@@ -92,54 +99,46 @@ if __name__ == "__main__":
         writer_batch_size=100
     )
 
+    # if TEST_ENV:
+    # dataset = load_dataset(
+    #     "webdataset", 
+    #     data_files=data_files,
+    #     trust_remote_code=True,
+    #     # streaming=True
+    # )
     print('load with webdataset')
-
-
+    # else:
+        
+    #     dataset = load_dataset(
+    #         "arrow", 
+    #         data_files=data_files,
+    #         trust_remote_code=True,
+    #         # streaming=True
+    #     )
+    #     print('load with arrow')
+    
+    
     print(dataset, flush=True)
+    # dataset = load_dataset('webdataset', data_files={)
     
     print('------------ Now you can do your solution ---------------')
     solution = []
-    
-    def convert(obj):
-        if isinstance(obj, np.ndarray):
-            return obj.tolist()
-        if isinstance(obj, np.integer):
-            return int(obj)
-        if isinstance(obj, np.floating):
-            return float(obj)
-        if isinstance(obj, dict):
-            return {k: convert(v) for k, v in obj.items()}
-        if isinstance(obj, (list, tuple)):
-            return [convert(i) for i in obj]
-        return obj
-
-    def process_sample(sample, i):
-        try:
-            pred_vertices, pred_edges = predict_wireframe(sample)
-        except:
-            pred_vertices, pred_edges = empty_solution()
-        if i %10 == 0:
-            gc.collect()
-        return {
-            'order_id': sample['order_id'], 
-            'wf_vertices': pred_vertices.tolist(),
-            'wf_edges': pred_edges
-        }
-    num_cores = 4
-    
-    for subset_name in dataset.keys():
-        print (f"Predicting {subset_name}")
-        for i, sample in enumerate(tqdm(dataset[subset_name])):
-            res = process_sample(sample, i)
-            solution.append(res)
-
+    for subset_name in dataset:
+        print (f"Predicitng on {subset_name}")
+        preds = Parallel(n_jobs=-1, prefer="processes")(
+            delayed(predict_wireframe_safely)(a) for a in tqdm(dataset[subset_name])
+        )
+        print ("Converting")
+        for p in preds:
+            pred_vertices, pred_edges, order_id = p
+            print (f'{order_id}: {len(pred_vertices)} verts, {len(pred_edges)} edges')
+            solution.append({
+                            'order_id': order_id, 
+                            'wf_vertices': pred_vertices.tolist(),
+                            'wf_edges': pred_edges
+                        })
     print('------------ Saving results ---------------')
-    # sub = pd.DataFrame(solution, columns=["order_id", "wf_vertices", "wf_edges"])
-    # sub.to_parquet("submission.parquet")
-
-    solution_as_lists = convert(solution)
-    
     with open("submission.json", "w") as f:
-        json.dump(solution_as_lists, f)
+        json.dump(solution, f)
 
     print("------------ Done ------------ ")