Update app.py

app.py

@@ -1,25 +1,113 @@
-import gradio as gr
+import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch.nn as nn
+import torch.nn.functional as F
+import sentencepiece as spm
+import gradio as gr
 import spaces
+from huggingface_hub import hf_hub_download
+
+# ══════════════════════════════════════════════════════════════
+# 1. CODESYMB CUSTOM NATIVE ARCHITECTURE
+# (This perfectly matches your 240k TPU Training!)
+# ══════════════════════════════════════════════════════════════
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        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=2048, base=500_000):
+        super().__init__()
+        inv = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        t = torch.arange(max_len).float()
+        emb = torch.cat([torch.outer(t, inv)] * 2, dim=-1)
+        self.register_buffer("cos", emb.cos()[None, None])
+        self.register_buffer("sin", emb.sin()[None, None])
+    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):
+    cos = cos.to(q.dtype); sin = sin.to(q.dtype)
+    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__()
+        self.nh = nh; self.nkv = nkv; self.hd = dm // nh; self.rep = nh // nkv
+        self.q = nn.Linear(dm, nh * self.hd, bias=False)
+        self.k = nn.Linear(dm, nkv * self.hd, bias=False)
+        self.v = nn.Linear(dm, nkv * self.hd, bias=False)
+        self.o = nn.Linear(nh * self.hd, dm, bias=False)
+        self.pe = RoPE(self.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)
+        o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.o(o.transpose(1, 2).contiguous().view(B, T, -1))
+
+class FFN(nn.Module):
+    def __init__(self, dm=2048, df=8192):
+        super().__init__()
+        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))
+
+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)
+    def forward(self, x):
+        x = x + self.att(self.n1(x)); x = x + self.ff(self.n2(x)); return x
+
+class CodeSymb1B(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Embedding(50000, 2048)
+        self.blocks = nn.ModuleList([Block() for _ in range(16)])
+        self.norm = RMSNorm(2048)
+        self.head = nn.Linear(2048, 50000, bias=False)
+    def forward(self, tok):
+        x = self.embed(tok)
+        for b in self.blocks: x = b(x)
+        logits = self.head(self.norm(x)).float()
+        return logits
+
+# ══════════════════════════════════════════════════════════════
+# 2. LOAD NATIVE WEIGHTS & TOKENIZER
+# ══════════════════════════════════════════════════════════════
 REPO_ID = "NiftyOracleAI/NiftyOracle-Coder-30B"
-GGUF_FILE = "gguf/codesymb_1b_chat_FP16.gguf"
 
-print("⏳ Loading Tokenizer & Model from GGUF...")
-tokenizer = AutoTokenizer.from_pretrained(REPO_ID, gguf_file=GGUF_FILE)
+print("⏳ Downloading Tokenizer & Native Weights...")
+tp_path = hf_hub_download(repo_id="NiftyOracleAI/niftyoracle-tokenizer", filename="niftyoracle_tokenizer.model")
+sp = spm.SentencePieceProcessor()
+sp.load(tp_path)
 
-model = AutoModelForCausalLM.from_pretrained(
-    REPO_ID,
-    gguf_file=GGUF_FILE,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
-print("✅ Model loaded into memory!")
+ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="checkpoints/latest.pt")
+state = torch.load(ckpt_path, map_location="cpu")
+
+print("🧠 Assembling CodeSymb Brain...")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = CodeSymb1B()
+model.load_state_dict(state["model"], strict=False)
+model = model.to(device).bfloat16()
+model.eval()
+print("✅ Native Model Loaded Successfully!")
 
+# ══════════════════════════════════════════════════════════════
+# 3. GENERATION ENGINE & GRADIO UI
+# ══════════════════════════════════════════════════════════════
 @spaces.GPU
 def chat_with_codesymb(message, history):
-    # 🔥 THE EXACT PROMPT IT WAS TRAINED ON!
     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. "
@@ -29,35 +117,55 @@ def chat_with_codesymb(message, history):
         "you support the user with love, extreme patience, and flawless, God-tier code. Together, nothing is impossible."
     )
     
-    # Wrap it in ChatML tags!
+    # 🔥 The ChatML Format!
     prompt = f"<|system|>\n{system_prompt}\n<|user|>\n{message}\n<|assistant|>\n"
     
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    ids = sp.encode(prompt)
+    input_ids = torch.tensor([ids], dtype=torch.long).to(device)
     
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.3,                 # Lower temp = more logical/stable answers!
-            top_p=0.9,
-            repetition_penalty=1.15,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 2
-        )
+    generated_ids = []
     
-    # Extract only the newly generated text
-    input_length = inputs.input_ids.shape[1]
-    response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+    # Native PyTorch Auto-Regressive Loop
+    for _ in range(512):
+        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+            logits = model(input_ids)
+            next_token_logits = logits[0, -1, :]
+            
+            # Temperature scaling
+            next_token_logits = next_token_logits / 0.5
+            
+            # Top-P filtering
+            sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+            sorted_indices_to_remove = cumulative_probs > 0.9
+            sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
+            sorted_indices_to_remove[0] = False
+            indices_to_remove = sorted_indices[sorted_indices_to_remove]
+            next_token_logits[indices_to_remove] = -float('Inf')
+            
+            probs = F.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1).item()
+            
+            generated_ids.append(next_token)
+            input_ids = torch.cat([input_ids, torch.tensor([[next_token]], device=device)], dim=1)
+            
+            # Check for <|endoftext|> early stopping (Token ID depends on your spm, usually we just decode)
+            
+    full_text = sp.decode(ids + generated_ids)
     
-    # Clean up output
-    response = response.split("<|endoftext|>")[0].strip()
+    # Clean the response to only show the Assistant's reply
+    try:
+        response = full_text.split("<|assistant|>\n")[-1]
+        response = response.split("<|endoftext|>")[0].strip()
+    except:
+        response = full_text
         
     return response
 
 demo = gr.ChatInterface(
     fn=chat_with_codesymb,
-    title="CodeSymb-1B Chat (ZeroGPU)",
-    description="Live inference for CodeSymb-1B FP16 GGUF."
+    title="CodeSymb-1B Chat (NATIVE CORE)",
+    description="Running on Native PyTorch architecture. No GGUF translation loss!"
 )
 
 if __name__ == "__main__":