Spaces:
Sleeping
Sleeping
| # ================================================================ | |
| # CodeSymb AI — THE GRAND AWAKENING 🔥 | |
| # ✅ Loads 240k Base Model | |
| # ✅ Loads 50,000-Step SFT Chat LoRA (Just finished!) | |
| # ✅ Mathematically Fuses for Native Speed | |
| # NiftyOracleAI 🇮🇳 | codesymb.com 🌎 | |
| # ================================================================ | |
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import sentencepiece as spm | |
| import spaces | |
| import os | |
| from huggingface_hub import hf_hub_download, login | |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") | |
| USERNAME = "NiftyOracleAI" | |
| REPO_ID = f"{USERNAME}/NiftyOracle-Coder-30B" | |
| TOK_REPO = f"{USERNAME}/niftyoracle-tokenizer" | |
| if HF_TOKEN: login(token=HF_TOKEN) | |
| # ══════════════════════════════════════════════════════════════ | |
| # CORE ARCHITECTURE (Fixed 2048 RoPE) | |
| # ══════════════════════════════════════════════════════════════ | |
| 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)) | |
| emb = torch.cat([torch.outer(torch.arange(max_len).float(), inv)] * 2, dim=-1) | |
| 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) | |
| 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) | |
| 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) | |
| 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, vocab=50000): | |
| super().__init__() | |
| DM=2048; DF=8192; NL=16; NH=16; NKV=4 | |
| self.embed = nn.Embedding(vocab, DM) | |
| self.blocks = nn.ModuleList([Block(DM,DF,NH,NKV) for _ in range(NL)]) | |
| self.norm = RMSNorm(DM) | |
| self.head = nn.Linear(DM, vocab, bias=False) | |
| def forward(self, tok): | |
| x = self.embed(tok).to(torch.float16) | |
| for b in self.blocks: x = b(x) | |
| return self.head(self.norm(x)) | |
| # ══════════════════════════════════════════════════════════════ | |
| # LOAD & FUSE THE 50,000-STEP CHAT LORA! | |
| # ══════════════════════════════════════════════════════════════ | |
| 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) | |
| print("⏳ Loading 240k Base Model...") | |
| 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) | |
| base_weights = {k: v for k, v in state.get("model", state).items() if not k.endswith((".pe.cos", ".pe.sin"))} | |
| model.load_state_dict(base_weights, strict=False) | |
| print("✅ Base Model Loaded.") | |
| print("⚡ FUSING THE NEW 50,000-STEP SFT LORA...") | |
| try: | |
| # 🔥 PULLING THE EXACT FILE YOU JUST UPLOADED FROM LIGHTNING AI! | |
| lora_ckpt = hf_hub_download(REPO_ID, "checkpoints/sft_full_latest.pt", token=HF_TOKEN or None) | |
| lora_state_full = torch.load(lora_ckpt, map_location="cpu", weights_only=False) | |
| lora_state = lora_state_full.get("model", lora_state_full) | |
| scale = 200 / 100 # Alpha(200) / Rank(100) | |
| with torch.no_grad(): | |
| for i, block in enumerate(model.blocks): | |
| for layer_name in ["att.q", "att.k", "att.v", "att.o", "ff.g", "ff.u", "ff.d"]: | |
| A_key = f"blocks.{i}.{layer_name}.A" | |
| B_key = f"blocks.{i}.{layer_name}.B" | |
| # Check for PEFT format keys | |
| peft_A_key = f"base_model.model.blocks.{i}.{layer_name}.lora_A.weight" | |
| peft_B_key = f"base_model.model.blocks.{i}.{layer_name}.lora_B.weight" | |
| A_val = lora_state.get(peft_A_key, lora_state.get(A_key)) | |
| B_val = lora_state.get(peft_B_key, lora_state.get(B_key)) | |
| if A_val is not None and B_val is not None: | |
| A = A_val.float() | |
| B = B_val.float() | |
| delta = (B @ A) * scale if A.shape[1] == 100 else (A @ B).T * scale | |
| target_layer = getattr(block, layer_name.split(".")[0]) | |
| sub_layer = getattr(target_layer, layer_name.split(".")[1]) | |
| sub_layer.weight.data += delta.to(sub_layer.weight.dtype) | |
| print("✅ 50k Chat LoRA Fused! The AI is fully conversational! 🗣️") | |
| except Exception as e: | |
| print(f"⚠️ Failed to fuse Chat LoRA: {e}") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = model.to(device).half().eval() | |
| print(f"✅ CodeSymb is ALIVE on {device}!") | |
| # ══════════════════════════════════════════════════════════════ | |
| # INFERENCE ENGINE | |
| # ══════════════════════════════════════════════════════════════ | |
| def chat_with_codesymb(message: str, history: list) -> str: | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) | |
| # EXACT SAME SYSTEM PROMPT FROM TRAINING! | |
| system_prompt = ( | |
| "You are CodeSymb AI — expert coding assistant built by Deepakkumar Malkar " | |
| "at NiftyOracle India. Write clean code, explain simply, be friendly. " | |
| "Nothing is impossible. codesymb.com" | |
| ) | |
| full_prompt = f"<|system|>\n{system_prompt}\n" | |
| for msg in history[-3:]: | |
| if isinstance(msg, dict): | |
| role = "<|user|>" if msg.get("role") == "user" else "<|assistant|>" | |
| full_prompt += f"{role}\n{msg.get('content', '')}\n" | |
| elif hasattr(msg, "role"): | |
| role = "<|user|>" if msg.role == "user" else "<|assistant|>" | |
| full_prompt += f"{role}\n{msg.content}\n" | |
| else: | |
| full_prompt += f"<|user|>\n{msg[0]}\n<|assistant|>\n{msg[1]}\n" | |
| full_prompt += f"<|user|>\n{message}\n<|assistant|>\n" | |
| prompt_ids = sp.encode(full_prompt) | |
| if len(prompt_ids) > 1024: prompt_ids = prompt_ids[-1024:] | |
| inp = torch.tensor([prompt_ids], dtype=torch.long).to(device) | |
| generated = [] | |
| END_SEQ = sp.encode("<|endoftext|>") | |
| with torch.no_grad(): | |
| for _ in range(512): | |
| with torch.autocast("cuda", dtype=torch.float16): | |
| logits = model(inp)[0, -1, :].float() | |
| for gid in set(generated[-50:]): | |
| logits[gid] -= 2.0 | |
| logits = logits / 0.5 | |
| top_k_values, _ = torch.topk(logits, 40) | |
| logits[logits < top_k_values[-1]] = -float('inf') | |
| 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] = False | |
| 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) | |
| decoded = sp.decode(generated) | |
| if "<|endoftext|>" in decoded or "<|user|>" in decoded: | |
| break | |
| response = sp.decode(generated).strip() | |
| for tag in ["<|user|>", "<|system|>", "<|assistant|>", "<|endoftext|>"]: | |
| response = response.split(tag)[0].strip() | |
| return response if response else "I am CodeSymb! How can I help you? 🔥" | |
| with gr.Blocks(theme=gr.themes.Base(primary_hue="cyan", secondary_hue="blue"), css="body { background: #060D1F !important; color: white !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> | |
| </p> | |
| """) | |
| gr.ChatInterface( | |
| fn=chat_with_codesymb, | |
| examples=[ | |
| ["Hi! Who are you?"], | |
| ["Write a Python function to reverse a string."], | |
| ["What is an API?"], | |
| ] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |