Upload scripts/train.py with huggingface_hub

scripts/train.py

@@ -0,0 +1,402 @@
+# ==============================================================================
+# 🚀 ViuTranslate — Dedicated Neural Translation Master Training Engine
+# ==============================================================================
+# Model Architecture: ViuAI Sarus-500M
+# Target Repository:  ViuAI/ViuTranslate
+# Dataset Repository: ViuAI/ViuTranslate-Data
+# Hardware:           Auto-Tuned (Kaggle T4 x 2, P100, A100, RTX 3090/4090)
+# ==============================================================================
+
+import os
+import sys
+import math
+import time
+import shutil
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, Sampler
+from huggingface_hub import HfApi, hf_hub_download
+
+# Fix stdout encoding for Windows & Cloud
+if hasattr(sys.stdout, "reconfigure"):
+    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+if hasattr(sys.stderr, "reconfigure"):
+    sys.stderr.reconfigure(encoding="utf-8", errors="replace")
+
+cur_dir = os.path.dirname(os.path.abspath(__file__)) if "__file__" in locals() else os.getcwd()
+if cur_dir not in sys.path:
+    sys.path.insert(0, cur_dir)
+
+from model import ViuAI
+from config import ViuAIConfig
+
+PAD_TOKEN_ID = 64000
+EOT_ID = 64002
+
+DOMAIN_NAMES = {
+    0: "en_to_hi_direct",
+    1: "en_to_hi_command",
+    2: "hi_to_en_direct",
+    3: "hi_to_en_command"
+}
+
+# ------------------------------------------------------------------------------
+# 1. Hardware Profiler
+# ------------------------------------------------------------------------------
+def auto_profile_hardware():
+    if not torch.cuda.is_available():
+        return {
+            "tier": "CPU", "device_name": "CPU", "vram_gb": 0.0,
+            "micro_batch": 1, "grad_accum": 64, "dtype": torch.float32,
+            "desc": "CPU fallback mode"
+        }
+        
+    props = torch.cuda.get_device_properties(0)
+    device_name = props.name
+    vram_gb = props.total_memory / (1024 ** 3)
+    major, minor = props.major, props.minor
+    bf16_supported = torch.cuda.is_bf16_supported()
+    dtype = torch.bfloat16 if bf16_supported else torch.float16
+
+    if vram_gb >= 30:
+        micro_batch, grad_accum = 32, 2
+        desc = "NVIDIA RTX 5090 / High-Tier 32GB Blackwell Beast"
+    elif vram_gb >= 20:
+        micro_batch, grad_accum = 12, 6
+        desc = "Pro-Tier GPU (24GB)"
+    elif vram_gb >= 12:
+        micro_batch, grad_accum = 8, 8
+        desc = "Standard Cloud GPU / 16GB (Kaggle T4 / P100)"
+    else:
+        micro_batch, grad_accum = 4, 16
+        desc = "Budget GPU (< 12GB)"
+
+    return {
+        "tier": "GPU",
+        "device_name": device_name,
+        "vram_gb": vram_gb,
+        "compute_cap": f"{major}.{minor}",
+        "micro_batch": micro_batch,
+        "grad_accum": grad_accum,
+        "effective_batch": micro_batch * grad_accum,
+        "dtype": dtype,
+        "desc": desc
+    }
+
+# ------------------------------------------------------------------------------
+# 2. Dataset & Length Grouping
+# ------------------------------------------------------------------------------
+class TranslationDataset(Dataset):
+    def __init__(self, ids_path: str, labels_path: str, offsets_path: str, domains_path: str = None):
+        self.tokens_mmap = np.load(ids_path, mmap_mode="r")
+        self.labels_mmap = np.load(labels_path, mmap_mode="r")
+        self.offsets = np.load(offsets_path)
+        self.domains = np.load(domains_path) if (domains_path and os.path.exists(domains_path)) else None
+        self.num_samples = len(self.offsets) - 1
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        start_idx = int(self.offsets[idx])
+        end_idx = int(self.offsets[idx + 1])
+        tokens = torch.from_numpy(self.tokens_mmap[start_idx:end_idx].astype(np.int64))
+        labels = torch.from_numpy(self.labels_mmap[start_idx:end_idx].astype(np.int64))
+        domain_id = int(self.domains[idx]) if self.domains is not None else 0
+        return tokens, labels, domain_id
+
+class LengthGroupedBatchSampler(Sampler):
+    def __init__(self, dataset, batch_size: int, mega_batch_mult: int = 40, shuffle: bool = True):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.mega_batch_mult = mega_batch_mult
+        self.shuffle = shuffle
+        self.lengths = dataset.offsets[1:] - dataset.offsets[:-1]
+
+    def __iter__(self):
+        indices = np.random.permutation(len(self.dataset)) if self.shuffle else np.arange(len(self.dataset))
+        mega_batch_size = self.batch_size * self.mega_batch_mult
+        for i in range(0, len(indices), mega_batch_size):
+            mega_batch = indices[i:i + mega_batch_size]
+            mega_batch = mega_batch[np.argsort(self.lengths[mega_batch])]
+            for j in range(0, len(mega_batch), self.batch_size):
+                yield mega_batch[j:j + self.batch_size].tolist()
+
+    def __len__(self):
+        return math.ceil(len(self.dataset) / self.batch_size)
+
+def collate_fn(batch):
+    tokens_list, labels_list, domain_list = zip(*batch)
+    max_len = min(512, max(len(t) for t in tokens_list))
+    
+    padded_tokens = torch.full((len(batch), max_len), PAD_TOKEN_ID, dtype=torch.long)
+    padded_labels = torch.full((len(batch), max_len), -100, dtype=torch.long)
+    
+    for i, (tok, lab) in enumerate(zip(tokens_list, labels_list)):
+        l = min(len(tok), max_len)
+        padded_tokens[i, :l] = tok[:l]
+        padded_labels[i, :l] = lab[:l]
+        
+    return padded_tokens, padded_labels, torch.tensor(domain_list, dtype=torch.long)
+
+# ------------------------------------------------------------------------------
+# 3. Learning Rate Scheduler (Cosine with Warmup)
+# ------------------------------------------------------------------------------
+def get_lr(step, warmup_steps, total_steps, max_lr, min_lr):
+    if step < warmup_steps:
+        return max_lr * (step + 1) / max(1, warmup_steps)
+    if step > total_steps:
+        return min_lr
+    decay_ratio = (step - warmup_steps) / max(1, (total_steps - warmup_steps))
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (max_lr - min_lr)
+
+# ------------------------------------------------------------------------------
+# 4. Live Evaluation Previews
+# ------------------------------------------------------------------------------
+@torch.no_grad()
+def run_live_eval_previews(model, tokenizer, device):
+    if tokenizer is None:
+        return
+    test_cases = [
+        ("Direct EN -> HI",  "<|user|>\nThe sun rises in the east and sets in the west.<|endofturn|>\n<|assistant|>\n"),
+        ("Direct HI -> EN",  "<|user|>\nसूरज पूर्व में उगता है और पश्चिम में डूबता है।<|endofturn|>\n<|assistant|>\n"),
+        ("Command EN -> HI", "<|user|>\nTranslate to Hindi: 'Consistency and discipline are the keys to long term success.'<|endofturn|>\n<|assistant|>\n"),
+        ("Command HI -> EN", "<|user|>\nTranslate to English: 'सफलता का कोई शॉर्टकट नहीं होता, निरंतर प्रयास ही कुंजी है।'<|endofturn|>\n<|assistant|>\n")
+    ]
+    print("\n   💬 --- [LIVE TRANSLATION PREVIEWS] ---")
+    model.eval()
+    for label, prompt in test_cases:
+        ids = torch.tensor([tokenizer.encode(prompt).ids], dtype=torch.long, device=device)
+        out = model.generate(ids, max_new_tokens=45, temperature=0.2, eos_token_id=EOT_ID)
+        gen = tokenizer.decode(out[0][ids.shape[1]:].tolist()).replace("<|endofturn|>", "").strip()
+        print(f"      • [{label:18s}]: \"{gen}\"")
+    model.train()
+
+# ------------------------------------------------------------------------------
+# 5. Cloud Auto-Download Helper
+# ------------------------------------------------------------------------------
+def ensure_dataset_and_base_ckpt(data_dir: str, base_ckpt_path: str, token: str = None):
+    os.makedirs(data_dir, exist_ok=True)
+    shards = [
+        "train_tokens.npy", "train_labels.npy", "train_offsets.npy", "train_domains.npy",
+        "val_tokens.npy", "val_labels.npy", "val_offsets.npy", "val_domains.npy",
+        "metadata.json"
+    ]
+    missing = [s for s in shards if not os.path.exists(os.path.join(data_dir, s))]
+    if missing:
+        print(f"\n🌐 Downloading ViuTranslate-Data shards from Hugging Face Hub (ViuAI/ViuTranslate-Data)...")
+        for s in shards:
+            target = os.path.join(data_dir, s)
+            if not os.path.exists(target):
+                print(f"   • Fetching {s}...")
+                dl = hf_hub_download(repo_id="ViuAI/ViuTranslate-Data", filename=s, repo_type="dataset", token=token)
+                shutil.copy(dl, target)
+        print("   ✅ All dataset shards downloaded.")
+
+    if not os.path.exists(base_ckpt_path):
+        print(f"\n🌐 Base checkpoint not found. Downloading base weights from ViuAI/ViuAI-500M...")
+        os.makedirs(os.path.dirname(base_ckpt_path), exist_ok=True)
+        dl_b = hf_hub_download(repo_id="ViuAI/ViuAI-500M", filename="checkpoints/ckpt_latest.pt", token=token)
+        shutil.copy(dl_b, base_ckpt_path)
+        print("   ✅ Base checkpoint ready.")
+
+# ------------------------------------------------------------------------------
+# 6. Main Training Function
+# ------------------------------------------------------------------------------
+def main():
+    hw = auto_profile_hardware()
+
+    parser = argparse.ArgumentParser(description="ViuTranslate Dedicated Training Engine")
+    parser.add_argument("--data_dir", type=str, default="data/tokenized", help="Tokenized dataset directory")
+    parser.add_argument("--base_ckpt", type=str, default="checkpoints/ckpt_latest.pt", help="Base pretrained weights")
+    parser.add_argument("--output_dir", type=str, default="checkpoints", help="Output directory for trained model")
+    parser.add_argument("--batch_size", type=int, default=None, help="Micro batch size")
+    parser.add_argument("--grad_accum", type=int, default=None, help="Gradient accumulation steps")
+    parser.add_argument("--epochs", type=int, default=3, help="Training epochs (Default: 3)")
+    parser.add_argument("--max_lr", type=float, default=3.2e-5, help="Peak learning rate")
+    parser.add_argument("--min_lr", type=float, default=2.0e-6, help="Min learning rate")
+    parser.add_argument("--neftune_alpha", type=float, default=5.0, help="NEFTune noise scale")
+    parser.add_argument("--eval_interval", type=int, default=250, help="Evaluation interval")
+    parser.add_argument("--push_to_hf", action="store_true", default=False, help="Upload directly to ViuAI/ViuTranslate")
+    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face API token")
+    args = parser.parse_args()
+
+    micro_b = args.batch_size or hw["micro_batch"]
+    grad_acc = args.grad_accum or hw["grad_accum"]
+    eff_batch = micro_b * grad_acc
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("=" * 80)
+    print("🚀 ViuTranslate-500M — Dedicated Neural Translation Training Engine")
+    print(f"   • Device:            {hw['device_name']} ({hw['vram_gb']:.2f} GB VRAM)")
+    print(f"   • Batch Config:      Micro-Batch {micro_b} × Accum {grad_acc} = Effective Batch {eff_batch}")
+    print(f"   • Target Epochs:     {args.epochs}")
+    print(f"   • Precision:         {hw['dtype']}")
+    print(f"   • Target HF Repo:    ViuAI/ViuTranslate")
+    print("=" * 80)
+
+    # Cloud Sync
+    ensure_dataset_and_base_ckpt(args.data_dir, args.base_ckpt, args.hf_token)
+
+    # Load Tokenizer
+    tokenizer = None
+    tok_candidates = ["tokenizer.json", "tokenizer/tokenizer.json", os.path.join(cur_dir, "tokenizer.json")]
+    for tc in tok_candidates:
+        if os.path.exists(tc):
+            try:
+                from tokenizers import Tokenizer
+                tokenizer = Tokenizer.from_file(tc)
+                print(f"✅ Tokenizer loaded successfully ({tokenizer.get_vocab_size():,} vocab)")
+                break
+            except Exception:
+                pass
+
+    # Datasets & Loaders
+    train_ds = TranslationDataset(
+        ids_path=os.path.join(args.data_dir, "train_tokens.npy"),
+        labels_path=os.path.join(args.data_dir, "train_labels.npy"),
+        offsets_path=os.path.join(args.data_dir, "train_offsets.npy"),
+        domains_path=os.path.join(args.data_dir, "train_domains.npy")
+    )
+    val_ds = TranslationDataset(
+        ids_path=os.path.join(args.data_dir, "val_tokens.npy"),
+        labels_path=os.path.join(args.data_dir, "val_labels.npy"),
+        offsets_path=os.path.join(args.data_dir, "val_offsets.npy"),
+        domains_path=os.path.join(args.data_dir, "val_domains.npy")
+    )
+
+    train_sampler = LengthGroupedBatchSampler(train_ds, batch_size=micro_b, shuffle=True)
+    train_loader = DataLoader(train_ds, batch_sampler=train_sampler, collate_fn=collate_fn, num_workers=2, pin_memory=True)
+    val_loader = DataLoader(val_ds, batch_size=micro_b * 2, shuffle=False, collate_fn=collate_fn, num_workers=2)
+
+    # Initialize Model
+    cfg = ViuAIConfig.sft(vocab_size=64003, context_length=2048, neftune_alpha=args.neftune_alpha)
+    model = ViuAI(cfg).to(device)
+
+    print(f"\n📦 Loading base pretrained weights from {args.base_ckpt}...")
+    base_state = torch.load(args.base_ckpt, map_location=device, weights_only=False)
+    weights = base_state.get("model_state_dict", base_state)
+    model.load_state_dict(weights, strict=False)
+    print("✅ Pretrained weights loaded.")
+
+    # Optimizer
+    fused_available = 'fused' in torch.optim.AdamW.__init__.__code__.co_varnames and torch.cuda.is_available()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.max_lr, weight_decay=0.01, betas=(0.9, 0.95), fused=fused_available)
+
+    total_steps = (len(train_loader) // grad_acc) * args.epochs
+    warmup_steps = int(total_steps * 0.04)
+    print(f"📊 Total Optimization Steps: {total_steps:,} | Warmup Steps: {warmup_steps:,}")
+
+    autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=hw["dtype"]) if torch.cuda.is_available() else contextlib.nullcontext()
+
+    # Validation Function
+    @torch.no_grad()
+    def evaluate():
+        model.eval()
+        total_loss, total_tokens = 0.0, 0
+        domain_losses = {k: 0.0 for k in DOMAIN_NAMES.keys()}
+        domain_counts = {k: 0 for k in DOMAIN_NAMES.keys()}
+        
+        for inputs, labels, doms in val_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            with autocast_ctx:
+                logits, loss = model(inputs, targets=labels, pad_id=PAD_TOKEN_ID, shift_labels=True)
+            tok_count = (labels != -100).sum().item()
+            total_loss += loss.item() * tok_count
+            total_tokens += tok_count
+            
+            # Domain-level tracking
+            for d in doms.unique():
+                d_val = d.item()
+                mask = (doms == d)
+                if mask.sum() > 0:
+                    with autocast_ctx:
+                        _, d_l = model(inputs[mask], targets=labels[mask], pad_id=PAD_TOKEN_ID, shift_labels=True)
+                    domain_losses[d_val] += d_l.item()
+                    domain_counts[d_val] += 1
+                    
+        avg_loss = total_loss / max(1, total_tokens)
+        ppl = math.exp(min(20.0, avg_loss))
+        d_summary = {DOMAIN_NAMES[k]: (domain_losses[k] / max(1, domain_counts[k])) for k in DOMAIN_NAMES.keys()}
+        model.train()
+        return avg_loss, ppl, d_summary
+
+    # Training Loop
+    os.makedirs(args.output_dir, exist_ok=True)
+    save_path = os.path.join(args.output_dir, "viutranslate_final.pt")
+    best_val_loss = float("inf")
+    global_step = 0
+    start_time = time.time()
+    total_tokens_trained = 0
+
+    print("\n⚡ Starting Training...")
+    for epoch in range(1, args.epochs + 1):
+        accum_loss = 0.0
+        model.train()
+        for micro_idx, (inputs, labels, _) in enumerate(train_loader):
+            inputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True)
+            active_tokens = (labels != -100).sum().item()
+            total_tokens_trained += active_tokens
+
+            with autocast_ctx:
+                logits, loss = model(inputs, targets=labels, pad_id=PAD_TOKEN_ID, shift_labels=True)
+                loss_scaled = loss / grad_acc
+
+            loss_scaled.backward()
+            accum_loss += loss.item()
+
+            if (micro_idx + 1) % grad_acc == 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                lr = get_lr(global_step, warmup_steps, total_steps, args.max_lr, args.min_lr)
+                for param_group in optimizer.param_groups:
+                    param_group["lr"] = lr
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+
+                step_loss = accum_loss / grad_acc
+                accum_loss = 0.0
+
+                if global_step % 10 == 0 or global_step == 1:
+                    elapsed = time.time() - start_time
+                    tok_s = total_tokens_trained / max(1.0, elapsed)
+                    vram = torch.cuda.memory_allocated() / (1024**3) if torch.cuda.is_available() else 0.0
+                    print(f"Step {global_step:4d}/{total_steps} | Epoch {epoch} | Loss: {step_loss:.4f} | LR: {lr:.2e} | Speed: {tok_s:,.0f} tok/s | VRAM: {vram:.1f}GB")
+
+                if global_step % args.eval_interval == 0:
+                    v_loss, v_ppl, d_losses = evaluate()
+                    print(f"\n🌟 [Eval @ Step {global_step}] Val Loss: {v_loss:.4f} | Perplexity: {v_ppl:.2f}")
+                    print("   📊 Direction Losses: " + " | ".join([f"{k}: {v:.3f}" for k, v in d_losses.items()]))
+                    run_live_eval_previews(model, tokenizer, device)
+                    
+                    if v_loss < best_val_loss:
+                        best_val_loss = v_loss
+                        torch.save({"model_state_dict": model.state_dict(), "global_step": global_step, "val_loss": best_val_loss}, save_path)
+                        print(f"  🏆 Saved New Best Checkpoint -> {save_path}\n")
+
+    # Final Save
+    torch.save({"model_state_dict": model.state_dict(), "global_step": global_step, "best_val_loss": best_val_loss}, save_path)
+    print(f"\n🎉 ViuTranslate Training Complete! Final checkpoint: {save_path}")
+
+    # Direct Push to ViuAI/ViuTranslate
+    if args.push_to_hf:
+        print("\n🚀 Pushing model weights to Hugging Face Model Repository (ViuAI/ViuTranslate)...")
+        token = args.hf_token or os.environ.get("HF_TOKEN")
+        if token:
+            api = HfApi(token=token)
+            api.upload_file(
+                path_or_fileobj=save_path,
+                path_in_repo="viutranslate_final.pt",
+                repo_id="ViuAI/ViuTranslate",
+                repo_type="model"
+            )
+            print("✅ Successfully uploaded viutranslate_final.pt to ViuAI/ViuTranslate!")
+        else:
+            print("⚠️ Skipping upload: No HF_TOKEN provided.")
+
+if __name__ == "__main__":
+    main()