Create app.py

app.py

@@ -0,0 +1,370 @@
+# -------------------------------
+# app.py — Sam-3.5: The Reasoning AI (Updated Architecture)
+# -------------------------------
+
+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
+from transformers import AutoTokenizer
+from dataclasses import dataclass
+from typing import Dict, List
+import gradio as gr
+import os
+from huggingface_hub import hf_hub_download
+import json
+
+# -------------------------------
+# 1) Configuration & Special Tokens
+# -------------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+SPECIAL_TOKENS = {
+    "bos": "<|bos|>",
+    "eot": "<|eot|>",
+    "user": "<|user|>",
+    "assistant": "<|assistant|>",
+    "system": "<|system|>",
+    "think": "<|think|>",  # Keep this for reasoning display
+}
+
+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())})
+
+SPECIAL_IDS = {k: tokenizer.convert_tokens_to_ids(v) for k, v in SPECIAL_TOKENS.items()}
+EOT_ID = SPECIAL_IDS.get("eot", tokenizer.eos_token_id)
+THINK_ID = SPECIAL_IDS.get("think")
+assert THINK_ID is not None, "Tokenizer must include <|think|> token"
+
+MAX_LENGTH = 1024
+
+# -------------------------------
+# 2) Model Architecture (Sam-3.5)
+# -------------------------------
+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__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        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)
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=True,
+            dropout_p=self.dropout.p if self.training else 0.0
+        )
+        return self.out_proj(out.transpose(1, 2).contiguous().view(B, T, C))
+
+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(F.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
+
+@dataclass
+class Sam3Config:
+    vocab_size: int
+    d_model: int = 468
+    n_layers: int = 14
+    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"
+
+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  # Weight tying
+
+    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)
+
+# -------------------------------
+# 3) Load Model from Hugging Face Hub
+# -------------------------------
+def load_sam3_model_from_hf(repo_id: str, filename: str = "sam3-epoch1-best.safetensors"):
+    print(f"📥 Loading config and weights from: {repo_id}")
+    config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+    weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+    with open(config_path, "r") as f:
+        config_dict = json.load(f)
+
+    # Ensure vocab_size matches tokenizer after adding special tokens
+    config_dict["vocab_size"] = len(tokenizer)
+    config = Sam3Config(**config_dict)
+
+    model = Sam3(config).to(device)
+    state_dict = load_file(weights_path)
+    model.load_state_dict(state_dict, strict=False)
+
+    model.eval()
+    print(f"✅ Model loaded successfully from Hugging Face Hub: {repo_id}")
+    return model
+
+# Load model
+model = load_sam3_model_from_hf("Smilyai-labs/Sam-3.5-1")
+
+# -------------------------------
+# 4) Sampling Function (Enhanced from your original)
+# -------------------------------
+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 Interface — WITH STYLED THINKING STEPS
+# -------------------------------
+def predict(message, history):
+    # Build prompt
+    chat_history = []
+    for human, assistant in history:
+        chat_history.append(f"{SPECIAL_TOKENS['user']} {human} {SPECIAL_TOKENS['eot']}")
+        if assistant:
+            # Assistant responses may contain <|think|>...<|eot|> blocks — we don't reconstruct them here
+            chat_history.append(f"{SPECIAL_TOKENS['assistant']} {assistant} {SPECIAL_TOKENS['eot']}")
+
+    chat_history.append(f"{SPECIAL_TOKENS['user']} {message} {SPECIAL_TOKENS['eot']}")
+    
+    system_prompt = "You are Sam-3.5, an advanced reasoning AI. You think step-by-step, analyze deeply, and respond with precision. You do not guess — you deduce. Avoid medical or legal advice."
+    prompt = f"{SPECIAL_TOKENS['system']} {system_prompt} {SPECIAL_TOKENS['eot']}\n" + "\n".join(chat_history) + f"\n{SPECIAL_TOKENS['assistant']} {SPECIAL_TOKENS['think']}"
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_LENGTH).to(device)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+
+    generated_text = ""
+    thinking_mode = False
+    thinking_buffer = ""
+
+    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=False)
+
+        # Append to sequence
+        input_ids = torch.cat([input_ids, next_token], dim=1)
+        attention_mask = torch.cat([attention_mask, torch.ones((1, 1), device=device, dtype=attention_mask.dtype)], dim=1)
+
+        # Handle thinking mode
+        if not thinking_mode and token_str.strip() == "<|think|>":
+            thinking_mode = True
+            thinking_buffer = ""
+            continue
+
+        if thinking_mode:
+            if token_str.strip() == "<|eot|>":
+                # End thinking block → yield styled output
+                thinking_buffer = thinking_buffer.strip()
+                if thinking_buffer:
+                    yield f"<div style='background-color:#f8f9fa; padding:12px; border-left:4px solid #007bff; border-radius:0 8px 8px 0; margin:10px 0; font-style:italic; color:#495057; font-size:0.95em;'>💡 <strong>Thinking:</strong> {thinking_buffer}</div>"
+                thinking_mode = False
+                continue
+            else:
+                thinking_buffer += token_str
+                continue
+
+        # Normal output
+        if not thinking_mode:
+            # Clean token for display (optional: handle GPT-2 space artifacts)
+            clean_token = token_str.replace('Ġ', ' ').replace('Ċ', '\n')
+            generated_text += clean_token
+            yield generated_text
+
+        # Stop if final EOT (outside thinking block)
+        if token_id == EOT_ID and not thinking_mode:
+            break
+
+# -------------------------------
+# 6) Launch Gradio Interface
+# -------------------------------
+CSS = """
+.gradio-container .message-bubble {
+    border-radius: 12px !important;
+    padding: 10px 14px !important;
+    font-size: 16px !important;
+}
+.gradio-container .message-bubble.user {
+    background-color: #007bff !important;
+    color: white !important;
+}
+.gradio-container .message-bubble.assistant {
+    background-color: #f8f9fa !important;
+    color: #212529 !important;
+    border: 1px solid #e9ecef;
+}
+"""
+
+demo = gr.ChatInterface(
+    fn=predict,
+    title="🧠 Sam-3.5: The Reasoning AI",
+    description="""
+    Sam-3.5 doesn’t just answer — it **thinks first**.  
+    Watch its internal reasoning unfold in real time — step by step, clearly shown.  
+    No guessing. No fluff. Just pure deduction.
+    
+    Try asking:
+    → “Why does a mirror reverse left and right but not up and down?”  
+    → “If I have 3 apples and give away half, then buy 5 more, how many do I have?”  
+    → “Explain quantum entanglement like I’m 10.”  
+    → “What’s wrong with 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.5 🤔",
+        bubble_full_width=False,
+        height=600,
+        avatar_images=(
+            "https://huggingface.co/datasets/huggingface/branding/resolve/main/avatar-bot.jpg",
+            "https://huggingface.co/datasets/huggingface/branding/resolve/main/avatar-user.jpg"
+        )
+    ),
+    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?"
+    ],
+    css=CSS,
+    cache_examples=False
+)
+
+if __name__ == "__main__":
+    demo.launch(show_api=True)