import gradio as gr
import torch
import tiktoken
from supplementary import GPTModel, generate_text_simple

# Load model configuration
GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}

# Initialize model
model = GPTModel(GPT_CONFIG_124M)

# Load the trained weights
model.load_state_dict(torch.load("my_gpt_model.pth", map_location=torch.device('cpu')))
model.eval()

tokenizer = tiktoken.get_encoding("gpt2")

def generate(prompt, max_new_tokens):
    token_ids = tokenizer.encode(prompt)
    input_ids = torch.tensor(token_ids).unsqueeze(0)
    output_ids = generate_text_simple(
        model=model,
        idx=input_ids,
        max_new_tokens=max_new_tokens,
        context_size=GPT_CONFIG_124M["context_length"]
    )
    return tokenizer.decode(output_ids.squeeze(0).tolist())

iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Max New Tokens")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="SamGPT Text Generation",
    description="Enter a prompt to generate text with the custom language model."
)

iface.launch()