fix: skip ONNX export on HF, download checkpoint zip, export locally

app.py

@@ -102,38 +102,18 @@ def train_and_export(encoder: str, epochs: int, batch_size: int, lr: float, max_
             eval_results = json.load(f)
         logs.append(f"\n📈 Results: {json.dumps(eval_results, indent=2)}")
 
-    # ── Export ONNX ─────────────────────────────────────────────────────
-    logs.append("\n📦 Exporting to ONNX...")
-    onnx_dir = os.path.join(tmpdir, "onnx_export")
-    export_cmd = _build_export_cmd(model_dir, onnx_dir)
-    proc = subprocess.run(export_cmd, capture_output=True, text=True, cwd=str(Path(__file__).parent))
-    logs.append(proc.stdout[-1000:] if proc.stdout else "")
-    if proc.returncode != 0:
-        logs.append(f"❌ Export failed:\n{proc.stderr[-1000:]}")
-        return "\n".join(logs), None, None
-
-    onnx_path = os.path.join(onnx_dir, "tiny_router.onnx")
-    if not os.path.exists(onnx_path):
-        logs.append("❌ ONNX file not found after export")
-        return "\n".join(logs), None, None
-
-    size_mb = os.path.getsize(onnx_path) / (1024 * 1024)
-    logs.append(f"✅ ONNX exported! Size: {size_mb:.1f} MB")
-
-    # Also grab temperature scaling config
-    temp_path = os.path.join(model_dir, "temperature_scaling.json")
-    temp_dest = os.path.join(tmpdir, "temperature_scaling.json")
-    if os.path.exists(temp_path):
-        shutil.copy2(temp_path, temp_dest)
+    # ── Export ONNX (skip on HF — export locally after download) ─────
+    logs.append("\n📦 ONNX export skipped on HF (do locally after download).")
 
-    # Copy config
-    config_path = os.path.join(model_dir, "config.json")
-    config_dest = os.path.join(tmpdir, "config.json")
-    if os.path.exists(config_path):
-        shutil.copy2(config_path, config_dest)
+    # Zip the checkpoint for download
+    zip_path = os.path.join(tmpdir, "tiny-router-checkpoint")
+    shutil.make_archive(zip_path, "zip", model_dir)
+    zip_file = zip_path + ".zip"
+    size_mb = os.path.getsize(zip_file) / (1024 * 1024)
+    logs.append(f"✅ Checkpoint ready! ({size_mb:.1f} MB)")
 
     summary = json.dumps(eval_results, indent=2) if eval_results else "No eval results"
-    return "\n".join(logs), onnx_path, summary
+    return "\n".join(logs), zip_file, summary
 
 
 def quick_predict(text: str):