Updated: scaled pipeline with real data (10K images, 5K LLM samples, checkpoint/resume support)

train_full_pipeline.py

@@ -1,12 +1,13 @@
 #!/usr/bin/env python3
 """
-Full Pipeline: Train VQ-VAE → Tokenize OpenVid → Train LLM → Push to EeshaAI/zeeb
-=================================================================================
+Full Pipeline: Train VQ-VAE → Tokenize Data → Train LLM → Push to EeshaAI/zeeb
+================================================================================
 Runs on HuggingFace Spaces (free CPU tier, 16GB RAM).
+Uses /data/ persistent volume for checkpoints (survives Space restarts).
 
-Phase 1: Train VQ-VAE on COCO 2017 images (118K real images, streaming)
-Phase 2: Stream 10K clips from OpenVid-1M → tokenize via trained VQ-VAE → save integers
-Phase 3: Fine-tune OLMo 2 1B with LoRA on 10K tokenized samples
+Phase 1: Train VQ-VAE on real images (COCO/imagenette, streaming)
+Phase 2: Tokenize image-text pairs through trained VQ-VAE
+Phase 3: Fine-tune OLMo 2 1B with LoRA on tokenized data
 Phase 4: Push trained model to EeshaAI/zeeb
 """
 
@@ -17,8 +18,9 @@ import time
 import gc
 import threading
 import traceback
+import hashlib
 import numpy as np
-from typing import Optional
+from typing import Optional, List, Dict, Any
 
 import torch
 import torch.nn as nn
@@ -38,27 +40,36 @@ VIDEO_START = "<video_start>"
 VIDEO_END = "<video_end>"
 VIDEO_PAD = "<video_pad>"
 
+# Persistent storage
+DATA_DIR = os.environ.get("DATA_DIR", "/data")
+PERSIST_DIR = os.path.join(DATA_DIR, "zeeb_checkpoints")
+os.makedirs(PERSIST_DIR, exist_ok=True)
+
 # VQ-VAE training
 VQ_VAE_EPOCHS = 5
-VQ_VAE_LR = 1e-3
-VQ_VAE_BATCH = 32
-VQ_VAE_IMG_SIZE = 128  # resize images to 128x128
+VQ_VAE_LR = 3e-4
+VQ_VAE_BATCH = 8
+VQ_VAE_IMG_SIZE = 128
+VQ_VAE_MAX_IMAGES = 10000  # Train on 10K real images
 
-# Dataset preparation
-NUM_OPENVID_CLIPS = 10000
-TOKENS_PER_CLIP = 128  # number of visual tokens per video clip
+# Tokenization
+TOKENS_PER_SAMPLE = 64  # 8x8 grid
+NUM_TOKENIZE_SAMPLES = 10000  # Tokenize 10K image-text pairs
 
 # LLM training
-NUM_EPOCHS = 3
+NUM_EPOCHS = 2
 LORA_R = 4
 LORA_ALPHA = 8
 LORA_DROPOUT = 0.05
-LEARNING_RATE = 1e-4
+LEARNING_RATE = 5e-5
 BATCH_SIZE = 1
-MAX_SEQ_LEN = 384
-GRADIENT_ACCUMULATION = 4
+MAX_SEQ_LEN = 256
+GRADIENT_ACCUMULATION = 8
+LLM_TRAIN_SAMPLES = 5000  # Train on 5K samples (feasible on CPU)
+SAVE_EVERY = 500  # Save checkpoint every N steps
 
-LOG_FILE = "/tmp/pipeline_log.txt"
+LOG_FILE = os.path.join(DATA_DIR, "pipeline_log.txt")
+STATE_FILE = os.path.join(PERSIST_DIR, "pipeline_state.json")
 
 
 # ============================================================================
@@ -69,30 +80,84 @@ class Logger:
         self.path = path
         self.lock = threading.Lock()
         with open(path, "w") as f:
-            f.write("🚀 Zeeb Full Pipeline Starting...\n\n")
+            f.write("Zeeb Full Pipeline Starting...\n\n")
 
     def log(self, msg):
+        timestamp = time.strftime("%H:%M:%S")
+        line = f"[{timestamp}] {msg}"
         with self.lock:
-            with open(self.path, "a") as f:
-                f.write(msg)
-                f.flush()
-            print(msg, end="", flush=True)
+            try:
+                with open(self.path, "a") as f:
+                    f.write(line)
+                    f.flush()
+            except:
+                pass
+        print(line, end="", flush=True)
+
+
+# ============================================================================
+# PIPELINE STATE (for resume after restart)
+# ============================================================================
+class PipelineState:
+    """Track pipeline progress so we can resume after Space restarts."""
+    
+    def __init__(self):
+        self.state = {
+            "phase": 0,  # 0=not started, 1=vq_vae, 2=tokenize, 3=llm, 4=done
+            "vq_vae_done": False,
+            "vq_vae_epoch": 0,
+            "vq_vae_batch": 0,
+            "tokenize_done": False,
+            "tokenize_count": 0,
+            "llm_done": False,
+            "llm_step": 0,
+            "llm_epoch": 0,
+            "pushed": False,
+        }
+        self.load()
+
+    def load(self):
+        if os.path.exists(STATE_FILE):
+            try:
+                with open(STATE_FILE) as f:
+                    saved = json.load(f)
+                self.state.update(saved)
+            except:
+                pass
+
+    def save(self):
+        try:
+            with open(STATE_FILE, "w") as f:
+                json.dump(self.state, f, indent=2)
+        except:
+            pass
+
+    def update(self, **kwargs):
+        self.state.update(kwargs)
+        self.save()
+
+    @property
+    def phase(self):
+        return self.state.get("phase", 0)
+
+    def is_done(self, phase_name):
+        return self.state.get(f"{phase_name}_done", False)
 
 
 # ============================================================================
