File size: 1,494 Bytes
e7ffa07 |
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 |
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from model import GPT, GPTConfig # Import your model definition
import tiktoken
# Load the model from Hugging Face Hub
repo_id = "fridayfringe/my-gpt" # Change this to your Hugging Face repo ID
model_path = hf_hub_download(repo_id=repo_id, filename="model.pth")
# Initialize the model
config = GPTConfig()
model = GPT(config)
model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
model.eval()
# Load tokenizer (GPT-2 tokenizer)
enc = tiktoken.get_encoding("gpt2")
# Define the chatbot function
def generate_text(prompt, max_length=100):
tokens = enc.encode(prompt)
tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0) # (1, T)
with torch.no_grad():
for _ in range(max_length):
logits, _ = model(tokens)
next_token = torch.argmax(logits[:, -1, :], dim=-1)
tokens = torch.cat((tokens, next_token.unsqueeze(0)), dim=1)
output = enc.decode(tokens[0].tolist())
return output
# Create Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(label="Enter your prompt", placeholder="Type something..."),
gr.Slider(10, 200, value=100, step=10, label="Max Length"),
],
outputs="text",
title="Custom GPT Chatbot",
description="A lightweight GPT model trained on Spotify lyrics, deployed using Hugging Face Spaces.",
)
# Launch the app
iface.launch()
|