Push checkpoint to HF Hub instead of download (Dexifried/tiny-router-checkpoint)

app.py

@@ -84,7 +84,7 @@ def train_and_export(encoder: str, epochs: int, batch_size: int, lr: float, max_
     logs.append(proc.stdout[-3000:] if proc.stdout else "")
     if proc.returncode != 0:
         logs.append(f"❌ Training failed:\n{proc.stderr[-2000:]}")
-        return "\n".join(logs), None, None
+        return "\n".join(logs), None
 
     logs.append("✅ Training complete!")
 
@@ -102,18 +102,25 @@ 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 (skip on HF — export locally after download) ─────
-    logs.append("\n📦 ONNX export skipped on HF (do locally after download).")
-
-    # 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)")
+    # ── Push checkpoint to HF Hub ──────────────────────────────────────
+    logs.append("\n📤 Pushing checkpoint to HF Hub...")
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        repo_id = "Dexifried/tiny-router-checkpoint"
+        api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
+        api.upload_folder(
+            folder_path=model_dir,
+            repo_id=repo_id,
+            repo_type="model",
+            commit_message=f"Tiny-router checkpoint (encoder={encoder}, epochs={epochs})",
+        )
+        logs.append(f"✅ Pushed to https://huggingface.co/{repo_id}")
+    except Exception as e:
+        logs.append(f"⚠️ Push failed: {e}")
 
     summary = json.dumps(eval_results, indent=2) if eval_results else "No eval results"
-    return "\n".join(logs), zip_file, summary
+    return "\n".join(logs), summary
 
 
 def quick_predict(text: str):
@@ -151,15 +158,14 @@ with gr.Blocks(title="⚡ Dex Neural Bake — Tiny-Router Trainer", theme=gr.the
             lr = gr.Number(label="Learning rate", value=DEFAULT_LR)
             max_length = gr.Slider(64, 512, value=DEFAULT_MAX_LENGTH, step=32, label="Max token length")
 
-        train_btn = gr.Button("🚀 Train on H200", variant="primary")
+        train_btn = gr.Button("🚀 Train & Push to HF", variant="primary")
         train_log = gr.Textbox(label="Training log", lines=25, max_lines=50, interactive=False)
         eval_summary = gr.JSON(label="Evaluation results")
-        model_file = gr.File(label="Download ONNX model")
 
         train_btn.click(
             fn=train_and_export,
             inputs=[encoder, epochs, batch_size, lr, max_length],
-            outputs=[train_log, model_file, eval_summary],
+            outputs=[train_log, eval_summary],
         )
 
     with gr.Tab("📋 Dataset Info"):

requirements.txt

@@ -9,3 +9,4 @@ torch>=2.2.0
 transformers>=4.40,<5
 tqdm>=4.66.0
 tiktoken>=0.12.0
+huggingface_hub>=0.22.0