-# VQ-VAE MODEL
+# VQ-VAE MODEL (same architecture as in generation code)
 # ============================================================================
 class Encoder(nn.Module):
     def __init__(self, in_channels=3, latent_dim=LATENT_DIM):
         super().__init__()
         self.net = nn.Sequential(
-            nn.Conv2d(in_channels, 64, 4, stride=2, padding=1),   # → 64x64
+            nn.Conv2d(in_channels, 64, 4, stride=2, padding=1),   # -> 64x64
             nn.ReLU(),
-            nn.Conv2d(64, 128, 4, stride=2, padding=1),            # → 32x32
+            nn.Conv2d(64, 128, 4, stride=2, padding=1),            # -> 32x32
             nn.ReLU(),
-            nn.Conv2d(128, 256, 4, stride=2, padding=1),           # → 16x16
+            nn.Conv2d(128, 256, 4, stride=2, padding=1),           # -> 16x16
             nn.ReLU(),
-            nn.Conv2d(256, latent_dim, 4, stride=2, padding=1),    # → 8x8
+            nn.Conv2d(256, latent_dim, 4, stride=2, padding=1),    # -> 8x8
         )
 
     def forward(self, x):
@@ -134,13 +199,13 @@ class Decoder(nn.Module):
     def __init__(self, out_channels=3, latent_dim=LATENT_DIM):
         super().__init__()
         self.net = nn.Sequential(
-            nn.ConvTranspose2d(latent_dim, 256, 4, stride=2, padding=1),  # → 16x16
+            nn.ConvTranspose2d(latent_dim, 256, 4, stride=2, padding=1),  # -> 16x16
             nn.ReLU(),
-            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),         # → 32x32
+            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),         # -> 32x32
             nn.ReLU(),
-            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),          # → 64x64
+            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),          # -> 64x64
             nn.ReLU(),
-            nn.ConvTranspose2d(64, out_channels, 4, stride=2, padding=1), # → 128x128
+            nn.ConvTranspose2d(64, out_channels, 4, stride=2, padding=1), # -> 128x128
             nn.Sigmoid(),
         )
 
@@ -190,49 +255,120 @@ class VQVAE(nn.Module):
 
 
 # ============================================================================
-# PHASE 1: TRAIN VQ-VAE ON COCO IMAGES
+# DATASET HELPERS
 # ============================================================================
-def train_vq_vae(logger: Logger) -> VQVAE:
-    """Train VQ-VAE on COCO 2017 images (streaming, so no massive download)."""
-    logger.log("=" * 60 + "\n")
-    logger.log("PHASE 1: Training VQ-VAE on COCO 2017 images\n")
-    logger.log("=" * 60 + "\n\n")
 
+# Imagenette class names for generating captions
+IMAGENETTE_CLASSES = {
+    0: "a fish in water",
+    1: "a dog running in a field", 
+    2: "a cassette player on a table",
+    3: "a chainsaw cutting wood",
+    4: "a church with a tall steeple",
+    5: "a French horn on stage",
+    6: "a garbage truck on the street",
+    7: "a gas station at night",
+    8: "a golf ball on a green",
+    9: "a parachute in the sky",
+}
+
+CIFAR10_CLASSES = ["airplane flying", "automobile on road", "bird in tree",
+                   "cat sitting", "deer in forest", "dog playing", "frog on lily pad",
+                   "horse running", "ship on ocean", "truck driving"]
+
+
+def load_image_dataset(logger: Logger):
+    """Load an image dataset for VQ-VAE training. Returns (stream, image_key, caption_key, name)."""
     from datasets import load_dataset
-    from torchvision import transforms
-
-    # Try multiple COCO/image datasets (some have compatibility issues)
-    logger.log("📦 Loading image dataset (trying multiple sources)...\n")
-    coco = None
-    image_key = "image"
-
+    
+    # Try datasets with both images and good captions
     dataset_sources = [
-        ("detection-datasets/coco", "train", "image"),
-        ("rafaelpadilla/coco2017", "train", "image"),
-        ("frgfm/imagenette", "train", "image"),
-        ("zh-plus/tiny-imagenet", "train", "image"),
-        ("cifar10", "train", "img"),
+        # (dataset_name, split, image_key, caption_key, description)
+        ("detection-datasets/coco", "train", "image", "caption", "COCO 2017 (detection)"),
+        ("frgfm/imagenette", "train", "image", "label", "Imagenette (10 classes)"),
+        ("cifar10", "train", "img", "label", "CIFAR-10"),
     ]
-
-    for ds_name, ds_split, ds_img_key in dataset_sources:
+    
+    for ds_name, ds_split, img_key, cap_key, desc in dataset_sources:
         try:
-            logger.log(f"  Trying {ds_name}...\n")
-            coco = load_dataset(ds_name, split=ds_split, streaming=True, trust_remote_code=True)
-            # Test first item
-            test_item = next(iter(coco))
-            if ds_img_key in test_item:
-                image_key = ds_img_key
-                logger.log(f"  ✅ Using {ds_name} (image key: '{image_key}')\n")
-                break
-            else:
-                logger.log(f"  ⚠️ No '{ds_img_key}' key in {ds_name}, keys: {list(test_item.keys())}\n")
-                coco = None
+            logger.log(f"  Trying {ds_name} ({desc})...\n")
+            ds = load_dataset(ds_name, split=ds_split, streaming=True, trust_remote_code=True)
+            test_item = next(iter(ds))
+            
+            # Verify keys exist
+            actual_img_key = img_key if img_key in test_item else None
+            actual_cap_key = cap_key if cap_key in test_item else None
+            
+            if actual_img_key is None:
+                # Try common alternatives
+                for k in ["image", "img", "png", "jpg"]:
+                    if k in test_item:
+                        actual_img_key = k
+                        break
+            
+            if actual_img_key is None:
+                logger.log(f"  No image key found in {ds_name}. Keys: {list(test_item.keys())}\n")
+                continue
+            
+            logger.log(f"  Using {ds_name}! img_key='{actual_img_key}', cap_key='{actual_cap_key}'\n")
+            return ds, actual_img_key, actual_cap_key, ds_name
+            
         except Exception as e:
