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("""

⚡ VERTEX-1B-SNN

1B Parameter Spiking Neural Network | Distilled from Llama 3, Qwen 3.6, CodeLlama

Built by Mohamed Amine • CEO of VERTEX AI • Morocco

""") 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("""

📊 Model Specs

""") # Controls gr.Markdown("

🎛️ Generation Settings

") 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("""

🔗 Links

Organization
GitHub
LinkedIn

""") # Footer gr.Markdown("""
VERTEX AI © 2026 • Built with PyTorch, Transformers & Gradio • Download Model
""") # 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)