Create app.py

app.py

@@ -0,0 +1,351 @@
+# -------------------------------
+# app.py
+#
+# This file contains the backend logic and Gradio UI for the chatbot.
+# Now using Sam-3.0-3 from Smilyai-labs/Sam-3.0-3 — a model that thinks, reasons, and responds with clarity.
+# -------------------------------
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pathlib import Path
+from safetensors.torch import load_file, safe_open
+from transformers import AutoTokenizer
+from dataclasses import dataclass
+import gradio as gr
+import os
+from huggingface_hub import hf_hub_download
+
+# -------------------------------
+# 1) Sam-3.0-3 Architecture (from your second code)
+# -------------------------------
+@dataclass
+class Sam3Config:
+    vocab_size: int = 50257
+    d_model: int = 384
+    n_layers: int = 10
+    n_heads: int = 6
+    ff_mult: float = 4.0
+    dropout: float = 0.1
+    input_modality: str = "text"
+    head_type: str = "causal_lm"
+    version: str = "0.1"
+
+    def __init__(self, vocab_size=50257, d_model=384, n_layers=10, n_heads=6, ff_mult=4.0, dropout=0.1, input_modality="text", head_type="causal_lm", version="0.1", **kwargs):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_mult = ff_mult
+        self.dropout = dropout
+        self.input_modality = input_modality
+        self.head_type = head_type
+        self.version = version
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+    def forward(self, x):
+        return self.weight * x * (x.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()
+
+class MHA(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.0):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+    def forward(self, x, attn_mask=None):
+        B, T, C = x.shape
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        causal = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
+        scores = scores.masked_fill(causal, float("-inf"))
+        if attn_mask is not None:
+            scores = scores.masked_fill(~attn_mask.unsqueeze(1).unsqueeze(2).bool(), float("-inf"))
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(self.dropout(attn), v).transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(out)
+
+class SwiGLU(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.w1 = nn.Linear(d_model, d_ff, bias=False)
+        self.w2 = nn.Linear(d_model, d_ff, bias=False)
+        self.w3 = nn.Linear(d_ff, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+    def forward(self, x):
+        return self.w3(self.dropout(torch.nn.functional.silu(self.w1(x)) * self.w2(x)))
+
+class Block(nn.Module):
+    def __init__(self, d_model, n_heads, ff_mult, dropout=0.0):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.attn = MHA(d_model, n_heads, dropout=dropout)
+        self.norm2 = RMSNorm(d_model)
+        self.ff = SwiGLU(d_model, int(ff_mult * d_model), dropout=dropout)
+        self.drop = nn.Dropout(dropout)
+    def forward(self, x, attn_mask=None):
+        x = x + self.drop(self.attn(self.norm1(x), attn_mask=attn_mask))
+        x = x + self.drop(self.ff(self.norm2(x)))
+        return x
+
+class Sam3(nn.Module):
+    def __init__(self, config: Sam3Config):
+        super().__init__()
+        self.config = config
+        self.embed = nn.Embedding(config.vocab_size, config.d_model)
+        self.blocks = nn.ModuleList([Block(config.d_model, config.n_heads, config.ff_mult, dropout=config.dropout) for _ in range(config.n_layers)])
+        self.norm = RMSNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.embed.weight
+    def forward(self, input_ids, attention_mask=None):
+        x = self.embed(input_ids)
+        for blk in self.blocks:
+            x = blk(x, attn_mask=attention_mask)
+        x = self.norm(x)
+        return self.lm_head(x)
+
+# -------------------------------
+# 2) Load tokenizer & special tokens (Sam-3.0-3 style)
+# -------------------------------
+SPECIAL_TOKENS = {
+    "bos": "<|bos|>",
+    "eot": "<|eot|>",
+    "user": "<|user|>",
+    "assistant": "<|assistant|>",
+    "system": "<|system|>",
+    "think": "<|think|>",
+}
+
+# Use GPT-2 tokenizer and add special tokens
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.add_special_tokens({"additional_special_tokens": list(SPECIAL_TOKENS.values())})
+
+EOT_ID = SPECIAL_TOKENS["eot"]
+EOT_ID = tokenizer.convert_tokens_to_ids(EOT_ID) or tokenizer.eos_token_id
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# -------------------------------
+# 3) Download model weights from Hugging Face Hub
+# -------------------------------
+hf_repo = "Smilyai-labs/Sam-3.0-3"
+weights_filename = "model.safetensors"
+
+print(f"Loading model '{hf_repo}' from Hugging Face Hub...")
+
+try:
+    # Download weights
+    weights_path = hf_hub_download(repo_id=hf_repo, filename=weights_filename)
+    print(f"✅ Downloaded weights to: {weights_path}")
+
+    # Verify file size
+    if not os.path.exists(weights_path):
+        raise FileNotFoundError(f"Downloaded file not found at {weights_path}")
+    file_size = os.path.getsize(weights_path)
+    print(f"📄 File size: {file_size} bytes")
+
+except Exception as e:
+    raise RuntimeError(f"❌ Failed to download model weights: {e}")
+
+# Initialize model with correct vocab size
+cfg = Sam3Config(vocab_size=len(tokenizer))
+model = Sam3(cfg).to(device)
+
+# Load state dict safely
+print("Loading state dict...")
+try:
+    # Try safe_open first (preferred)
+    state_dict = {}
+    with safe_open(weights_path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            state_dict[key] = f.get_tensor(key)
+    print("✅ Loaded via safe_open")
+
+except Exception as e:
+    print(f"⚠️ safe_open failed: {e}. Falling back to torch.load...")
+    try:
+        state_dict = torch.load(weights_path, map_location="cpu")
+        print("✅ Loaded via torch.load")
+    except Exception as torch_e:
+        raise RuntimeError(f"❌ Could not load model weights: {torch_e}")
+
+# Filter state_dict to match model keys
+model_state_dict = model.state_dict()
+filtered_state_dict = {k: v for k, v in state_dict.items() if k in model_state_dict}
+
+# Warn about missing/extra keys
+missing_keys = set(model_state_dict.keys()) - set(filtered_state_dict.keys())
+extra_keys = set(state_dict.keys()) - set(model_state_dict.keys())
+if missing_keys:
+    print(f"⚠️ Missing keys in loaded state dict: {missing_keys}")
+if extra_keys:
+    print(f"⚠️ Extra keys in loaded state dict: {extra_keys}")
+
+model.load_state_dict(filtered_state_dict, strict=False)
+model.eval()
+print("✅ Model loaded successfully!")
+
+# -------------------------------
+# 4) Sampling function (unchanged from Sam-3.0-3 code)
+# -------------------------------
+def sample_next_token(
+    logits,
+    past_tokens,
+    temperature=0.8,
+    top_k=60,
+    top_p=0.9,
+    repetition_penalty=1.1,
+    max_repeat=5,
+    no_repeat_ngram_size=3
+):
+    if logits.dim() == 3:
+        logits = logits[:, -1, :].clone()
+    else:
+        logits = logits.clone()
+    batch_size, vocab_size = logits.size(0), logits.size(1)
+    orig_logits = logits.clone()
+
+    if temperature != 1.0:
+        logits = logits / float(temperature)
+
+    past_list = past_tokens.tolist() if isinstance(past_tokens, torch.Tensor) else list(past_tokens)
+
+    for token_id in set(past_list):
+        if 0 <= token_id < vocab_size:
+            logits[:, token_id] /= repetition_penalty
+
+    if len(past_list) >= max_repeat:
+        last_token = past_list[-1]
+        count = 1
+        for i in reversed(past_list[:-1]):
+            if i == last_token:
+                count += 1
+            else:
+                break
+        if count >= max_repeat:
+            if 0 <= last_token < vocab_size:
+                logits[:, last_token] = -float("inf")
+
+    if no_repeat_ngram_size > 0 and len(past_list) >= no_repeat_ngram_size:
+        for i in range(len(past_list) - no_repeat_ngram_size + 1):
+            ngram = tuple(past_list[i : i + no_repeat_ngram_size])
+            if len(past_list) >= no_repeat_ngram_size - 1:
+                prefix = tuple(past_list[-(no_repeat_ngram_size - 1):])
+                for token_id in range(vocab_size):
+                    if tuple(list(prefix) + [token_id]) == ngram and 0 <= token_id < vocab_size:
+                        logits[:, token_id] = -float("inf")
+
+    if top_k is not None and top_k > 0:
+        tk = min(max(1, int(top_k)), vocab_size)
+        topk_vals, topk_indices = torch.topk(logits, tk, dim=-1)
+        min_topk = topk_vals[:, -1].unsqueeze(-1)
+        logits[logits < min_topk] = -float("inf")
+
+    if top_p is not None and 0.0 < top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        sorted_probs = F.softmax(sorted_logits, dim=-1)
+        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+        for b in range(batch_size):
+            sorted_mask = cumulative_probs[b] > top_p
+            if sorted_mask.numel() > 0:
+                sorted_mask[0] = False
+                tokens_to_remove = sorted_indices[b][sorted_mask]
+                logits[b, tokens_to_remove] = -float("inf")
+
+    for b in range(batch_size):
+        if torch.isneginf(logits[b]).all():
+            logits[b] = orig_logits[b]
+
+    probs = F.softmax(logits, dim=-1)
+    if torch.isnan(probs).any():
+        probs = torch.ones_like(logits) / logits.size(1)
+
+    next_token = torch.multinomial(probs, num_samples=1)
+    return next_token.to(device)
+
+# -------------------------------
+# 5) Gradio Chat UI and API Logic (Updated with truthful, compelling UI)
+# -------------------------------
+SPECIAL_TOKENS_CHAT = {"bos": "<|bos|>", "eot": "<|eot|>", "user": "<|user|>", "assistant": "<|assistant|>", "system": "<|system|>"}
+
+def predict(message, history):
+    # Construct the chat history with special tokens
+    chat_history = []
+    for human, assistant in history:
+        chat_history.append(f"{SPECIAL_TOKENS_CHAT['user']} {human} {SPECIAL_TOKENS_CHAT['eot']}")
+        if assistant:
+            chat_history.append(f"{SPECIAL_TOKENS_CHAT['assistant']} {assistant} {SPECIAL_TOKENS_CHAT['eot']}")
+    
+    chat_history.append(f"{SPECIAL_TOKENS_CHAT['user']} {message} {SPECIAL_TOKENS_CHAT['eot']}")
+    
+    system_prompt = "You are Sam-3, an advanced reasoning AI. You think step by step, analyze deeply, and answer with precision. You do not guess — you deduce. Avoid medical or legal advice."
+    prompt = f"{SPECIAL_TOKENS_CHAT['system']} {system_prompt} {SPECIAL_TOKENS_CHAT['eot']}\n" + "\n".join(chat_history) + f"\n{SPECIAL_TOKENS_CHAT['assistant']}"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    
+    generated_text = ""
+    for _ in range(256):
+        with torch.no_grad():
+            logits = model(input_ids, attention_mask=attention_mask)
+        next_token = sample_next_token(logits, input_ids[0], temperature=0.4, top_k=50, top_p=0.9, repetition_penalty=1.1)
+        
+        token_id = int(next_token.squeeze().item())
+        token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+        
+        input_ids = torch.cat([input_ids, next_token], dim=1)
+        attention_mask = torch.cat([attention_mask, torch.ones((attention_mask.size(0), 1), device=device, dtype=attention_mask.dtype)], dim=1)
+        
+        generated_text += token_str
+        yield generated_text
+        
+        if token_id == EOT_ID:
+            break
+
+# Gradio Interface — Now Truthfully Representing the Model’s Capabilities
+demo = gr.ChatInterface(
+    fn=predict,
+    title="🌟 Sam-3: The Reasoning AI",
+    description="""
+    Sam-3 is not just a language model — it **thinks before it speaks**.  
+    Built with deep architectural integrity, it analyzes problems step-by-step, uncovers hidden patterns, and delivers precise, logical answers.  
+    No fluff. No guessing. Just reasoning.
+    
+    Try asking it:  
+    → “If I have 3 apples and give away half of them, then buy 5 more, how many do I have?”  
+    → “Explain quantum entanglement like I’m 10.”  
+    → “What’s the flaw in this argument: ‘All birds fly; penguins are birds; therefore penguins can fly’?”
+    """,
+    theme=gr.themes.Soft(
+        primary_hue="indigo",
+        secondary_hue="blue"
+    ),
+    chatbot=gr.Chatbot(
+        label="Sam-3 🤔",
+        bubble_full_width=False,
+        height=600,
+    ),
+    examples=[
+        "What is the capital of France?",
+        "Explain why the sky is blue.",
+        "If a train leaves at 2 PM going 60 mph, and another leaves 30 minutes later at 80 mph, when does the second catch up?",
+        "What are the ethical implications of AI making medical diagnoses?"
+    ],
+    cache_examples=False
+).launch(
+    show_api=True
+)