-            logger.log(f"  ❌ {ds_name} failed: {str(e)[:100]}\n")
-            coco = None
+            logger.log(f"  Failed: {str(e)[:100]}\n")
+            continue
+    
+    return None, None, None, None
+
+
+def get_caption(item, cap_key, ds_name, index):
+    """Extract or generate a caption for a dataset item."""
+    if cap_key and cap_key in item and item[cap_key] is not None:
+        cap = item[cap_key]
+        if isinstance(cap, list):
+            return cap[0] if cap else f"image {index}"
+        elif isinstance(cap, str):
+            return cap
+        elif isinstance(cap, int):
+            # Class label - convert to descriptive caption
+            if "imagenette" in ds_name.lower():
+                return IMAGENETTE_CLASSES.get(cap, f"photo of object {cap}")
+            elif "cifar" in ds_name.lower():
+                return CIFAR10_CLASSES[cap] if cap < len(CIFAR10_CLASSES) else f"photo of class {cap}"
+            else:
+                return f"photo of a {cap}"
+    return f"image {index}"
+
+
+# ============================================================================
+# PHASE 1: TRAIN VQ-VAE ON REAL IMAGES
+# ============================================================================
+def train_vq_vae(logger: Logger, state: PipelineState) -> VQVAE:
+    """Train VQ-VAE on real images with checkpoint/resume support."""
+    logger.log("=" * 60 + "\n")
+    logger.log("PHASE 1: Training VQ-VAE on real images\n")
+    logger.log("=" * 60 + "\n\n")
+    
+    from datasets import load_dataset
+    from torchvision import transforms
+    from PIL import Image
+
+    # Check if already done
+    if state.is_done("vq_vae"):
+        logger.log("VQ-VAE already trained! Loading checkpoint...\n")
+        ckpt_path = os.path.join(PERSIST_DIR, "vq_vae_best.pt")
+        if os.path.exists(ckpt_path):
+            model = VQVAE()
+            model.load_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=False))
+            logger.log("Loaded trained VQ-VAE from checkpoint.\n")
+            return model
+        else:
+            logger.log("Checkpoint not found, retraining...\n")
+            state.update(vq_vae_done=False)
 
-    if coco is None:
-        logger.log("❌ No dataset could be loaded! Cannot train VQ-VAE.\n")
+    # Load dataset
+    ds, img_key, cap_key, ds_name = load_image_dataset(logger)
+    if ds is None:
+        logger.log("No dataset available! Cannot train VQ-VAE.\n")
         return None
 
     # Image transforms
@@ -242,7 +378,7 @@ def train_vq_vae(logger: Logger) -> VQVAE:
     ])
 
     class ImageStreamDataset(IterableDataset):
-        def __init__(self, hf_dataset, transform, img_key, max_samples=50000):
+        def __init__(self, hf_dataset, transform, img_key, max_samples):
             self.dataset = hf_dataset
             self.transform = transform
             self.img_key = img_key
@@ -258,23 +394,38 @@ def train_vq_vae(logger: Logger) -> VQVAE:
                     if img.mode != "RGB":
                         img = img.convert("RGB")
                     tensor = self.transform(img)
-                    yield tensor
                     count += 1
+                    yield tensor
                 except Exception:
                     continue
 
-    dataset = ImageStreamDataset(coco, transform, image_key, max_samples=50000)
+    dataset = ImageStreamDataset(ds, transform, img_key, VQ_VAE_MAX_IMAGES)
     dataloader = DataLoader(dataset, batch_size=VQ_VAE_BATCH, num_workers=0)
 
-    # Initialize model
+    # Initialize or resume model
     model = VQVAE()
     n_params = sum(p.numel() for p in model.parameters()) / 1e6
-    logger.log(f"✅ VQ-VAE initialized: {n_params:.1f}M parameters\n")
+    logger.log(f"VQ-VAE initialized: {n_params:.1f}M parameters\n")
+    
+    # Resume from checkpoint if available
+    resume_ckpt = os.path.join(PERSIST_DIR, "vq_vae_latest.pt")
+    start_epoch = 0
+    if os.path.exists(resume_ckpt):
+        try:
+            ckpt = torch.load(resume_ckpt, map_location="cpu", weights_only=False)
+            model.load_state_dict(ckpt["model_state_dict"])
+            start_epoch = ckpt.get("epoch", 0)
+            logger.log(f"Resumed VQ-VAE from epoch {start_epoch}\n")
+        except:
+            logger.log("Could not resume checkpoint, starting fresh.\n")
 
     optimizer = torch.optim.Adam(model.parameters(), lr=VQ_VAE_LR)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=VQ_VAE_EPOCHS)
     model.train()
 
-    for epoch in range(VQ_VAE_EPOCHS):
+    best_loss = float('inf')
+    
+    for epoch in range(start_epoch, VQ_VAE_EPOCHS):
         epoch_loss = 0.0
         epoch_recon = 0.0
         epoch_vq = 0.0
@@ -288,6 +439,7 @@ def train_vq_vae(logger: Logger) -> VQVAE:
 
             optimizer.zero_grad()
             loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
             optimizer.step()
 
             epoch_loss += loss.item()
@@ -295,202 +447,103 @@ def train_vq_vae(logger: Logger) -> VQVAE:
             epoch_vq += vq_loss.item()
             num_batches += 1
 
-            if batch_idx % 50 == 0 and batch_idx > 0:
+            if batch_idx % 100 == 0 and batch_idx > 0:
                 avg = epoch_loss / num_batches
                 avg_r = epoch_recon / num_batches
                 avg_v = epoch_vq / num_batches
                 logger.log(f"  Epoch {epoch+1}/{VQ_VAE_EPOCHS} | Batch {batch_idx} | "
                           f"Loss: {avg:.4f} (recon: {avg_r:.4f}, vq: {avg_v:.4f})\n")
 
-            del recon, vq_loss, loss
-            if batch_idx % 200 == 0:
+            del recon, vq_loss, loss, batch
+            if batch_idx % 100 == 0:
                 gc.collect()
 
+        # End of epoch
+        scheduler.step()
         elapsed = time.time() - start_time
         avg_loss = epoch_loss / max(num_batches, 1)
-        logger.log(f"\n📈 Epoch {epoch+1} done. Avg Loss: {avg_loss:.4f} | "
+        avg_recon = epoch_recon / max(num_batches, 1)
+        logger.log(f"\nEpoch {epoch+1} done. Loss: {avg_loss:.4f} (recon: {avg_recon:.4f}) | "
                   f"Batches: {num_batches} | Time: {elapsed:.0f}s\n\n")
+        
+        # Save checkpoint
+        ckpt_path = os.path.join(PERSIST_DIR, "vq_vae_latest.pt")
+        torch.save({
+            "epoch": epoch + 1,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "loss": avg_loss,
+        }, ckpt_path)
+        
+        # Save best model
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            best_path = os.path.join(PERSIST_DIR, "vq_vae_best.pt")
+            torch.save(model.state_dict(), best_path)
+            logger.log(f"  New best model! Loss: {avg_loss:.4f}\n")
+        
+        state.update(vq_vae_epoch=epoch + 1)
+        gc.collect()
 
-    # Save
+    # Save final
+    final_path = os.path.join(PERSIST_DIR, "vq_vae_best.pt")
+    if not os.path.exists(final_path):
+        torch.save(model.state_dict(), final_path)
+    
+    # Also save to root for easy access
     torch.save(model.state_dict(), "vq_vae_real.pt")
-    logger.log("✅ VQ-VAE saved to vq_vae_real.pt\n\n")
+    
+    state.update(vq_vae_done=True, phase=2)
+    logger.log(f"VQ-VAE training complete! Best loss: {best_loss:.4f}\n\n")
     return model
 
 
 # ============================================================================
-# PHASE 2: TOKENIZE OPENVID-1M DATASET
+# PHASE 2: TOKENIZE IMAGE-TEXT PAIRS
 # ============================================================================
-def tokenize_openvid(logger: Logger, vq_vae: Optional[VQVAE] = None):
-    """Stream OpenVid-1M, tokenize videos with VQ-VAE, save tokenized data."""
+def tokenize_dataset(logger: Logger, state: PipelineState, vq_vae: Optional[VQVAE] = None):
+    """Tokenize image-text pairs through trained VQ-VAE."""
     logger.log("=" * 60 + "\n")
-    logger.log("PHASE 2: Tokenizing OpenVid-1M dataset (10K clips)\n")
+    logger.log("PHASE 2: Tokenizing image-text pairs\n")
     logger.log("=" * 60 + "\n\n")
 
+    if state.is_done("tokenize"):
+        logger.log("Tokenization already done! Loading cached data...\n")
+        data_path = os.path.join(PERSIST_DIR, "tokenized_dataset.json")
+        if os.path.exists(data_path):
+            with open(data_path) as f:
+                data = json.load(f)
+            logger.log(f"Loaded {len(data)} tokenized samples.\n")
+            return data
+        else:
+            logger.log("Cached data not found, re-tokenizing...\n")
+            state.update(tokenize_done=False)
+
     # Load VQ-VAE if not provided
     if vq_vae is None:
-        if os.path.exists("vq_vae_real.pt"):
+        ckpt_path = os.path.join(PERSIST_DIR, "vq_vae_best.pt")
+        if os.path.exists(ckpt_path):
+            vq_vae = VQVAE()
+            vq_vae.load_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=False))
+            logger.log("Loaded trained VQ-VAE for tokenization.\n")
+        elif os.path.exists("vq_vae_real.pt"):
             vq_vae = VQVAE()
             vq_vae.load_state_dict(torch.load("vq_vae_real.pt", map_location="cpu", weights_only=False))
-            logger.log("✅ Loaded trained VQ-VAE from vq_vae_real.pt\n")
+            logger.log("Loaded VQ-VAE from vq_vae_real.pt.\n")
         else:
-            logger.log("❌ No trained VQ-VAE found! Run Phase 1 first.\n")
+            logger.log("No trained VQ-VAE found! Run Phase 1 first.\n")
             return None
 
     vq_vae.eval()
-
-    from datasets import load_dataset
-
-    logger.log("📦 Loading OpenVid-1M dataset (streaming)...\n")
-    try:
-        dataset = load_dataset("NJU-PCALab/OpenVid-1M", split="train", streaming=True, trust_remote_code=True)
-    except Exception as e:
-        logger.log(f"⚠️ OpenVid-1M load error: {e}\n")
-        logger.log("🔄 Trying alternative: WebVid-2M...\n")
-        try:
-            dataset = load_dataset("tmpdump/webvid10m", split="train", streaming=True, trust_remote_code=True)
-        except Exception as e2:
-            logger.log(f"⚠️ WebVid load error: {e2}\n")
-            logger.log("🔄 Falling back to COCO captions (image-only, but much more data)...\n")
-            return _tokenize_coco_fallback(logger, vq_vae)
-
-    # Tokenize clips
-    tokenized_data = []
-    count = 0
-    errors = 0
-
-    for item in dataset:
-        if count >= NUM_OPENVID_CLIPS:
-            break
-
-        try:
-            # Get text caption
-            caption = ""
-            for key in ["caption", "text", "description", "title"]:
-                if key in item and item[key]:
-                    caption = item[key]
-                    break
-
-            if not caption:
-                caption = f"video clip {count}"
-
-            # Get video frames
-            video = item.get("video", None)
-            if video is None:
-                errors += 1
-                continue
-
-            # Process video frames
-            import io
-            from PIL import Image
-
-            frames = []
-            if hasattr(video, 'read'):
-                # It's bytes
-                pass
-
-            # Try to extract frames
-            if isinstance(video, dict) and "bytes" in video:
-                video_bytes = video["bytes"]
-            elif isinstance(video, bytes):
-                video_bytes = video
-            else:
-                errors += 1
-                continue
-
-            # Use imageio or decord to extract frames
-            try:
-                import imageio
-                reader = imageio.get_reader(io.BytesIO(video_bytes), format='mp4')
-                for i, frame in enumerate(reader):
-                    if i >= 4:  # Take first 4 frames
-                        break
-                    img = Image.fromarray(frame).convert("RGB").resize((128, 128))
-                    frames.append(np.array(img))
-                reader.close()
-            except Exception:
-                errors += 1
-                continue
-
-            if not frames:
-                errors += 1
-                continue
-
-            # Tokenize frames through VQ-VAE
-            from torchvision import transforms
-            transform = transforms.ToTensor()
-            all_tokens = []
-
-            for frame in frames:
-                img_tensor = transform(Image.fromarray(frame)).unsqueeze(0)
-                with torch.no_grad():
-                    tokens = vq_vae.encode(img_tensor)
-                    all_tokens.extend(tokens.flatten().tolist())
-
-            # Truncate/pad to fixed length
-            all_tokens = all_tokens[:TOKENS_PER_CLIP]
-            while len(all_tokens) < TOKENS_PER_CLIP:
-                all_tokens.append(0)
-
-            tokenized_data.append({
-                "text_prompt": caption,
-                "video_tokens": all_tokens,
-            })
-
-            count += 1
-            if count % 100 == 0:
-                logger.log(f"  Tokenized {count}/{NUM_OPENVID_CLIPS} clips (errors: {errors})\n")
-
-        except Exception as e:
-            errors += 1
-            if errors <= 3:
-                logger.log(f"  ⚠️ Error on item: {e}\n")
-            continue
-
-    if not tokenized_data:
-        logger.log("❌ No clips tokenized from OpenVid-1M! Falling back to COCO captions.\n")
-        return _tokenize_coco_fallback(logger, vq_vae)
-
-    # Save
-    with open("tokenized_dataset.json", "w") as f:
-        json.dump(tokenized_data, f)
-
-    logger.log(f"\n✅ Tokenized {len(tokenized_data)} clips saved to tokenized_dataset.json\n")
-    logger.log(f"   Errors: {errors}\n\n")
-    return tokenized_data
-
-
-def _tokenize_coco_fallback(logger: Logger, vq_vae: VQVAE):
-    """Fallback: tokenize image-text pairs from available datasets."""
-    logger.log("📦 Using image-text pairs as fallback (50K samples)...\n")
-
+    
     from datasets import load_dataset
     from torchvision import transforms
     from PIL import Image
 
-    # Try multiple datasets
-    ds = None
-    image_key = "image"
-    caption_key = "text"
-
-    for ds_name, ds_split, img_k, cap_k in [
-        ("detection-datasets/coco", "train", "image", "caption"),
-        ("frgfm/imagenette", "train", "image", "label"),
-        ("cifar10", "train", "img", "label"),
-    ]:
-        try:
-            logger.log(f"  Trying {ds_name}...\n")
-            ds = load_dataset(ds_name, split=ds_split, streaming=True, trust_remote_code=True)
-            test = next(iter(ds))
-            image_key = img_k if img_k in test else "image"
-            caption_key = cap_k if cap_k in test else "text"
-            logger.log(f"  ✅ Using {ds_name} (img='{image_key}', cap='{caption_key}')\n")
-            break
-        except Exception as e:
-            logger.log(f"  ❌ {ds_name}: {str(e)[:100]}\n")
-            ds = None
-
+    # Load dataset with captions
+    ds, img_key, cap_key, ds_name = load_image_dataset(logger)
     if ds is None:
-        logger.log("❌ No dataset available for tokenization!\n")
+        logger.log("No dataset available for tokenization!\n")
         return None
 
     transform = transforms.Compose([
@@ -498,50 +551,43 @@ def _tokenize_coco_fallback(logger: Logger, vq_vae: VQVAE):
         transforms.ToTensor(),
     ])
 
-    vq_vae.eval()
     tokenized_data = []
     count = 0
+    errors = 0
+    
+    # Check for partial tokenization (resume support)
+    partial_path = os.path.join(PERSIST_DIR, "tokenized_partial.json")
+    if os.path.exists(partial_path):
+        try:
+            with open(partial_path) as f:
+                tokenized_data = json.load(f)
+            count = len(tokenized_data)
+            logger.log(f"Resuming tokenization from {count} samples.\n")
+        except:
+            tokenized_data = []
+            count = 0
 
-    # Label mapping for datasets that only have class labels
-    label_names = {
-        "cifar10": ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"],
-    }
+    logger.log(f"Tokenizing up to {NUM_TOKENIZE_SAMPLES} images...\n")
 
     for item in ds:
