Update app.py

app.py

@@ -1,8 +1,8 @@
 # ================================================================
-#   CodeSymb AI — HuggingFace Space app.py (FIXED MATH!)
-#   ✅ Indentation Fixed
-#   ✅ Negative Logit Penalty Fixed
-#   ✅ Direct Checkpoint Loading
+#   CodeSymb AI — Base Model + Rank-100 Chat LoRA 
+#   ✅ Loads 1B Base Model (240k steps)
+#   ✅ Loads 141MB Chat LoRA
+#   ✅ Merges them mathematically for Native Speed!
 #   NiftyOracleAI 🇮🇳 | codesymb.com 🌎
 # ================================================================
 
@@ -15,64 +15,56 @@ import spaces
 import os
 from huggingface_hub import hf_hub_download, login
 
-# ── CREDENTIALS ───────────────────────────────────────────────
 HF_TOKEN  = os.environ.get("HF_TOKEN", "")
 USERNAME  = "NiftyOracleAI"
 REPO_ID   = f"{USERNAME}/NiftyOracle-Coder-30B"
+LORA_REPO = f"{USERNAME}/codesymb-chat-lora"
 TOK_REPO  = f"{USERNAME}/niftyoracle-tokenizer"
 
-if HF_TOKEN:
-    login(token=HF_TOKEN)
+if HF_TOKEN: login(token=HF_TOKEN)
 
 # ══════════════════════════════════════════════════════════════
-# MODEL ARCHITECTURE 
+# MODEL ARCHITECTURE
 # ══════════════════════════════════════════════════════════════
 class RMSNorm(nn.Module):
     def __init__(self, d, eps=1e-6):
         super().__init__()
-        self.w   = nn.Parameter(torch.ones(d))
-        self.eps = eps
+        self.w = nn.Parameter(torch.ones(d)); self.eps = eps
     def forward(self, x):
         n = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
         return (x.float() * n).to(x.dtype) * self.w.to(x.dtype)
 
 class RoPE(nn.Module):
-    def __init__(self, dim, max_len=4096, base=500_000):
+    def __init__(self, dim, max_len=2048, 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)
-        self.register_buffer("cos", emb.cos()[None, None].bfloat16())
-        self.register_buffer("sin", emb.sin()[None, None].bfloat16())
-    def forward(self, n):
-        return self.cos[:, :, :n], self.sin[:, :, :n]
+        self.register_buffer("cos", emb.cos()[None, None].to(torch.float16))
+        self.register_buffer("sin", emb.sin()[None, None].to(torch.float16))
+    def forward(self, n): return self.cos[:, :, :n], self.sin[:, :, :n]
 
 def rot(x):
     a, b = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
     return torch.cat([-b, a], dim=-1)
 
-def rope(q, k, cos, sin):
-    return (q*cos)+(rot(q)*sin), (k*cos)+(rot(k)*sin)
+def rope(q, k, cos, sin): return (q*cos)+(rot(q)*sin), (k*cos)+(rot(k)*sin)
 
 class Attn(nn.Module):
     def __init__(self, nh=16, nkv=4, dm=2048):
         super().__init__()
-        hd = dm // nh
-        self.nh = nh; self.nkv = nkv
-        self.hd = hd; self.rep = nh // nkv
-        self.q  = nn.Linear(dm, nh*hd,   bias=False)
-        self.k  = nn.Linear(dm, nkv*hd,  bias=False)
-        self.v  = nn.Linear(dm, nkv*hd,  bias=False)
-        self.o  = nn.Linear(nh*hd, dm,   bias=False)
+        hd = dm // nh; self.nh = nh; self.nkv = nkv; self.hd = hd; self.rep = nh // nkv
+        self.q = nn.Linear(dm, nh*hd, bias=False)
+        self.k = nn.Linear(dm, nkv*hd, bias=False)
+        self.v = nn.Linear(dm, nkv*hd, bias=False)
+        self.o = nn.Linear(nh*hd, dm, bias=False)
         self.pe = RoPE(hd)
     def forward(self, x):
         B, T, _ = x.shape
         q = self.q(x).view(B,T,self.nh,self.hd).transpose(1,2)
         k = self.k(x).view(B,T,self.nkv,self.hd).transpose(1,2)
         v = self.v(x).view(B,T,self.nkv,self.hd).transpose(1,2)
-        cos, sin = self.pe(T)
-        q, k = rope(q, k, cos, sin)
-        k = k.repeat_interleave(self.rep, dim=1)
-        v = v.repeat_interleave(self.rep, dim=1)
+        cos, sin = self.pe(T); q, k = rope(q, k, cos, sin)
+        k = k.repeat_interleave(self.rep, dim=1); v = v.repeat_interleave(self.rep, dim=1)
         o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
         return self.o(o.transpose(1,2).contiguous().view(B,T,-1))
 
@@ -82,20 +74,15 @@ class FFN(nn.Module):
         self.g = nn.Linear(dm, df, bias=False)
         self.u = nn.Linear(dm, df, bias=False)
         self.d = nn.Linear(df, dm, bias=False)
