Update app.py

app.py

@@ -1,5 +1,9 @@
 # ================================================================
-#   CodeSymb AI — HuggingFace Space app.py (FIXED INDENTATION)
+#   CodeSymb AI — HuggingFace Space app.py (FIXED MATH!)
+#   ✅ Indentation Fixed
+#   ✅ Negative Logit Penalty Fixed
+#   ✅ Direct Checkpoint Loading
+#   NiftyOracleAI 🇮🇳 | codesymb.com 🌎
 # ================================================================
 
 import gradio as gr
@@ -21,7 +25,7 @@ if HF_TOKEN:
     login(token=HF_TOKEN)
 
 # ══════════════════════════════════════════════════════════════
-# MODEL ARCHITECTURE (must match training exactly)
+# MODEL ARCHITECTURE 
 # ══════════════════════════════════════════════════════════════
 class RMSNorm(nn.Module):
     def __init__(self, d, eps=1e-6):
@@ -33,7 +37,7 @@ class RMSNorm(nn.Module):
         return (x.float() * n).to(x.dtype) * self.w.to(x.dtype)
 
 class RoPE(nn.Module):
-    def __init__(self, dim, max_len=2048, base=500_000):
+    def __init__(self, dim, max_len=4096, base=500_000):
         super().__init__()
         inv = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
         emb = torch.cat([torch.outer(torch.arange(max_len).float(), inv)] * 2, dim=-1)
