| import gradio as gr |
| import torch |
| from huggingface_hub import hf_hub_download |
| from tokenizers import Tokenizer |
|
|
| REPO = "JustScriptzz/nexus-smAll-v1" |
|
|
| def load_model(): |
| import sys |
| sys.path.insert(0, REPO) |
| from src.model import Nexus |
| from src.config import NexusConfig |
|
|
| device = torch.device("cpu") |
| config = NexusConfig() |
| model = Nexus(config) |
| weights_path = hf_hub_download(repo_id=REPO, filename="weights/nexus_instruct.pt") |
| checkpoint = torch.load(weights_path, map_location=device, weights_only=False) |
| model.load_state_dict(checkpoint["model_state_dict"]) |
| model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8) |
| model.eval() |
| tokenizer_path = hf_hub_download(repo_id=REPO, filename="data/tokenizer.json") |
| tokenizer = Tokenizer.from_file(tokenizer_path) |
| return model, tokenizer, config |
|
|
| model, tokenizer, config = load_model() |
| bos_id = tokenizer.token_to_id("<bos>") or 1 |
| eos_id = tokenizer.token_to_id("<eos>") or 2 |
|
|
| def chat(message, history): |
| messages = history + [{"role": "user", "content": message}] |
| tokens = [bos_id] |
| for msg in messages: |
| if msg["role"] == "user": |
| tokens.extend(tokenizer.encode(f"User: {msg['content']}\nAssistant:").ids) |
| elif msg["role"] == "assistant": |
| tokens.extend(tokenizer.encode(f" {msg['content']}").ids + [eos_id]) |
| input_tensor = torch.tensor([tokens[-config.max_seq_len:]], dtype=torch.long) |
| with torch.no_grad(): |
| for _ in range(128): |
| seq_len = input_tensor.shape[1] |
| if seq_len > config.max_seq_len: |
| input_tensor = input_tensor[:, -config.max_seq_len:] |
| logits = model(input_tensor, 0) |
| logits = logits[:, -1, :] / 0.2 |
| probs = torch.softmax(logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| input_tensor = torch.cat([input_tensor, next_token], dim=-1) |
| if next_token.item() == eos_id: |
| break |
| new_ids = input_tensor[0].tolist()[len(tokens):] |
| reply = tokenizer.decode(new_ids) |
| for tok in ["<|assistant|>", "<|user|>", "<|system|>"]: |
| reply = reply.replace(tok, "") |
| reply = reply.split("<eos>")[0].split("User:")[0].replace("Assistant:", "").strip() |
| return reply or "..." |
|
|
| demo = gr.ChatInterface( |
| fn=chat, |
| title="Nexus SmAll v1", |
| description="89.8M parameter transformer built from scratch", |
| theme=gr.themes.Base( |
| primary_hue="purple", |
| neutral_hue="slate", |
| font=gr.themes.GoogleFont("Inter"), |
| ).set( |
| body_background_fill="#0a0a0f", |
| body_text_color="#e1e1e6", |
| block_background_fill="#111118", |
| block_border_color="#1e1e2e", |
| block_label_text_color="#888", |
| input_background_fill="#111118", |
| input_background_fill_focus="#111118", |
| button_primary_background_fill="#7c3aed", |
| button_primary_background_fill_hover="#6d28d9", |
| ), |
| ) |
|
|
| demo.launch() |
|
|