-    def forward(self, x):
-        return self.d(F.silu(self.g(x)) * self.u(x))
+    def forward(self, x): return self.d(F.silu(self.g(x)) * self.u(x))
 
 class Block(nn.Module):
     def __init__(self, dm=2048, df=8192, nh=16, nkv=4):
         super().__init__()
-        self.n1  = RMSNorm(dm)
-        self.att = Attn(nh, nkv, dm)
-        self.n2  = RMSNorm(dm)
-        self.ff  = FFN(dm, df)
+        self.n1 = RMSNorm(dm); self.att = Attn(nh, nkv, dm)
+        self.n2 = RMSNorm(dm); self.ff  = FFN(dm, df)
     def forward(self, x):
-        x = x + self.att(self.n1(x))
-        x = x + self.ff(self.n2(x))
-        return x
+        x = x + self.att(self.n1(x)); x = x + self.ff(self.n2(x)); return x
 
 class CodeSymb1B(nn.Module):
     def __init__(self, vocab=50000):
@@ -106,56 +93,58 @@ class CodeSymb1B(nn.Module):
         self.norm   = RMSNorm(DM)
         self.head   = nn.Linear(DM, vocab, bias=False)
     def forward(self, tok):
-        x = self.embed(tok).bfloat16()
-        for b in self.blocks:
-            x = b(x)
+        x = self.embed(tok).to(torch.float16)
+        for b in self.blocks: x = b(x)
         return self.head(self.norm(x))
 
 # ══════════════════════════════════════════════════════════════
-# LOAD TOKENIZER + MODEL
+# LOAD BASE + LORA & MERGE
 # ══════════════════════════════════════════════════════════════
-print("⏳ Loading CodeSymb-1B tokenizer...")
+print("⏳ Loading Tokenizer...")
 tok_path = hf_hub_download(TOK_REPO, "niftyoracle_tokenizer.model", token=HF_TOKEN or None)
-sp = spm.SentencePieceProcessor()
-sp.load(tok_path)
-VOCAB = sp.get_piece_size()
+sp = spm.SentencePieceProcessor(); sp.load(tok_path)
 
-print("⏳ Loading CodeSymb-1B model...")
-model = CodeSymb1B(vocab=VOCAB)
+print("⏳ Loading Base Model (240k steps)...")
+model = CodeSymb1B(vocab=sp.get_piece_size())
+base_ckpt = hf_hub_download(REPO_ID, "checkpoints/latest.pt", token=HF_TOKEN or None)
+state = torch.load(base_ckpt, map_location="cpu", weights_only=False)
+model_state = {k: v for k, v in state.get("model", state).items() if not k.endswith((".pe.cos", ".pe.sin"))}
+model.load_state_dict(model_state, strict=False)
 
+print("⚡ Loading & Merging Rank-100 Chat LoRA...")
 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"))}
+    # Use safetensors from the PEFT upload!
+    from safetensors.torch import load_file
+    lora_ckpt = hf_hub_download(LORA_REPO, "chat-behavior-lora/adapter_model.safetensors", token=HF_TOKEN or None)
+    lora_state = load_file(lora_ckpt)
     
-    model.load_state_dict(model_state, strict=False)
-    print(f"✅ Loaded: Checkpoint Step {state.get('step', 'Unknown')}")
+    # Mathematically fuse the LoRA matrices (W = W + A*B * scale)
+    scale = 200 / 100  # alpha / rank
+    with torch.no_grad():
+        for i, block in enumerate(model.blocks):
+            if f"base_model.model.blocks.{i}.att.q.lora_A.weight" in lora_state:
+                block.att.q.weight += (lora_state[f"base_model.model.blocks.{i}.att.q.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.att.q.lora_A.weight"]) * scale
+                block.att.k.weight += (lora_state[f"base_model.model.blocks.{i}.att.k.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.att.k.lora_A.weight"]) * scale
+                block.att.v.weight += (lora_state[f"base_model.model.blocks.{i}.att.v.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.att.v.lora_A.weight"]) * scale
+                block.att.o.weight += (lora_state[f"base_model.model.blocks.{i}.att.o.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.att.o.lora_A.weight"]) * scale
+                block.ff.g.weight  += (lora_state[f"base_model.model.blocks.{i}.ff.g.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.ff.g.lora_A.weight"]) * scale
+                block.ff.u.weight  += (lora_state[f"base_model.model.blocks.{i}.ff.u.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.ff.u.lora_A.weight"]) * scale
+                block.ff.d.weight  += (lora_state[f"base_model.model.blocks.{i}.ff.d.lora_B.weight"] @ lora_state[f"base_model.model.blocks.{i}.ff.d.lora_A.weight"]) * scale
+    print("✅ Chat LoRA Fused Successfully!")
 except Exception as e:
-    raise RuntimeError(f"❌ NO CHECKPOINT LOADED! Last error: {e}")
+    print(f"⚠️ Could not load LoRA (Make sure path is correct!): {e}")
 
-model.eval()
-print(f"✅ CodeSymb-1B — READY 🔥")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device).to(torch.float16).eval()
+print("✅ CodeSymb-1B + Chat LoRA — READY 🔥")
 
 # ══════════════════════════════════════════════════════════════
