Spaces:
Sleeping
Sleeping
File size: 1,378 Bytes
4dde291 56f888d 4dde291 56f888d 4dde291 56f888d 4dde291 56f888d 4dde291 56f888d 4dde291 56f888d 4dde291 56f888d |
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 |
import gradio as gr
import torch
from transformers import PreTrainedTokenizerFast
import os
# -----------------------------
# Load tokenizer
# -----------------------------
tokenizer_path = "./Dave/tokenizer.json"
tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)
# -----------------------------
# Load custom model
# -----------------------------
from dave_model import SimpleTransformer # put your model class in dave_model.py
vocab_size = tokenizer.vocab_size
model_path = "./Dave/pytorch_model.bin"
model = SimpleTransformer(vocab_size=vocab_size)
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()
# -----------------------------
# Inference function
# -----------------------------
def generate_response(prompt):
input_ids = torch.tensor([tokenizer.encode(prompt).ids])
with torch.no_grad():
output = model(input_ids)
predicted_ids = output.argmax(-1)[0].tolist()
response = tokenizer.decode(predicted_ids)
return response
# -----------------------------
# Gradio interface
# -----------------------------
iface = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(lines=2, placeholder="Type your prompt here..."),
outputs="text",
title="Dave – Fully Custom AI",
description="Interact with your fully custom AI model."
)
if __name__ == "__main__":
iface.launch()
|