app.py

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()