Upload demo_train.py with huggingface_hub

demo_train.py

@@ -12,14 +12,30 @@
 from datasets import load_dataset
 from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
+import trackio
 import os
 
 print("🚀 Starting TRL + Trackio Demo")
 print("=" * 50)
 
-# Load a tiny dataset (just 50 examples for demo)
+# Initialize Trackio with Space sync for remote viewing
+# Trackio will auto-create the Space if it doesn't exist
+print("\n📊 Initializing Trackio...")
+trackio.init(
+    project="trl-demo",
+    space_id="evalstate/trl-trackio-dashboard",  # Auto-creates if needed!
+    config={
+        "model": "Qwen/Qwen2.5-0.5B",
+        "dataset": "trl-lib/Capybara",
+        "max_steps": 50,  # Longer for better visualization
+        "learning_rate": 2e-5,
+    }
+)
+print("✅ Trackio initialized! Dashboard: https://huggingface.co/spaces/evalstate/trl-trackio-dashboard")
+
+# Load a small dataset (200 examples for better visualization)
 print("\n📊 Loading dataset...")
-dataset = load_dataset("trl-lib/Capybara", split="train[:50]")
+dataset = load_dataset("trl-lib/Capybara", split="train[:200]")
 print(f"✅ Dataset loaded: {len(dataset)} examples")
 
 # Get username for hub push
@@ -33,12 +49,12 @@ config = SFTConfig(
     push_to_hub=True,
     hub_model_id=f"{username}/trl-trackio-demo",
 
-    # Quick demo settings
-    max_steps=10,  # Very short for demo
+    # Training settings (longer for better metrics)
+    max_steps=50,  # More steps for visualization
     per_device_train_batch_size=2,
 
-    # Logging
-    logging_steps=2,
+    # Logging (log frequently for real-time monitoring)
+    logging_steps=5,
 
     # Trackio monitoring - this is the key!
     report_to="trackio",
@@ -76,6 +92,10 @@ trainer.train()
 print("\n💾 Pushing to Hub...")
 trainer.push_to_hub()
 
+# Finish Trackio logging
+print("\n📊 Finalizing Trackio...")
+trackio.finish()
+
 print("\n✅ Demo complete!")
 print(f"📦 Model saved to: https://huggingface.co/{username}/trl-trackio-demo")
 print("📊 Check Trackio for training metrics and visualizations!")