Update train_on_hf_spaces.py: auto-start training, file-based logging

train_on_hf_spaces.py

@@ -5,13 +5,8 @@ HuggingFace Spaces Training Script for EeshaAI/zeeb
 Runs on HuggingFace Spaces (free CPU tier, 16GB RAM).
 Fine-tunes OLMo 2 1B Instruct with LoRA to generate video tokens.
 
-Steps:
-1. Load OLMo 2 1B Instruct (full HuggingFace model, fp32)
-2. Expand vocabulary with visual tokens (<v_0> ... <v_1023>)
-3. Apply LoRA (r=4, alpha=8) to q_proj and v_proj
-4. Train on tokenized video data (3 epochs)
-5. Merge LoRA weights back into base model
-6. Push merged model to EeshaAI/zeeb
+Writes all logs to a file for the Gradio UI to read.
+Auto-pushes the trained model to EeshaAI/zeeb when done.
 """
 
 import os
@@ -20,7 +15,7 @@ import json
 import time
 import traceback
 import gc
-from typing import Generator
+import threading
 
 import torch
 from torch.utils.data import DataLoader, Dataset
@@ -46,6 +41,25 @@ GRADIENT_ACCUMULATION = 4
 MAX_GRAD_NORM = 1.0
 LOG_EVERY = 1
 
+
+class _Logger:
+    """Thread-safe logger that writes to both stdout and a log file."""
+    def __init__(self, log_path):
+        self.log_path = log_path
+        self.lock = threading.Lock()
+        # Initialize log file
+        with open(log_path, "w") as f:
+            f.write("🚀 Zeeb Training Pipeline Starting...\n\n")
+
+    def log(self, msg):
+        with self.lock:
+            with open(self.log_path, "a") as f:
+                f.write(msg)
+                f.flush()
+            # Also print to stdout for HF Spaces logs
+            print(msg, end="", flush=True)
+
+
 # ---------------------------------------------------------------------------
 # Dataset
 # ---------------------------------------------------------------------------
@@ -56,7 +70,6 @@ class VideoTokenDataset(Dataset):
         with open(data_path) as f:
             self.data = json.load(f)
         self.max_tokens = max_tokens
-        print(f"[Dataset] Loaded {len(self.data)} samples from {data_path}")
 
     def __len__(self):
         return len(self.data)
@@ -74,92 +87,79 @@ class VideoTokenDataset(Dataset):
 
 
 # ---------------------------------------------------------------------------
-# Training
+# Training (file-based logging)
 # ---------------------------------------------------------------------------
-def train(data_path: str = "tokenized_dataset.json") -> Generator[str, None, None]:
-    """
-    Main training loop. Yields log messages for the Gradio UI.
-    """
-    yield "🚀 Starting training pipeline...\n"
-
-    # ── 1. Load tokenizer & model ──────────────────────────────────────────
-    yield "📦 Loading OLMo 2 1B Instruct tokenizer...\n"
+def run_training_to_file(log_path: str = "/tmp/training_log.txt"):
+    """Run the full training pipeline, logging to a file."""
+    logger = _Logger(log_path)
 
     try:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-    except ImportError as e:
-        yield f"❌ transformers not installed: {e}\n"
-        raise
-
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        yield f"✅ Tokenizer loaded. Vocab size: {len(tokenizer)}\n"
+        _run_training(logger)
     except Exception as e:
-        yield f"❌ Failed to load tokenizer: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
+        logger.log(f"\n❌ FATAL ERROR: {e}\n")
+        logger.log(traceback.format_exc() + "\n")
 
-    yield "📦 Loading model in float32 on CPU (this takes ~2-3 min)...\n"
-    try:
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            trust_remote_code=True,
-            torch_dtype=torch.float32,
-        )
-        yield f"✅ Model loaded. Parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M\n"
-    except Exception as e:
-        yield f"❌ Failed to load model: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
 
-    # ── 2. Expand vocabulary ───────────────────────────────────────────────
-    yield f"🔤 Adding {CODEBOOK_SIZE} visual tokens + special tokens...\n"
+def _run_training(logger: _Logger):
+    """Core training logic."""
+
+    # ── 1. Load tokenizer ──────────────────────────────────────────────────
+    logger.log("📦 Loading OLMo 2 1B Instruct tokenizer...\n")
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    logger.log(f"✅ Tokenizer loaded. Vocab size: {len(tokenizer)}\n")
+
+    # ── 2. Load model ───────────────────────────────────────────────���──────
+    logger.log("📦 Loading model in float32 on CPU (this takes ~2-3 min)...\n")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        trust_remote_code=True,
+        torch_dtype=torch.float32,
+    )
+    n_params = sum(p.numel() for p in model.parameters()) / 1e6
+    logger.log(f"✅ Model loaded. Parameters: {n_params:.1f}M\n")
+
+    # ── 3. Expand vocabulary ───────────────────────────────────────────────
+    logger.log(f"🔤 Adding {CODEBOOK_SIZE} visual tokens + special tokens...\n")
     visual_tokens = [VIDEO_START, VIDEO_END, VIDEO_PAD]
     for i in range(CODEBOOK_SIZE):
         visual_tokens.append(f"<v_{i}>")
 
     num_added = tokenizer.add_tokens(visual_tokens)
     model.resize_token_embeddings(len(tokenizer))
-    yield f"✅ Added {num_added} tokens. New vocab size: {len(tokenizer)}\n"
-
-    # ── 3. Apply LoRA ─────────────────────────────────────────────────────
-    yield f"🔧 Applying LoRA (r={LORA_R}, alpha={LORA_ALPHA})...\n"
-    try:
-        from peft import LoraConfig, get_peft_model, TaskType
-
-        lora_config = LoraConfig(
-            r=LORA_R,
-            lora_alpha=LORA_ALPHA,
-            target_modules=["q_proj", "v_proj"],
-            lora_dropout=LORA_DROPOUT,
-            bias="none",
-            task_type=TaskType.CAUSAL_LM,
-        )
-        model = get_peft_model(model, lora_config)
-        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        total = sum(p.numel() for p in model.parameters())
-        yield f"✅ LoRA applied. Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)\n"
-    except Exception as e:
-        yield f"❌ Failed to apply LoRA: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
-
-    # ── 4. Load dataset ───────────────────────────────────────────────────
-    yield f"📊 Loading dataset from {data_path}...\n"
-    try:
-        dataset = VideoTokenDataset(data_path, max_tokens=256)
-        dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
-        total_steps = NUM_EPOCHS * len(dataloader)
-        yield f"📊 {len(dataset)} samples × {NUM_EPOCHS} epochs = {total_steps} steps\n"
-    except Exception as e:
-        yield f"❌ Failed to load dataset: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
-
-    # ── 5. Train ──────────────────────────────────────────────────────────
-    yield "🔥 Starting training loop...\n\n"
+    logger.log(f"✅ Added {num_added} tokens. New vocab size: {len(tokenizer)}\n")
+
+    # ── 4. Apply LoRA ─────────────────────────────────────────────────────
+    logger.log(f"🔧 Applying LoRA (r={LORA_R}, alpha={LORA_ALPHA})...\n")
+    from peft import LoraConfig, get_peft_model, TaskType
+
+    lora_config = LoraConfig(
+        r=LORA_R,
+        lora_alpha=LORA_ALPHA,
+        target_modules=["q_proj", "v_proj"],
+        lora_dropout=LORA_DROPOUT,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+    )
+    model = get_peft_model(model, lora_config)
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    logger.log(f"✅ LoRA applied. Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)\n")
+
+    # ── 5. Load dataset ───────────────────────────────────────────────────
+    data_path = "tokenized_dataset.json"
+    logger.log(f"📊 Loading dataset from {data_path}...\n")
+    dataset = VideoTokenDataset(data_path, max_tokens=256)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+    total_steps = NUM_EPOCHS * len(dataloader)
+    logger.log(f"📊 {len(dataset)} samples × {NUM_EPOCHS} epochs = {total_steps} steps\n")
+
+    # ── 6. Train ──────────────────────────────────────────────────────────
+    logger.log("🔥 Starting training loop...\n\n")
 
     optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
     model.train()
@@ -168,123 +168,139 @@ def train(data_path: str = "tokenized_dataset.json") -> Generator[str, None, Non
     running_loss = 0.0
     start_time = time.time()
 
-    try:
-        for epoch in range(NUM_EPOCHS):
-            epoch_loss = 0.0
-            num_batches = 0
-
-            for batch_idx, batch in enumerate(dataloader):
-                prompt = batch["prompt"][0]
-                video_tokens = batch["video_tokens"][0]
-
-                # Format training text
-                token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens[:64])  # limit tokens for memory
-                text = f"Create a video of: {prompt} {VIDEO_START} {token_str} {VIDEO_END}"
-
-                inputs = tokenizer(
-                    text,
-                    return_tensors="pt",
-                    truncation=True,
-                    max_length=MAX_SEQ_LEN,
-                    padding="max_length",
+    for epoch in range(NUM_EPOCHS):
+        epoch_loss = 0.0
+        num_batches = 0
+
+        for batch_idx, batch in enumerate(dataloader):
+            prompt = batch["prompt"][0]
+            video_tokens = batch["video_tokens"][0]
+
+            # Format training text (limit to 64 visual tokens for memory)
+            token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens[:64])
+            text = f"Create a video of: {prompt} {VIDEO_START} {token_str} {VIDEO_END}"
+
+            inputs = tokenizer(
+                text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=MAX_SEQ_LEN,
+                padding="max_length",
+            )
+
+            # Forward
+            outputs = model(**inputs, labels=inputs["input_ids"])
+            loss = outputs.loss / GRADIENT_ACCUMULATION
+
+            # Backward
+            loss.backward()
+
+            if (batch_idx + 1) % GRADIENT_ACCUMULATION == 0 or (batch_idx + 1) == len(dataloader):
+                torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
+                optimizer.step()
+                optimizer.zero_grad()
+
+            global_step += 1
+            batch_loss = loss.item() * GRADIENT_ACCUMULATION
+            epoch_loss += batch_loss
+            running_loss += batch_loss
+            num_batches += 1
+
+            elapsed = time.time() - start_time
+            steps_per_sec = global_step / elapsed if elapsed > 0 else 0
+
+            if batch_idx % LOG_EVERY == 0:
+                logger.log(
+                    f"  Epoch {epoch + 1}/{NUM_EPOCHS} | "
+                    f"Step {batch_idx + 1}/{len(dataloader)} | "
+                    f"Loss: {batch_loss:.4f} | "
+                    f"Avg: {epoch_loss / num_batches:.4f} | "
+                    f"Speed: {steps_per_sec:.2f} steps/s\n"
                 )
 
-                # Forward pass
-                outputs = model(**inputs, labels=inputs["input_ids"])
-                loss = outputs.loss / GRADIENT_ACCUMULATION
-
-                # Backward pass
-                loss.backward()
-
-                if (batch_idx + 1) % GRADIENT_ACCUMULATION == 0 or (batch_idx + 1) == len(dataloader):
-                    torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
-                    optimizer.step()
-                    optimizer.zero_grad()
+            del outputs, loss
+            gc.collect()
 
-                global_step += 1
-                batch_loss = loss.item() * GRADIENT_ACCUMULATION
-                epoch_loss += batch_loss
-                running_loss += batch_loss
-                num_batches += 1
+        avg_epoch_loss = epoch_loss / num_batches
+        logger.log(f"\n📈 Epoch {epoch + 1} complete. Avg Loss: {avg_epoch_loss:.4f}\n\n")
 
-                elapsed = time.time() - start_time
-                steps_per_sec = global_step / elapsed if elapsed > 0 else 0
+    total_time = time.time() - start_time
+    logger.log(f"✅ Training complete in {total_time:.0f}s ({total_time / 60:.1f} min)\n")
+    logger.log(f"   Final avg loss: {running_loss / global_step:.4f}\n\n")
 
-                if batch_idx % LOG_EVERY == 0:
-                    msg = (
-                        f"  Epoch {epoch + 1}/{NUM_EPOCHS} | "
-                        f"Step {batch_idx + 1}/{len(dataloader)} | "
-                        f"Loss: {batch_loss:.4f} | "
-                        f"Avg: {epoch_loss / num_batches:.4f} | "
-                        f"Speed: {steps_per_sec:.2f} steps/s\n"
-                    )
-                    yield msg
+    # ── 7. Merge & push ──────────────────────────────────────────────────
+    logger.log("🔀 Merging LoRA weights back into base model...\n")
+    model = model.merge_and_unload()
+    logger.log("✅ LoRA merged.\n")
 
-                # Free memory
-                del outputs, loss
-                gc.collect()
+    logger.log("💾 Saving model locally...\n")
+    save_dir = "./trained_model"
+    model.save_pretrained(save_dir, safe_serialization=True)
+    tokenizer.save_pretrained(save_dir)
+    logger.log("✅ Model saved locally.\n")
 
-            avg_epoch_loss = epoch_loss / num_batches
-            yield f"\n📈 Epoch {epoch + 1} complete. Avg Loss: {avg_epoch_loss:.4f}\n\n"
+    logger.log(f"🚀 Pushing to {REPO_ID}...\n")
+    from huggingface_hub import HfApi
 
-    except Exception as e:
-        yield f"\n❌ Training error: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
-
-    total_time = time.time() - start_time
-    yield f"✅ Training complete in {total_time:.0f}s ({total_time / 60:.1f} min)\n"
-    yield f"   Final avg loss: {running_loss / global_step:.4f}\n\n"
+    api = HfApi(token=HF_TOKEN)
 
-    # ── 6. Merge & push ──────────────────────────────────────────────────
-    yield "🔀 Merging LoRA weights back into base model...\n"
     try:
-        model = model.merge_and_unload()
-        yield "✅ LoRA merged.\n"
+        api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
     except Exception as e:
-        yield f"⚠️ Merge note: {e}\n"
+        logger.log(f"⚠️ Repo note: {e}\n")
 
-    yield "💾 Saving model locally...\n"
-    save_dir = "./trained_model"
-    try:
-        model.save_pretrained(save_dir, safe_serialization=True)
-        tokenizer.save_pretrained(save_dir)
-        yield "✅ Model saved locally.\n"
-    except Exception as e:
-        yield f"❌ Save failed: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
-
-    yield f"🚀 Pushing to {REPO_ID}...\n"
-    try:
-        from huggingface_hub import HfApi
+    api.upload_folder(
+        folder_path=save_dir,
+        repo_id=REPO_ID,
+        repo_type="model",
+        commit_message=f"LoRA-trained OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs)",
+    )
+    logger.log(f"✅ Model pushed to https://huggingface.co/{REPO_ID}\n")
+    logger.log("\n🎉 All done! The trained model is now available on HuggingFace.\n")
 
-        api = HfApi(token=HF_TOKEN)
 
-        # Create model repo if it doesn't exist
+# ---------------------------------------------------------------------------
+# Generator version (for Gradio streaming if needed)
+# ---------------------------------------------------------------------------
+def train(data_path: str = "tokenized_dataset.json"):
+    """Generator version that yields log messages."""
+    import tempfile
+    log_path = tempfile.mktemp(suffix=".txt")
+    logger = _Logger(log_path)
+
+    # Start training in a thread
+    t = threading.Thread(target=lambda: _run_training(logger), daemon=True)
+    t.start()
+
+    # Stream log file
+    last_pos = 0
+    while t.is_alive():
+        time.sleep(1)
         try:
-            api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
-        except Exception as e:
-            yield f"⚠️ Repo creation note: {e}\n"
-
-        api.upload_folder(
-            folder_path=save_dir,
-            repo_id=REPO_ID,
-            repo_type="model",
-            commit_message=f"LoRA-trained OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs)",
-        )
-        yield f"✅ Model pushed to https://huggingface.co/{REPO_ID}\n"
-        yield "\n🎉 All done! The trained model is now available on HuggingFace.\n"
-    except Exception as e:
-        yield f"❌ Push failed: {e}\n"
-        yield traceback.format_exc() + "\n"
-        raise
+            with open(log_path, "r") as f:
+                f.seek(last_pos)
+                new_content = f.read()
+                last_pos = f.tell()
+            if new_content:
+                yield new_content
+        except:
+            pass
+
+    # Final read
+    time.sleep(1)
+    try:
+        with open(log_path, "r") as f:
+            f.seek(last_pos)
+            final = f.read()
+        if final:
+            yield final
+    except:
+        pass
 
 
 # ---------------------------------------------------------------------------
-# CLI entry point (for local testing)
+# CLI entry point
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
     data_path = sys.argv[1] if len(sys.argv) > 1 else "tokenized_dataset.json"
-    for log_msg in train(data_path):
-        print(log_msg, end="", flush=True)
+    run_training_to_file("/tmp/training_log.txt")