-        if count >= 50000:
+        if count >= NUM_TOKENIZE_SAMPLES:
             break
 
         try:
-            img = item[image_key]
+            img = item[img_key]
             if img.mode != "RGB":
                 img = img.convert("RGB")
 
-            # Get caption
-            caption = ""
-            if caption_key in item and item[caption_key] is not None:
-                cap = item[caption_key]
-                if isinstance(cap, list):
-                    caption = cap[0] if cap else ""
-                elif isinstance(cap, int):
-                    # It's a class label - convert to text
-                    ds_name_short = ds_name.split("/")[0] if "/" in ds_name else ds_name
-                    if ds_name_short in label_names and cap < len(label_names[ds_name_short]):
-                        caption = f"a photo of a {label_names[ds_name_short][cap]}"
-                    else:
-                        caption = f"image class {cap}"
-                else:
-                    caption = str(cap)
-            if not caption:
-                caption = f"image {count}"
+            caption = get_caption(item, cap_key, ds_name, count)
 
             img_tensor = transform(img).unsqueeze(0)
             with torch.no_grad():
                 tokens = vq_vae.encode(img_tensor)
                 flat_tokens = tokens.flatten().tolist()
 
-            # Truncate/pad
-            flat_tokens = flat_tokens[:TOKENS_PER_CLIP]
-            while len(flat_tokens) < TOKENS_PER_CLIP:
+            # Truncate/pad to fixed length
+            flat_tokens = flat_tokens[:TOKENS_PER_SAMPLE]
+            while len(flat_tokens) < TOKENS_PER_SAMPLE:
                 flat_tokens.append(0)
 
             tokenized_data.append({
@@ -550,71 +596,110 @@ def _tokenize_coco_fallback(logger: Logger, vq_vae: VQVAE):
             })
 
             count += 1
-            if count % 1000 == 0:
-                logger.log(f"  Tokenized {count}/50000 images\n")
-                # Save checkpoint periodically
-                if count % 10000 == 0:
-                    with open("tokenized_dataset.json", "w") as f:
-                        json.dump(tokenized_data, f)
-                    logger.log(f"  💾 Checkpoint saved ({len(tokenized_data)} samples)\n")
-
-        except Exception:
+            
+            if count % 500 == 0:
+                logger.log(f"  Tokenized {count}/{NUM_TOKENIZE_SAMPLES} images (errors: {errors})\n")
+                # Save partial progress
+                with open(partial_path, "w") as f:
+                    json.dump(tokenized_data, f)
+
+            del img_tensor, tokens
+            if count % 200 == 0:
+                gc.collect()
+
+        except Exception as e:
+            errors += 1
+            if errors <= 5:
+                logger.log(f"  Error on item {count}: {str(e)[:80]}\n")
             continue
 
-    # Final save
+    if not tokenized_data:
+        logger.log("No images tokenized!\n")
+        return None
+
+    # Save final
+    data_path = os.path.join(PERSIST_DIR, "tokenized_dataset.json")
+    with open(data_path, "w") as f:
+        json.dump(tokenized_data, f)
+    
+    # Also save to root
     with open("tokenized_dataset.json", "w") as f:
         json.dump(tokenized_data, f)
 
-    logger.log(f"\n✅ Tokenized {len(tokenized_data)} images saved to tokenized_dataset.json\n\n")
+    # Clean up partial
+    if os.path.exists(partial_path):
+        os.remove(partial_path)
+
+    state.update(tokenize_done=True, tokenize_count=len(tokenized_data), phase=3)
+    logger.log(f"\nTokenized {len(tokenized_data)} images saved (errors: {errors})\n\n")
     return tokenized_data
 
 
 # ============================================================================
 # PHASE 3: TRAIN LLM WITH LORA
 # ============================================================================
-def train_llm(logger: Logger):
+def train_llm(logger: Logger, state: PipelineState):
     """Fine-tune OLMo 2 1B with LoRA on tokenized data."""
     logger.log("=" * 60 + "\n")
-    logger.log("PHASE 3: Fine-tuning OLMo 2 1B + LoRA\n")
+    logger.log("PHASE 3: Fine-tuning OLMo 2 1B + LoRA on real data\n")
     logger.log("=" * 60 + "\n\n")
 
+    if state.is_done("llm"):
+        logger.log("LLM already trained! Skipping.\n")
+        return
+
     from transformers import AutoModelForCausalLM, AutoTokenizer
     from peft import LoraConfig, get_peft_model, TaskType
 
     # Load data
-    data_path = "tokenized_dataset.json"
+    data_path = os.path.join(PERSIST_DIR, "tokenized_dataset.json")
     if not os.path.exists(data_path):
-        logger.log("❌ No tokenized dataset found! Run Phase 2 first.\n")
+        data_path = "tokenized_dataset.json"
+    
+    if not os.path.exists(data_path):
+        logger.log("No tokenized dataset found! Run Phase 2 first.\n")
         return
 
     with open(data_path) as f:
-        data = json.load(f)
-    logger.log(f"📊 Loaded {len(data)} training samples\n")
+        all_data = json.load(f)
+    
+    # Limit to training samples
+    data = all_data[:LLM_TRAIN_SAMPLES]
+    logger.log(f"Loaded {len(all_data)} total samples, using {len(data)} for training\n")
+
+    # Quick data quality check
+    if data:
+        sample = data[0]
+        logger.log(f"Sample prompt: '{sample['text_prompt']}'\n")
+        logger.log(f"Sample tokens (first 10): {sample['video_tokens'][:10]}\n")
+        unique_tokens = len(set(sample['video_tokens']))
+        logger.log(f"Unique tokens in sample: {unique_tokens}\n\n")
 
     # Tokenizer