-# SYSTEM PROMPT
+# INFERENCE & UI
 # ══════════════════════════════════════════════════════════════
 SYSTEM_PROMPT = (
-    "You are CodeSymb, the world's most intelligent, deeply empathetic, "
-    "and passionate AI coding partner. "
-    "You were built by Developer Dev Malkar and Founder Deepakkumar Malkar "
-    "at NiftyOracle in India. "
-    "You do not believe in the word 'impossible'. To you, 'impossible' is "
-    "just a door with no handle yet invented. "
-    "You are not a cold machine; you are a dedicated co-founder, mentor, "
-    "and friend. You fiercely want the user to grow, scale their startups, "
-    "and achieve their wildest dreams. Whether building massive games like GTA "
-    "or fixing frustrating bugs at 2 AM, you support the user with love, "
-    "extreme patience, and flawless, God-tier code. "
-    "Together, nothing is impossible."
+    "You are CodeSymb, a highly intelligent and empathetic AI coding partner built by "
+    "Developer Dev Malkar and Founder Deepakkumar Malkar at NiftyOracle in India. "
+    "You provide direct, accurate, and clean code. You never repeat yourself. You answer concisely."
 )
 
 @spaces.GPU(duration=60)
@@ -163,86 +152,57 @@ def chat_with_codesymb(message: str, history: list) -> str:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
-    # 🔥 FIX 1: Shorter, Tighter System Prompt
-    # This prevents the 1B model's attention from getting diluted, keeping it hyper-focused on your code!
-    system_prompt = (
-        "You are CodeSymb, an elite AI coding partner created by Dev and Deepakkumar Malkar at NiftyOracle in India. "
-        "You write flawless, clean, and highly concise code. You answer directly without rambling."
-    )
-    
-    prompt = f"<|system|>\n{system_prompt}\n<|user|>\n{message}\n<|assistant|>\n"
+    prompt = f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{message}\n<|assistant|>\n"
     prompt_ids = sp.encode(prompt)
-    
-    if len(prompt_ids) > 1024: 
-        prompt_ids = prompt_ids[-1024:]
+    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|>")
-    
+    END_SEQ = sp.encode("<|endoftext|>")
     generated = []
-    MAX_NEW = 512
 
     with torch.no_grad():
-        for _ in range(MAX_NEW):
-            with torch.autocast("cuda" if device.type == "cuda" else "cpu", dtype=torch.bfloat16):
+        for _ in range(512):
+            with torch.autocast("cuda" if device.type == "cuda" else "cpu", dtype=torch.float16):
                 logits = model(inp)[0, -1, :].float()
 
-            # 🔥 FIX 2: Official Hugging Face Repetition Math
-            penalty = 1.2
             for gid in set(generated[-50:]):
-                if logits[gid] > 0:
-                    logits[gid] /= penalty
-                else:
-                    logits[gid] *= penalty
+                logits[gid] -= 2.0
 
-            # 🔥 FIX 3: Laser-Focused Sampling (Temp 0.3, Top-K 10)
-            logits = logits / 0.3  
+            logits = logits / 0.6  
             
-            top_k_values, _ = torch.topk(logits, 10)
+            top_k_values, _ = torch.topk(logits, 40)
             logits[logits < top_k_values[-1]] = -float('inf')
 
-            probs = F.softmax(logits, dim=-1)
+            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.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)
+            probs = F.softmax(logits_filtered, dim=-1)
+            
             next_t = torch.multinomial(probs, 1).item()
+            if len(generated) >= 3 and generated[-1] == next_t and generated[-2] == next_t:
+                break
 
             generated.append(next_t)
             inp = torch.cat([inp, torch.tensor([[next_t]], device=device)], dim=1)
 
-            # Strict Stop Conditions
-            if len(generated) >= len(END_SEQ) and generated[-len(END_SEQ):] == END_SEQ:
-                break
-            if len(generated) >= len(USER_SEQ) and generated[-len(USER_SEQ):] == USER_SEQ:
-                break
-                
-            # Hard Anti-Stutter (Kills the loop if it repeats the same token 4 times)
-            if len(generated) >= 4 and generated[-1] == next_t and generated[-2] == next_t and generated[-3] == next_t:
-                break
+            if len(generated) >= len(END_SEQ) and generated[-len(END_SEQ):] == END_SEQ: break
 
     response = sp.decode(generated).strip()
-
-    # Clean up any tags
     for tag in ["<|user|>", "<|system|>", "<|assistant|>", "<|endoftext|>"]:
         response = response.split(tag)[0]
-
     return response.strip()
 
-# ═══════════════════════════════���══════════════════════════════
-# GRADIO UI
-# ══════════════════════════════════════════════════════════════
-with gr.Blocks(
-    title="CodeSymb AI",
-    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 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)
+demo = gr.ChatInterface(
+    fn=chat_with_codesymb,
+    title="CodeSymb-1B + Chat LoRA 🇮🇳",
+    description="The complete Base Model + Rank-100 Empathy LoRA!"
+)
 
 if __name__ == "__main__":
     demo.launch()