Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,90 @@
+import gradio as gr
+import torch
+import sys
+import os
+
+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)
+
+weights_path = os.path.join(os.path.dirname(__file__), "weights", "nexus_instruct.pt")
+if not os.path.exists(weights_path):
+    weights_path = os.path.join(os.path.dirname(__file__), "weights", "nexus_final.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_path = os.path.join(os.path.dirname(__file__), "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()