Spaces:
Sleeping
Sleeping
File size: 1,975 Bytes
0eff763 4dde291 c46d75e 17df16c 56f888d 0eff763 56f888d 0eff763 c46d75e 17df16c c46d75e bf6da7a c46d75e 0eff763 c46d75e 0eff763 c46d75e 0eff763 c46d75e 0eff763 4dde291 17df16c 56f888d 17df16c 0eff763 17df16c 56f888d c46d75e 0eff763 17df16c 0eff763 17df16c 0eff763 17df16c |
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 |
# app.py
import torch
import torch.nn as nn
from tokenizers import Tokenizer
import gradio as gr
# -----------------------------
# Load tokenizer
# -----------------------------
tokenizer_path = "tokenizer.json"
tokenizer = Tokenizer.from_file(tokenizer_path)
vocab_size = tokenizer.get_vocab_size()
# -----------------------------
# Define the same transformer as used in training
# -----------------------------
class SimpleTransformer(nn.Module):
def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=4):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.fc = nn.Linear(d_model, vocab_size)
def forward(self, x):
x = self.embedding(x)
x = x.transpose(0, 1)
x = self.transformer(x)
x = x.transpose(0, 1)
return self.fc(x)
# -----------------------------
# Load model weights
# -----------------------------
model = SimpleTransformer(vocab_size)
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
model.eval()
# -----------------------------
# Generation function
# -----------------------------
def generate(prompt):
input_ids = tokenizer.encode(prompt).ids
input_tensor = torch.tensor([input_ids])
with torch.no_grad():
output = model(input_tensor)
predicted_ids = torch.argmax(output, dim=-1)[0].tolist()
response_text = tokenizer.decode(predicted_ids)
return response_text
# -----------------------------
# Gradio Interface
# -----------------------------
iface = gr.Interface(
fn=generate,
inputs=gr.Textbox(lines=2, placeholder="Type a prompt for Dave..."),
outputs="text",
title="Dave – Fully Custom AI",
description="Interact with your fully custom AI trained from scratch."
)
iface.launch()
|