File size: 3,044 Bytes
75a9771 0bdcf6d 75a9771 0bdcf6d 75a9771 0bdcf6d 75a9771 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import gradio as gr
import torch
import sys
import os
from huggingface_hub import hf_hub_download
sys.path.insert(0, os.path.dirname(__file__))
from src.model import Nexus
from src.config import NexusConfig
from tokenizers import Tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Loading on {device}")
config = NexusConfig()
model = Nexus(config)
REPO = "JustScriptzz/nexus-smAll-v1"
weights_local = os.path.join(os.path.dirname(__file__), "weights", "nexus_instruct.pt")
if os.path.exists(weights_local):
weights_path = weights_local
else:
print("Downloading weights from HuggingFace...")
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 = model.to(device)
model.eval()
print("Model loaded")
tokenizer_local = os.path.join(os.path.dirname(__file__), "data", "tokenizer.json")
if os.path.exists(tokenizer_local):
tokenizer_path = tokenizer_local
else:
tokenizer_path = hf_hub_download(repo_id=REPO, filename="data/tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_path)
bos_id = tokenizer.token_to_id("<bos>") or 1
eos_id = tokenizer.token_to_id("<eos>") or 2
def chat(message, history):
history_ids = [bos_id]
for user_msg, bot_msg in history:
if user_msg:
history_ids.extend(tokenizer.encode(f"User: {user_msg}\nAssistant:").ids)
if bot_msg:
history_ids.extend(tokenizer.encode(f" {bot_msg}").ids + [eos_id])
history_ids.extend(tokenizer.encode(f"User: {message}\nAssistant:").ids)
input_tensor = torch.tensor([history_ids[-config.max_seq_len:]], dtype=torch.long, device=device)
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
response_ids = input_tensor[0].tolist()
response_ids = response_ids[-(input_tensor.shape[1] - len(history_ids)):]
response = tokenizer.decode(response_ids)
response = response.split("<eos>")[0].split("User:")[0].replace("Assistant:", "").strip()
if len(response) < 2:
response = "[no response]"
return response
css = """
footer {visibility: hidden}
.message {font-size: 14px}
"""
demo = gr.ChatInterface(
fn=chat,
title="Nexus SmAll v1",
description="A 89.8M parameter transformer trained from scratch. May produce incoherent outputs — it's a tiny model!",
css=css,
theme="soft",
)
if __name__ == "__main__":
demo.launch()
|