app.py

import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
import os

# ═══════════════════════════════════════════════════════════════════════════════
# VERTEX-1B-SNN Demo Space
# Organization: vertex-snn
# ═══════════════════════════════════════════════════════════════════════════════

# ═══════════════════════════════════════════════════════════════════════════════
# Minimal SNN Layer (for inference only)
# ═══════════════════════════════════════════════════════════════════════════════

class SpikingLayer(torch.nn.Module):
    def __init__(self, dim, heads, num_steps=12):
        super().__init__()
        self.num_steps = num_steps
        self.norm1 = torch.nn.RMSNorm(dim)
        self.attn = torch.nn.MultiheadAttention(dim, heads, batch_first=True)
        self.norm2 = torch.nn.RMSNorm(dim)
        self.ffn = torch.nn.Sequential(
            torch.nn.Linear(dim, dim * 4, bias=False),
            torch.nn.GELU(),
            torch.nn.Linear(dim * 4, dim, bias=False),
        )
        self.beta = torch.nn.Parameter(torch.tensor(0.85))
        self.threshold = torch.nn.Parameter(torch.tensor(1.0))

    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x), need_weights=False)[0]
        x_norm = self.norm2(x)
        mem = torch.zeros_like(x_norm)
        spike_acc = torch.zeros_like(x_norm)
        beta = torch.sigmoid(self.beta)
        thr = torch.abs(self.threshold)
        for _ in range(self.num_steps):
            cur = self.ffn(x_norm)
            mem = beta * mem + cur
            spikes = (mem >= thr).float()
            mem = mem * (1 - spikes)
            spike_acc += spikes
        return x + spike_acc / self.num_steps


class VERTEXStudent(torch.nn.Module):
    def __init__(self, vocab_size=32000, dim=512, layers=8, heads=8):
        super().__init__()
        self.token_embed = torch.nn.Embedding(vocab_size, dim)
        self.pos_embed = torch.nn.Embedding(4096, dim)
        self.layers = torch.nn.ModuleList([SpikingLayer(dim, heads) for _ in range(layers)])
        self.norm = torch.nn.RMSNorm(dim)
        self.head = torch.nn.Linear(dim, vocab_size, bias=False)
        self.head.weight = self.token_embed.weight

    def forward(self, x):
        B, S = x.shape
        pos = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1)
        x = self.token_embed(x) + self.pos_embed(pos)
        for layer in self.layers:
            x = layer(x)
        return self.head(self.norm(x))

    @torch.no_grad()
    def generate(self, input_ids, max_new=50, temp=0.7, top_k=50, top_p=0.9, rep_penalty=1.1):
        self.eval()
        generated = input_ids.clone()
        for _ in range(max_new):
            logits = self(generated)[:, -1, :] / temp
            for tid in set(generated[0].tolist()):
                logits[0, tid] /= rep_penalty
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
            cumsum = torch.cumsum(sorted_probs, dim=-1)
            remove = cumsum > top_p
            remove[..., 1:] = remove[..., :-1].clone()
            remove[..., 0] = False
            logits[remove.scatter(-1, sorted_indices, remove)] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            next_tok = torch.multinomial(probs, 1)
            generated = torch.cat([generated, next_tok], dim=1)
            if next_tok.item() in [1, 2, 128001, 128009]:
                break
        return generated


# ═══════════════════════════════════════════════════════════════════════════════
# Model Loader - Uses PUBLIC tokenizer (no gated repo)
# ═══════════════════════════════════════════════════════════════════════════════

class VERTEXChat:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.tokenizer = None
        self.model = None
        self.load_model()

    def load_model(self):
        """Load public tokenizer and initialize model."""
        print("Loading VERTEX-1B-SNN...")

        # Use PUBLIC tokenizer - no gated repo needed!
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(
                "gpt2",
                trust_remote_code=True
            )
            print("✓ Loaded GPT-2 public tokenizer")
        except Exception as e:
            print(f"⚠ Tokenizer load failed: {e}")
            from transformers import GPT2Tokenizer
            self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        vocab_size = len(self.tokenizer)
        print(f"  Vocab size: {vocab_size}")

        self.model = VERTEXStudent(
            vocab_size=vocab_size,
            dim=512,
            layers=8,
            heads=8
        ).to(self.device)

        self.model.eval()
        print(f"✓ VERTEX-1B-SNN ready on {self.device}")
        print("⚠ Using randomly initialized weights — train model for real responses")

    def chat(self, message, history, temp, top_k, max_tokens):
        """Generate response from user message."""
        if not message.strip():
            return "", history

        prompt = self._format_prompt(message, history)
        inputs = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)

        with torch.no_grad():
            output = self.model.generate(
                inputs,
                max_new=max_tokens,
                temp=temp,
                top_k=top_k,
                top_p=0.9,
                rep_penalty=1.1
            )

        response = self.tokenizer.decode(output[0], skip_special_tokens=True)

        if prompt in response:
            response = response[len(prompt):].strip()

        history = history + [[message, response]]
        return "", history

    def _format_prompt(self, message, history):
        """Format conversation history into prompt."""
        parts = []
        parts.append("System: You are VERTEX, a 1B parameter Spiking Neural Network.\n")

        for human, assistant in history:
            parts.append(f"User: {human}\nAssistant: {assistant}\n")

        parts.append(f"User: {message}\nAssistant: ")

        return "".join(parts)


# ═══════════════════════════════════════════════════════════════════════════════
# Gradio Interface - Compatible with Gradio 6.x
# ═══════════════════════════════════════════════════════════════════════════════

