Update build_merged_submission.py to support compare-mode 3 (best=first/second/optimized)

medal-solvers/build_merged_submission.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python3
 """
-build_merged_submission.py — After running profile_best_of_two.py, use its output
-to create the merged submission zip picking best(first, second) per task.
+build_merged_submission.py — Build merged submission zip from profile_best_models.py output.
+
+Supports compare-mode 2 (best=first|second) and compare-mode 3 (best=first|second|optimized).
 
 Usage:
     python build_merged_submission.py \
-        --best-json /kaggle/working/best_per_task.json \
+        --best-json best_per_task.json \
         --first-zip ../submission-base.zip \
-        --second-zip ../submission-v81.zip \
+        --second-zip ../submission-v88.zip \
         --optimized-dir optimized \
         --output /kaggle/working/submission.zip
 """
@@ -19,7 +20,7 @@ import argparse
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--best-json', required=True, help='Output from profile_best_of_two.py')
+    parser.add_argument('--best-json', required=True, help='Output from profile_best_models.py')
     parser.add_argument('--first-zip', required=True, help='First submission zip')
     parser.add_argument('--second-zip', required=True, help='Second submission zip')
     parser.add_argument('--optimized-dir', required=True, help='optimized/ dir with models')
@@ -40,16 +41,27 @@ def main():
 
     from_first = 0
     from_second = 0
+    from_optimized = 0
 
     with zipfile.ZipFile(args.output, 'w', zipfile.ZIP_DEFLATED) as zf_out:
         for task_num in range(1, 401):
             fname = f'task{task_num:03d}.onnx'
             task_key = str(task_num)
-            
+
             best = results.get(task_key, {}).get('best', 'first')
 
+            # --- best = optimized: take from optimized dir ---
+            if best == 'optimized':
+                opt_path = os.path.join(args.optimized_dir, fname)
+                if os.path.exists(opt_path):
+                    zf_out.write(opt_path, fname)
+                    from_optimized += 1
+                    continue
+                # fallthrough to first if file missing
+                best = 'first'
+
+            # --- best = second: take from second zip ---
             if best == 'second':
-                # Use second model
                 try:
                     data = zf_second.read(fname)
                     zf_out.writestr(fname, data)
@@ -58,31 +70,28 @@ def main():
                 except:
                     pass  # fallthrough to first
 
-            # Use first model (optimized or from first zip)
-            if task_num in optimized_tasks:
-                opt_path = os.path.join(args.optimized_dir, fname)
-                if os.path.exists(opt_path):
-                    zf_out.write(opt_path, fname)
-                    from_first += 1
-                    continue
-
+            # --- best = first: take from first zip ---
             try:
                 data = zf_first.read(fname)
                 zf_out.writestr(fname, data)
                 from_first += 1
             except:
                 # Last resort: second
-                data = zf_second.read(fname)
-                zf_out.writestr(fname, data)
-                from_second += 1
+                try:
+                    data = zf_second.read(fname)
+                    zf_out.writestr(fname, data)
+                    from_second += 1
+                except:
+                    print(f"  WARNING: no model found for {fname}")
 
     zf_first.close()
     zf_second.close()
 
     print(f"Created {args.output}")
-    print(f"  From first: {from_first} models")
-    print(f"  From second: {from_second} models")
-    print(f"  Total: {from_first + from_second}")
+    print(f"  From first:     {from_first} models")
+    print(f"  From second:    {from_second} models")
+    print(f"  From optimized: {from_optimized} models")
+    print(f"  Total: {from_first + from_second + from_optimized}")
 
 
 if __name__ == '__main__':