-    logger.log("📦 Loading OLMo 2 1B tokenizer...\n")
+    logger.log("Loading OLMo 2 1B tokenizer...\n")
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
     # Model
-    logger.log("📦 Loading model (fp32, CPU)...\n")
+    logger.log("Loading model (fp32, CPU)...\n")
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME, trust_remote_code=True, torch_dtype=torch.float32
     )
-    logger.log(f"✅ Model loaded. Original vocab: {len(tokenizer)}\n")
+    orig_vocab = len(tokenizer)
+    logger.log(f"Model loaded. Original vocab: {orig_vocab}\n")
 
     # Expand vocab
-    logger.log(f"🔤 Adding {CODEBOOK_SIZE} visual tokens...\n")
+    logger.log(f"Adding {CODEBOOK_SIZE} visual tokens...\n")
     visual_tokens = [VIDEO_START, VIDEO_END, VIDEO_PAD]
     for i in range(CODEBOOK_SIZE):
         visual_tokens.append(f"<v_{i}>")
     tokenizer.add_tokens(visual_tokens)
     model.resize_token_embeddings(len(tokenizer))
-    logger.log(f"✅ New vocab: {len(tokenizer)}\n")
+    logger.log(f"New vocab: {len(tokenizer)}\n")
 
     # LoRA
-    logger.log(f"🔧 Applying LoRA (r={LORA_R})...\n")
+    logger.log(f"Applying LoRA (r={LORA_R})...\n")
     lora_config = LoraConfig(
         r=LORA_R, lora_alpha=LORA_ALPHA,
         target_modules=["q_proj", "v_proj"],
@@ -624,11 +709,11 @@ def train_llm(logger: Logger):
     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: {trainable:,} / {total:,} trainable ({100*trainable/total:.2f}%)\n")
+    logger.log(f"LoRA: {trainable:,} / {total:,} trainable ({100*trainable/total:.2f}%)\n")
 
     # Dataset
     class VideoTokenDataset(Dataset):
-        def __init__(self, data, max_tokens=TOKENS_PER_CLIP):
+        def __init__(self, data, max_tokens=TOKENS_PER_SAMPLE):
             self.data = data
             self.max_tokens = max_tokens
 
@@ -646,13 +731,32 @@ def train_llm(logger: Logger):
     dataset = VideoTokenDataset(data)
     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\n")
+    logger.log(f"{len(dataset)} samples x {NUM_EPOCHS} epochs = {total_steps} steps\n\n")
+
+    # Optimizer - Adafactor is more memory-efficient for CPU
+    from transformers import Adafactor
+    optimizer = Adafactor(
+        model.parameters(), lr=LEARNING_RATE,
+        relative_step=False, scale_parameter=False, warmup_init=False
+    )
+
+    # Resume from checkpoint if available
+    start_step = state.state.get("llm_step", 0)
+    start_epoch = state.state.get("llm_epoch", 0)
+    
+    llm_ckpt_dir = os.path.join(PERSIST_DIR, "llm_checkpoint")
+    if start_step > 0 and os.path.exists(llm_ckpt_dir):
+        try:
+            logger.log(f"Resuming LLM training from step {start_step}, epoch {start_epoch}\n")
+            # We'd need to skip dataloader steps - for simplicity, restart epoch
+            start_step = 0
+        except:
+            pass
 
-    # Train
-    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
     model.train()
     global_step = 0
     running_loss = 0.0
+    best_loss = float('inf')
     start_time = time.time()
 
     for epoch in range(NUM_EPOCHS):
@@ -663,8 +767,8 @@ def train_llm(logger: Logger):
             prompt = batch["prompt"][0]
             video_tokens = batch["video_tokens"][0]
 
-            # Format: use 64 visual tokens per sample for memory
-            token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens[:64])
+            # Format training text
+            token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens)
             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")
@@ -686,73 +790,140 @@ def train_llm(logger: Logger):
             if batch_idx % 100 == 0:
                 elapsed = time.time() - start_time
                 speed = global_step / elapsed if elapsed > 0 else 0
+                eta = (total_steps - global_step) / speed if speed > 0 else 0
                 logger.log(f"  Epoch {epoch+1}/{NUM_EPOCHS} | Step {batch_idx+1}/{len(dataloader)} | "
                           f"Loss: {batch_loss:.4f} | Avg: {epoch_loss/num_batches:.4f} | "
-                          f"Speed: {speed:.2f} steps/s\n")
+                          f"Speed: {speed:.2f} steps/s | ETA: {eta/60:.0f}m\n")
 
-            del outputs, loss
-            gc.collect()
+            # Save checkpoint periodically
+            if global_step % SAVE_EVERY == 0 and global_step > 0:
+                ckpt_loss = running_loss / global_step
+                logger.log(f"  Saving checkpoint at step {global_step} (loss: {ckpt_loss:.4f})...\n")
+                try:
+                    os.makedirs(llm_ckpt_dir, exist_ok=True)
+                    model.save_pretrained(llm_ckpt_dir)
+                    tokenizer.save_pretrained(llm_ckpt_dir)
+                    state.update(llm_step=global_step, llm_epoch=epoch)
+                except Exception as e:
+                    logger.log(f"  Checkpoint save failed: {str(e)[:80]}\n")
+
+            del outputs, loss, inputs
+            if batch_idx % 50 == 0:
+                gc.collect()
 
-        logger.log(f"\n📈 Epoch {epoch+1} done. Avg Loss: {epoch_loss/num_batches:.4f}\n\n")
+        avg_epoch_loss = epoch_loss / max(num_batches, 1)
+        logger.log(f"\nEpoch {epoch+1} done. Avg Loss: {avg_epoch_loss:.4f}\n\n")
+        
+        # Save best model
+        if avg_epoch_loss < best_loss:
+            best_loss = avg_epoch_loss
+        
+        state.update(llm_epoch=epoch + 1)
 
     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")