# CSS defined GLOBALLY (outside function)
VERTEX_CSS = """
.gradio-container {
    background: linear-gradient(135deg, #0a0a0a 0%, #1a1a2e 100%) !important;
}
.chatbot {
    background: rgba(20, 20, 40, 0.8) !important;
    border: 1px solid #00d4ff !important;
    border-radius: 12px !important;
    height: 500px;
}
.user-message {
    background: linear-gradient(90deg, #00d4ff22, transparent) !important;
    border-left: 3px solid #00d4ff !important;
    color: #e0e0e0 !important;
}
.bot-message {
    background: linear-gradient(90deg, #ff006622, transparent) !important;
    border-left: 3px solid #ff0066 !important;
    color: #e0e0e0 !important;
}
.input-textbox {
    background: #1a1a2e !important;
    border: 1px solid #00d4ff44 !important;
    color: #e0e0e0 !important;
}
.generate-btn {
    background: linear-gradient(90deg, #00d4ff, #0066ff) !important;
    color: white !important;
    font-weight: bold !important;
    border: none !important;
}
.generate-btn:hover {
    box-shadow: 0 0 20px #00d4ff66 !important;
}
h1 {
    color: #00d4ff !important;
    text-shadow: 0 0 20px #00d4ff44;
}
.info-box {
    background: rgba(0, 212, 255, 0.1) !important;
    border: 1px solid #00d4ff44 !important;
    border-radius: 8px !important;
    padding: 12px !important;
    color: #ccc !important;
}
.info-box h3, .info-box h4 {
    color: #00d4ff !important;
    margin-top: 0;
}
.info-box a {
    color: #00d4ff !important;
    text-decoration: none;
}
.info-box a:hover {
    text-decoration: underline;
}
.footer {
    text-align: center;
    padding: 20px;
    color: #666;
    font-size: 0.85em;
}
.footer a {
    color: #00d4ff;
    text-decoration: none;
}
"""

def create_demo():
    vertex = VERTEXChat()

    with gr.Blocks(title="VERTEX-1B-SNN") as demo:
        # Header
        gr.Markdown("""
        <div style="text-align: center; padding: 20px 0;">
            <h1 style="font-size: 2.5em; margin-bottom: 0;">⚡ VERTEX-1B-SNN</h1>
            <p style="color: #888; font-size: 1.1em; margin-top: 8px;">
                1B Parameter Spiking Neural Network | Distilled from Llama 3, Qwen 3.6, CodeLlama
            </p>
            <p style="color: #00d4ff; font-size: 0.9em;">
                Built by <strong>Mohamed Amine</strong> • CEO of VERTEX AI • Morocco
            </p>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=3):
                # Chat interface - Gradio 6.x compatible
                chatbot = gr.Chatbot(
                    height=500,
                    elem_classes=["chatbot"]
                )

                with gr.Row():
                    msg = gr.Textbox(
                        placeholder="Ask VERTEX anything... (code, reasoning, general knowledge)",
                        show_label=False,
                        container=False,
                        elem_classes=["input-textbox"],
                        scale=8
                    )
                    submit = gr.Button("⚡ Generate", elem_classes=["generate-btn"], scale=2)

                clear = gr.Button("🗑️ Clear Conversation", variant="secondary")

            with gr.Column(scale=1):
                # Info panel
                gr.Markdown("""
                <div class="info-box">
                    <h3>📊 Model Specs</h3>
                    <ul style="padding-left: 16px;">
                        <li><strong>Params:</strong> 1B</li>
                        <li><strong>Architecture:</strong> SNN Transformer</li>
                        <li><strong>Teachers:</strong> Llama 3 8B, Qwen 3.6 27B, CodeLlama 7B, Llama 2 7B</li>
                        <li><strong>Vocab:</strong> GPT-2 (public, 50k tokens)</li>
                        <li><strong>Spike Steps:</strong> 12</li>
                        <li><strong>Precision:</strong> BF16</li>
                    </ul>
                </div>
                """)

                # Controls
                gr.Markdown("<h3 style='color: #00d4ff;'>🎛️ Generation Settings</h3>")
                temp = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
                top_k = gr.Slider(1, 100, value=50, step=1, label="Top-K")
                max_tokens = gr.Slider(10, 512, value=128, step=10, label="Max Tokens")

                gr.Markdown("""
                <div class="info-box" style="margin-top: 20px;">
                    <h4 style="color: #ff0066;">🔗 Links</h4>
                    <p style="font-size: 0.85em;">
                        <a href="https://huggingface.co/vertex-snn">Organization</a><br>
                        <a href="https://github.com/mohamedamine/vertex-snn">GitHub</a><br>
                        <a href="https://linkedin.com/in/mohamedamine">LinkedIn</a>
                    </p>
                </div>
                """)

        # Footer
        gr.Markdown("""
        <div class="footer">
            VERTEX AI © 2026 • Built with PyTorch, Transformers & Gradio • 
            <a href="https://huggingface.co/vertex-snn/vertex-1b-snn-v2.1">Download Model</a>
        </div>
        """)

        # Event handlers
        def respond(message, history, temp, top_k, max_tokens):
            if not message.strip():
                return "", history

            return vertex.chat(message, history, temp, top_k, max_tokens)

        submit.click(
            respond,
            inputs=[msg, chatbot, temp, top_k, max_tokens],
            outputs=[msg, chatbot]
        )
        msg.submit(
            respond,
            inputs=[msg, chatbot, temp, top_k, max_tokens],
            outputs=[msg, chatbot]
        )
        clear.click(lambda: None, None, chatbot, queue=False)

    return demo


if __name__ == "__main__":
    demo = create_demo()
    demo.launch(css=VERTEX_CSS)