Upload 2 files

code/model.py

@@ -24,10 +24,14 @@ class RMSNorm(nn.Module):
         
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # T4 fp16 me stable rakhne ke liye float32 me norm
-        # Original se same logic, bas efficient
-        return F.rms_norm(x.float(), (x.shape[-1],), self.weight.float(), self.eps).to(x.dtype)
-        # Fallback agar F.rms_norm na ho:
-        # norm_x = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        try:
+            # PyTorch 2.4+ me available
+            if hasattr(F, 'rms_norm'):
+                return F.rms_norm(x.float(), (x.shape[-1],), self.weight.float(), self.eps).to(x.dtype)
+        except Exception:
+            pass
+        norm_x = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (norm_x * self.weight.float()).to(x.dtype)
         # return (norm_x * self.weight.float()).to(x.dtype)
 
 class ResonanceLayerKaggle(nn.Module):

code/train.py

@@ -71,27 +71,34 @@ def load_tokenizer_robust(hf_token: str, model_repo: str, subfolder: str, is_mai
     for fast in [True, False]:
         for p_str in local_candidates:
             p = Path(p_str)
-            if (p / "tokenizer.json").exists() or (p / "vocab.json").exists():
+            if (p / "tokenizer.json").exists() or (p / "vocab.json").exists() or (p / "tokenizer.model").exists():
                 try:
                     tok = AutoTokenizer.from_pretrained(str(p), local_files_only=True, use_fast=fast, trust_remote_code=True)
                     if is_main: logger.info(f"Loaded tokenizer from {p} fast={fast}")
                     return tok
                 except Exception as e:
                     if is_main: logger.warning(f"Local tokenizer {p} fast={fast} failed: {e}")
-                    if "sentencepiece" in str(e).lower() or "tiktoken" in str(e).lower():
+                    continue
+
+        # HF se try - pehle ViuRec tokenizer subfolder, phir alag tokenizer repo
+        hf_repos_to_try = [model_repo, "ViuAI/viuai-500m-tokenizer", "ViuAI/ViuRec"]
+        for repo_try in hf_repos_to_try:
+            for attempt_subfolder in [subfolder, None]:
+                try:
+                    if attempt_subfolder:
+                        tok = AutoTokenizer.from_pretrained(repo_try, subfolder=attempt_subfolder, token=hf_token, use_fast=fast, trust_remote_code=True)
+                    else:
+                        tok = AutoTokenizer.from_pretrained(repo_try, token=hf_token, use_fast=fast, trust_remote_code=True)
+                    if is_main: logger.info(f"Loaded tokenizer from HF {repo_try} subfolder={attempt_subfolder} fast={fast}")
+                    return tok
+                except Exception as e:
+                    err = str(e).lower()
+                    if "sentencepiece" in err or "tiktoken" in err or "protobuf" in err:
+                        if is_main: logger.warning(f"Tokenizer needs sentencepiece/tiktoken/protobuf: {e}")
+                        # fast=False pe dobara try karo
                         continue
-        # HF se try
-        for attempt_subfolder in [subfolder, None]:
-            try:
-                if attempt_subfolder:
-                    tok = AutoTokenizer.from_pretrained(model_repo, subfolder=attempt_subfolder, token=hf_token, use_fast=fast, trust_remote_code=True)
-                else:
-                    tok = AutoTokenizer.from_pretrained(model_repo, token=hf_token, use_fast=fast, trust_remote_code=True)
-                if is_main: logger.info(f"Loaded tokenizer from HF {model_repo} subfolder={attempt_subfolder} fast={fast}")
-                return tok
-            except Exception as e:
-                if is_main: logger.warning(f"HF tokenizer try failed fast={fast} subfolder={attempt_subfolder}: {e}")
-                continue
+                    if is_main: logger.debug(f"HF tokenizer try failed {repo_try} fast={fast} subfolder={attempt_subfolder}: {e}")
+                    continue
 
     if is_main:
         logger.warning("Tokenizer failed, using dummy vocab 64000")
@@ -253,7 +260,28 @@ def main():
     if is_main:
         logger.info(f"Using vocab_size={vocab_size}, ctx_len={args.ctx_len}, device={device}, is_ddp={is_ddp}")
 
-    from model_fixed_T4 import ViuResonance100M, ModelConfig
+    # FIX: Dono naam support karo - model.py ya model_fixed_T4.py
+    try:
+        from model_fixed_T4 import ViuResonance100M, ModelConfig
+    except ModuleNotFoundError:
+        try:
+            from model import ViuResonance100M, ModelConfig
+        except ModuleNotFoundError:
+            # Kaggle me file /kaggle/working/code/model.py ho sakti hai
+            import importlib.util, sys
+            for cand in ["./model_fixed_T4.py", "./model.py", "/kaggle/working/code/model_fixed_T4.py", "/kaggle/working/code/model.py"]:
+                p = Path(cand)
+                if p.exists():
+                    spec = importlib.util.spec_from_file_location("model", str(p))
+                    mod = importlib.util.module_from_spec(spec)
+                    sys.modules["model"] = mod
+                    spec.loader.exec_module(mod)
+                    ViuResonance100M = mod.ViuResonance100M
+                    ModelConfig = mod.ModelConfig
+                    break
+            else:
+                raise ModuleNotFoundError("model.py / model_fixed_T4.py nahi mila - dono files same folder me rakho")
+
     config = ModelConfig(vocab_size=vocab_size, chunk_size=256) # T4 ke liye 256
     unwrapped_model = ViuResonance100M(config).to(device)