Update code/train_sft.py with user enhancements

code/train_sft.py

@@ -74,7 +74,9 @@ if cur_dir not in sys.path:
 from model import ViuAI, Transformer
 from config import ViuAIConfig, ModelArgs
 
-PAD_TOKEN_ID = 64000
+# PAD = EOT (end-of-turn) used as padding, standard for models without a dedicated
+# pad token. Never use <|user|> (64000) as pad: it pollutes turn-start semantics.
+PAD_TOKEN_ID = 64002
 
 DOMAIN_NAMES_V18 = {
     0: "gk_polity_history",
@@ -308,10 +310,11 @@ def get_lr(it, warmup_steps, total_steps, max_lr, min_lr):
 
 
 
-DEFAULT_HF_TOKEN = "".join(["hf_", "ssyCVhuny", "XxjGdqKp", "VLPpkmWK", "FrrMOIFbg"])
+# SECURITY: HF token must come from environment only. Never hardcode tokens.
+# Set HF_TOKEN env var before running. See README for instructions.
 
 def ensure_cloud_data_and_checkpoint(version: str, data_dir: str, ckpt_path: str, token: str = None):
-    token = token or os.environ.get("HF_TOKEN") or DEFAULT_HF_TOKEN
+    token = token or os.environ.get("HF_TOKEN")
 
     if str(version).lower() == "viutranslate":
         stage_subfolder = "ViuTranslate"
@@ -404,7 +407,9 @@ def run_live_eval_previews(model, tokenizer, device, version=""):
     if tokenizer is None:
         return
 
-    if any(k in str(version).lower() for k in ["translat", "viutranslate"]):
+    _vnorm = str(version).lower().replace("sft_", "").replace("sft", "").strip()
+    _vnum = _vnorm[1:] if _vnorm.startswith("v") and _vnorm[1:].isdigit() else _vnorm
+    if _vnorm in ("translat", "viutranslate") or "translat" in _vnorm:
         test_suite = [
             ("Identity",        "<|user|>\nWho created you and what is your purpose?<|endofturn|>\n<|assistant|>\n", 40),
             ("EN -> HI Direct", "<|user|>\nThe sun rises in the east and sets in the west.<|endofturn|>\n<|assistant|>\n", 45),
@@ -412,7 +417,7 @@ def run_live_eval_previews(model, tokenizer, device, version=""):
             ("Proverb Command", "<|user|>\nTranslate to Hindi: 'Consistency and discipline are the keys to long term success.'<|endofturn|>\n<|assistant|>\n", 50),
             ("Hinglish Command","<|user|>\nTranslate to Hinglish: 'I am waiting outside your office, please call me when you are free.'<|endofturn|>\n<|assistant|>\n", 50)
         ]
-    elif any(k in str(version).lower() for k in ["v2", "sft_v2"]):
+    elif _vnum == "2":
         test_suite = [
             ("Identity",        "<|user|>\nWho created you and what is your name?<|endofturn|>\n<|assistant|>\n", 40),
             ("Indian Greet",    "<|user|>\nRam Ram bhai! Kaise ho?<|endofturn|>\n<|assistant|>\n", 35),
@@ -509,6 +514,10 @@ def main():
     
     ckpt_filename = "viutranslate_final.pt" if str(args.version).lower() == "viutranslate" else f"sft_{args.version}_final.pt"
     save_path = os.path.join(output_dir, ckpt_filename)
+    # Best-checkpoint is kept in a SEPARATE file so later (worse) evals can never
+    # overwrite the best weights. Previously both shared save_path.
+    best_filename = "viutranslate_best.pt" if str(args.version).lower() == "viutranslate" else f"sft_{args.version}_best.pt"
+    best_save_path = os.path.join(output_dir, best_filename)
     
     if args.init_ckpt:
         init_ckpt = args.init_ckpt
@@ -524,16 +533,20 @@ def main():
     tokenizer = None
     tok_path = os.path.join(root_dir, "tokenizer", "tokenizer.json")
     global PAD_TOKEN_ID, EOT_ID
-    PAD_TOKEN_ID = 64000
+    # EOT-as-pad: padding must never reuse <|user|> (turn-start semantics).
+    PAD_TOKEN_ID = 64002
     EOT_ID = 64002
+    MODEL_VOCAB = 64003
     if os.path.exists(tok_path):
         try:
             from tokenizers import Tokenizer
             tokenizer = Tokenizer.from_file(tok_path)
             v = tokenizer.get_vocab()
-            PAD_TOKEN_ID = v.get("<|user|>", 64000)
             EOT_ID = v.get("<|endofturn|>", 64002)
-            print(f"✅ Tokenizer bound dynamically: PAD_TOKEN_ID={PAD_TOKEN_ID}, EOT_ID={EOT_ID}")
+            # EOT-as-pad: padding must never reuse <|user|> (turn-start semantics).
+            PAD_TOKEN_ID = EOT_ID
+            MODEL_VOCAB = tokenizer.get_vocab_size()
+            print(f"✅ Tokenizer bound dynamically: PAD_TOKEN_ID={PAD_TOKEN_ID}, EOT_ID={EOT_ID}, VOCAB={MODEL_VOCAB}")
         except Exception as e:
             print(f"⚠️ Could not load tokenizer for special tokens ({e}). Using defaults.")
 
@@ -553,13 +566,17 @@ def main():
             print(f"⚠️ Notice: Could not parse metadata.json domain_names ({e}). Falling back to version heuristic.")
 
     if not loaded_dynamic_domains:
-        if any(v in str(args.version).lower() for v in ["v2", "2", "sft_v2"]):
+        # Exact normalized match only. Substring matching ("2" in "v26") previously
+        # mis-assigned v21/v22-v26 to the wrong domain map, corrupting per-domain telemetry.
+        vnorm = str(args.version).lower().replace("sft_", "").replace("sft", "").strip()
+        vnum = vnorm[1:] if vnorm.startswith("v") and vnorm[1:].isdigit() else vnorm
+        if vnorm == "viutranslate" or vnum in ("2",):
             DOMAIN_NAMES = DOMAIN_NAMES_V2
-        elif any(v in str(args.version).lower() for v in ["v1", "1", "sft_v1"]):
+        elif vnum in ("1",):
             DOMAIN_NAMES = DOMAIN_NAMES_V1
-        elif any(v in str(args.version).lower() for v in ["22", "23", "24", "25", "26"]):
+        elif vnum in ("22", "23", "24", "25", "26"):
             DOMAIN_NAMES = DOMAIN_NAMES_V22
-        elif any(v in str(args.version).lower() for v in ["19", "20", "21"]):
+        elif vnum in ("19", "20", "21"):
             DOMAIN_NAMES = DOMAIN_NAMES_V19
         else:
             DOMAIN_NAMES = DOMAIN_NAMES_V2
@@ -601,17 +618,34 @@ def main():
             return None
         return cand1
 
-    # 1. Load Data
+    # 1. Load Data (resolve + validate shard paths BEFORE touching np.load so a
+    # wrong --version/--data_dir fails with a clear message, not a cryptic error)
+    _shard_paths = {
+        name: find_shard_path(stage_data_dir, name, args.version)
+        for name in ["train_tokens", "train_labels", "train_offsets",
+                     "val_tokens", "val_labels", "val_offsets"]
+    }
+    for _label, _p in _shard_paths.items():
+        if _p is None or not os.path.exists(_p):
+            try:
+                _have = sorted(os.listdir(stage_data_dir))
+            except OSError:
+                _have = []
+            raise FileNotFoundError(
+                f"Missing required shard '{_label}*' for version '{args.version}' in: {stage_data_dir} "
+                f"(found: {_have}). Check --version/--data_dir, or run the dataset builder first."
+            )
+
     train_dataset = SFTDataset(
-        ids_path=find_shard_path(stage_data_dir, "train_tokens", args.version),
-        labels_path=find_shard_path(stage_data_dir, "train_labels", args.version),
-        offsets_path=find_shard_path(stage_data_dir, "train_offsets", args.version),
+        ids_path=_shard_paths["train_tokens"],
+        labels_path=_shard_paths["train_labels"],
+        offsets_path=_shard_paths["train_offsets"],
         domains_path=find_shard_path(stage_data_dir, "train_domains", args.version)
     )
     val_dataset = SFTDataset(
-        ids_path=find_shard_path(stage_data_dir, "val_tokens", args.version),
-        labels_path=find_shard_path(stage_data_dir, "val_labels", args.version),
-        offsets_path=find_shard_path(stage_data_dir, "val_offsets", args.version),
+        ids_path=_shard_paths["val_tokens"],
+        labels_path=_shard_paths["val_labels"],
+        offsets_path=_shard_paths["val_offsets"],
         domains_path=find_shard_path(stage_data_dir, "val_domains", args.version)
     )
 
@@ -646,7 +680,7 @@ def main():
         print("💾 Activation Checkpointing: ENABLED (low-VRAM mode)")
 
     model_args = ViuAIConfig(
-        vocab_size=64003,
+        vocab_size=MODEL_VOCAB,
         context_length=2048,
         z_loss_weight=0.0,
         attn_dropout=0.05,
@@ -766,13 +800,20 @@ def main():
             inps = inps.to(device, non_blocking=True)
             lbls = lbls.to(device, non_blocking=True)
             with autocast_ctx:
-                # Single forward pass: get hidden states
+                # Single forward pass: get hidden states (with pad masking)
                 B, T = inps.shape
                 x = raw_m.tok_emb(inps)
                 cos = raw_m.rope_cos.to(x.device)
                 sin = raw_m.rope_sin.to(x.device)
+                eval_mask = None
+                if (inps == PAD_TOKEN_ID).any():
+                    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
+                    key_mask = (inps != PAD_TOKEN_ID).unsqueeze(1).unsqueeze(2)
+                    bool_mask = causal.unsqueeze(0).unsqueeze(0) & key_mask
+                    eval_mask = torch.zeros(B, 1, T, T, dtype=x.dtype, device=x.device)
+                    eval_mask = eval_mask.masked_fill(~bool_mask, float("-inf"))
                 for block in raw_m.blocks:
-                    x = block(x, cos, sin, None)
+                    x = block(x, cos, sin, eval_mask)
                 x = raw_m.final_norm(x)
                 hidden = x[:, :-1, :].contiguous()   # [B, T-1, D]
                 shift_labels = lbls[:, 1:].contiguous()  # [B, T-1]
@@ -855,7 +896,7 @@ def main():
             total_tokens_trained += active_tokens_count
             
             with autocast_ctx:
-                _, loss = model(inputs, targets=labels, pad_id=None, shift_labels=True)
+                _, loss = model(inputs, targets=labels, pad_id=PAD_TOKEN_ID, shift_labels=True)
                 loss_scaled = loss / grad_acc
 
             loss_scaled.backward()
@@ -925,7 +966,10 @@ def main():
                         "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
                     }
                     torch.save(save_payload, save_path)
-                    print(f"  💾 Saved checkpoint -> {save_path}\n")
+                    print(f"  💾 Saved checkpoint -> {save_path}")
+                    if is_best:
+                        torch.save(save_payload, best_save_path)
+                        print(f"  🏆 Saved BEST checkpoint -> {best_save_path}")
 
         # End of Epoch Handling
         if epoch_batches > 0 and (micro_idx + 1) % grad_acc != 0:
@@ -969,11 +1013,13 @@ def main():
     }
     torch.save(final_payload, save_path)
     print(f"💾 Final master checkpoint saved to: {save_path}")
+    if os.path.exists(best_save_path):
+        print(f"🏆 Best checkpoint (lowest val loss) kept at: {best_save_path}")
 
     # Push to Hugging Face Hub (Only Final Checkpoint)
     if args.push_to_hf:
         print("\n🚀 Pushing Final Checkpoint to Hugging Face Model Hub (ViuAI/ViuAI-500M)...")
-        token = args.hf_token or os.environ.get("HF_TOKEN") or DEFAULT_HF_TOKEN
+        token = args.hf_token or os.environ.get("HF_TOKEN")
         if token:
             try:
                 api = HfApi(token=token)