@@ -115,42 +119,26 @@ tok_path = hf_hub_download(TOK_REPO, "niftyoracle_tokenizer.model", token=HF_TOK
 sp = spm.SentencePieceProcessor()
 sp.load(tok_path)
 VOCAB = sp.get_piece_size()
-print(f"✅ Tokenizer: {VOCAB:,} vocab")
 
 print("⏳ Loading CodeSymb-1B model...")
 model = CodeSymb1B(vocab=VOCAB)
 
-ckpt_found = False
-last_err = None
-for ckpt_name in ["checkpoints/sft_latest.pt",
-                   "checkpoints/latest.pt",
-                   "checkpoints/sft_chat_latest.pt"]:
-    try:
-        ckpt_path = hf_hub_download(REPO_ID, ckpt_name, token=HF_TOKEN or None)
-        print(f"⬇️ Downloaded {ckpt_name} — loading weights...")
-        try:
-            state = torch.load(ckpt_path, map_location="cpu")
-        except Exception:
-            state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-        
-        model_state = state.get("model", state)
-        model_state = {k: v for k, v in model_state.items() if not k.endswith((".pe.cos", ".pe.sin"))}
-        
-        model.load_state_dict(model_state, strict=False)
-        step = state.get("step", 0)
-        print(f"✅ Loaded: {ckpt_name} (step {step:,})")
-        ckpt_found = True
-        break
-    except Exception as e:
-        last_err = e
-        print(f"   ⚠️ {ckpt_name} failed: {type(e).__name__}: {e}")
-
-if not ckpt_found:
-    raise RuntimeError(f"❌ NO CHECKPOINT LOADED! Last error: {last_err}")
+try:
+    ckpt_path = hf_hub_download(REPO_ID, "checkpoints/latest.pt", token=HF_TOKEN or None)
+    print(f"⬇️ Downloaded latest.pt — loading weights...")
+    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    model_state = state.get("model", state)
+    
+    # Safely strip RoPE buffers
+    model_state = {k: v for k, v in model_state.items() if not k.endswith((".pe.cos", ".pe.sin"))}
+    
+    model.load_state_dict(model_state, strict=False)
+    print(f"✅ Loaded: Checkpoint Step {state.get('step', 'Unknown')}")
+except Exception as e:
+    raise RuntimeError(f"❌ NO CHECKPOINT LOADED! Last error: {e}")
 
 model.eval()
-n_params = sum(p.numel() for p in model.parameters())
-print(f"✅ CodeSymb-1B: {n_params/1e9:.2f}B params — READY 🔥")
+print(f"✅ CodeSymb-1B — READY 🔥")
 
 # ══════════════════════════════════════════════════════════════
 # SYSTEM PROMPT
@@ -170,127 +158,81 @@ SYSTEM_PROMPT = (
     "Together, nothing is impossible."
 )
 
-def build_prompt(message: str, history: list) -> list:
-    parts = ["<|system|>\n" + SYSTEM_PROMPT + "\n"]
-    for user_msg, asst_msg in history[-4:]:
-        parts.append(f"<|user|>\n{user_msg}\n<|assistant|>\n{asst_msg}\n")
-    parts.append(f"<|user|>\n{message}\n<|assistant|>\n")
-    return sp.encode("".join(parts))
-
-def tail_matches(generated: list, seq: list) -> bool:
-    seq = list(seq)
-    return len(generated) >= len(seq) and generated[-len(seq):] == seq
-
 @spaces.GPU(duration=60)
 def chat_with_codesymb(message: str, history: list) -> str:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
-    prompt_ids = build_prompt(message, history)
-    if len(prompt_ids) > 1024:
-        user_seq = sp.encode("<|user|>")
-        start = 0
-        for i in range(len(prompt_ids) - 6, 0, -1):
-            if prompt_ids[i:i+len(user_seq)] == list(user_seq):
-                start = i
-                break
-        prompt_ids = prompt_ids[start:][-1024:]
-
+    # Clean ChatML construction
+    prompt = f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{message}\n<|assistant|>\n"
+    prompt_ids = sp.encode(prompt)
+    
+    # Cap prompt size
+    if len(prompt_ids) > 1024: prompt_ids = prompt_ids[-1024:]
     inp = torch.tensor([prompt_ids], dtype=torch.long).to(device)
 
-    END_SEQ  = sp.encode("<|endoftext|>")
-    USER_SEQ = sp.encode("<|user|>")
-    SYS_SEQ  = sp.encode("<|system|>")
-
     generated = []
-    MAX_NEW = 512
-
+    
     with torch.no_grad():
-        for _ in range(MAX_NEW):
+        for _ in range(512):
             with torch.autocast("cuda" if device.type == "cuda" else "cpu", dtype=torch.bfloat16):
                 logits = model(inp)[0, -1, :].float()
 
-            logits = logits / 0.7
+            # Temp scaling
+            logits = logits / 0.5
 
+            # 🔥 THE MATHEMATICALLY CORRECT REPETITION PENALTY 🔥
+            for gid in set(generated[-30:]):
+                if logits[gid] > 0:
+                    logits[gid] /= 1.2
+                else:
+                    logits[gid] *= 1.2  # Multiply negatives to shrink probability!
+
+            # Top-p sampling
             sorted_logits, sorted_idx = torch.sort(logits, descending=True)
             cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-            remove = cum_probs > 0.92
+            remove = cum_probs > 0.90
             remove[..., 1:] = remove[..., :-1].clone()
             remove[..., 0]  = 0
             sorted_logits[remove] = -float("inf")
+            
             logits_filtered = torch.zeros_like(logits)
             logits_filtered.scatter_(0, sorted_idx, sorted_logits)
-
-            if generated:
-                for gid in set(generated[-20:]):
-                    if 0 <= gid < logits_filtered.shape[0]:
-                        logits_filtered[gid] /= 1.3
-
-            probs  = F.softmax(logits_filtered, dim=-1)
+            probs = F.softmax(logits_filtered, dim=-1)
+            
             next_t = torch.multinomial(probs, 1).item()
-
-            if tail_matches(generated, END_SEQ) or tail_matches(generated, USER_SEQ) or tail_matches(generated, SYS_SEQ):
-                break
-
-            if len(generated) > 2 and generated[-1] == generated[-2] == next_t:
-                break
-
             generated.append(next_t)
             inp = torch.cat([inp, torch.tensor([[next_t]], device=device)], dim=1)
+            
+            # Fast early stopping check
+            decoded_so_far = sp.decode(generated)
+            if "<|endoftext|>" in decoded_so_far or "<|user|>" in decoded_so_far:
+                break
 
     response = sp.decode(generated).strip()
 
+    # Strip tags
     for tag in ["<|user|>", "<|system|>", "<|assistant|>", "<|endoftext|>"]:
         response = response.split(tag)[0]
-    response = response.strip()
 
-    if not response:
-        response = (
-            "I'm CodeSymb! 🔥\n"
-            "Built by Dev & Deepakkumar Malkar 🇮🇳\n"
-            "Ask me anything about code!"
-        )
-    return response
+    return response.strip()
 
 # ══════════════════════════════════════════════════════════════
 # GRADIO UI
 # ══════════════════════════════════════════════════════════════
 with gr.Blocks(
     title="CodeSymb AI",
-    theme=gr.themes.Base(primary_hue="blue", secondary_hue="indigo"),
-    css="""
-    body { background: #060D1F !important; }
-    .gradio-container { max-width: 860px !important; }
-    h1 {
-        text-align: center;
-        background: linear-gradient(135deg, #00D4FF, #A855F7);
-        -webkit-background-clip: text;
-        -webkit-text-fill-color: transparent;
-        font-size: 2.5em !important;
-        font-weight: 900 !important;
-    }
-    .subtitle { text-align: center; color: #A0AEC0; margin-bottom: 20px; font-size: 1.1em; }
-    """,
+    theme=gr.themes.Base(primary_hue="cyan", secondary_hue="blue"),
+    css="body { background: #060D1F !important; color: white !important; } .gradio-container { max-width: 860px !important; }"
 ) as demo:
-
     gr.HTML("""
-    <h1>⚡ CodeSymb AI</h1>
-    <p class='subtitle'>
-        World's Most Powerful Coding AI 🌎<br>
+    <h1 style='text-align:center; color:#00D4FF;'>⚡ CodeSymb AI</h1>
+    <p style='text-align:center; color:#A0AEC0;'>
         Built by <strong>Deepakkumar Malkar</strong> at NiftyOracle 🇮🇳<br>
         <em>"Impossible is just a door with no handle yet invented."</em>
     </p>
     """)
-
-    gr.ChatInterface(
-        fn=chat_with_codesymb,
-        examples=[
-            ["Hi! Who are you and what can you build?"],
-            ["Build me a Python FastAPI with JWT auth"],
-            ["Fix this error: AttributeError: 'NoneType'"],
-            ["Write me a complete Android login screen in Kotlin"],
-        ],
-    )
+    gr.ChatInterface(fn=chat_with_codesymb)
 
 if __name__ == "__main__":
     demo.launch()