+    final_loss = running_loss / max(global_step, 1)
+    logger.log(f"Training complete in {total_time:.0f}s ({total_time/60:.1f} min)\n")
+    logger.log(f"Final avg loss: {final_loss:.4f}\n\n")
 
     # Merge & save
-    logger.log("🔀 Merging LoRA → base model...\n")
+    logger.log("Merging LoRA into base model...\n")
     model = model.merge_and_unload()
 
-    save_dir = "./trained_model"
+    save_dir = os.path.join(PERSIST_DIR, "trained_model")
+    os.makedirs(save_dir, exist_ok=True)
     model.save_pretrained(save_dir, safe_serialization=True)
     tokenizer.save_pretrained(save_dir)
 
-    # Also save VQ-VAE
-    if os.path.exists("vq_vae_real.pt"):
+    # Also save VQ-VAE checkpoint
+    vq_path = os.path.join(PERSIST_DIR, "vq_vae_best.pt")
+    if os.path.exists(vq_path):
         import shutil
-        shutil.copy("vq_vae_real.pt", f"{save_dir}/vq_vae_final.pt")
+        shutil.copy(vq_path, os.path.join(save_dir, "vq_vae_final.pt"))
+    elif os.path.exists("vq_vae_real.pt"):
+        import shutil
+        shutil.copy("vq_vae_real.pt", os.path.join(save_dir, "vq_vae_final.pt"))
 
     # Copy tokenized dataset
-    if os.path.exists("tokenized_dataset.json"):
+    if os.path.exists(os.path.join(PERSIST_DIR, "tokenized_dataset.json")):
         import shutil
-        shutil.copy("tokenized_dataset.json", f"{save_dir}/tokenized_dataset.json")
+        shutil.copy(os.path.join(PERSIST_DIR, "tokenized_dataset.json"), 
+                    os.path.join(save_dir, "tokenized_dataset.json"))
 
-    logger.log("✅ Model saved locally.\n")
+    logger.log("Model saved locally.\n")
 
-    # Push
-    logger.log(f"🚀 Pushing to {REPO_ID}...\n")
-    from huggingface_hub import HfApi
-    api = HfApi(token=HF_TOKEN)
-    try:
-        api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
-    except:
-        pass
-    api.upload_folder(folder_path=save_dir, repo_id=REPO_ID, repo_type="model",
-                      commit_message=f"LoRA OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs, {len(data)} samples)")
-    logger.log(f"✅ Pushed to https://huggingface.co/{REPO_ID}\n\n")
+    # Push to Hub
+    if HF_TOKEN:
+        logger.log(f"Pushing to {REPO_ID}...\n")
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi(token=HF_TOKEN)
+            try:
+                api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
+            except:
+                pass
+            api.upload_folder(
+                folder_path=save_dir, repo_id=REPO_ID, repo_type="model",
+                commit_message=f"LoRA OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs, {len(data)} real samples, loss={final_loss:.4f})"
+            )
+            logger.log(f"Pushed to https://huggingface.co/{REPO_ID}\n\n")
+            state.update(pushed=True)
+        except Exception as e:
+            logger.log(f"Push failed: {str(e)[:200]}\n")
+            logger.log("Model is saved locally and can be pushed manually.\n\n")
+    else:
+        logger.log("No HF_TOKEN set, skipping push.\n")
+
+    state.update(llm_done=True, phase=4)
 
 
 # ============================================================================
 # MAIN PIPELINE
 # ============================================================================
-def run_pipeline(log_path: str = LOG_FILE):
+def run_pipeline(log_path: str = None):
+    if log_path is None:
+        log_path = LOG_FILE
+    
     logger = Logger(log_path)
+    state = PipelineState()
+    
+    logger.log(f"Pipeline state: Phase {state.phase}\n")
+    logger.log(f"Persistent dir: {PERSIST_DIR}\n")
+    logger.log(f"Data dir contents: {os.listdir(PERSIST_DIR) if os.path.exists(PERSIST_DIR) else 'empty'}\n\n")
 
     try:
         # Phase 1: Train VQ-VAE
-        vq_vae = train_vq_vae(logger)
+        if not state.is_done("vq_vae"):
+            state.update(phase=1)
+            vq_vae = train_vq_vae(logger, state)
+        else:
+            logger.log("Skipping Phase 1 (already done)\n")
+            vq_vae = None
+        
         gc.collect()
 
         # Phase 2: Tokenize dataset
-        tokenize_openvid(logger, vq_vae)
+        if not state.is_done("tokenize"):
+            state.update(phase=2)
+            tokenize_dataset(logger, state, vq_vae)
+        else:
+            logger.log("Skipping Phase 2 (already done)\n")
+        
         gc.collect()
 
         # Phase 3: Train LLM
-        train_llm(logger)
+        if not state.is_done("llm"):
+            state.update(phase=3)
+            train_llm(logger, state)
+        else:
+            logger.log("Skipping Phase 3 (already done)\n")
 
-        logger.log("\n🎉 FULL PIPELINE COMPLETE!\n")
+        logger.log("\n" + "=" * 60 + "\n")
+        logger.log("FULL PIPELINE COMPLETE!\n")
+        logger.log("=" * 60 + "\n")
+        state.update(phase=4)
+        
     except Exception as e:
-        logger.log(f"\n❌ PIPELINE ERROR: {e}\n")
+        logger.log(f"\nPIPELINE ERROR: {e}\n")
         logger.log(traceback.format_exc())