import gradio as gr
import torch
from gpt_dev import GPTLanguageModel, encode, decode, generate_text  # Assuming these are in gpt_dev.py

# Initialize model parameters and load the pre-trained model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Parameters (adjust based on your model's architecture)
block_size = 256
n_embd = 384
n_head = 6
n_layer = 6
vocab_size = 95

# Initialize the model
model = GPTLanguageModel()
model.to(device)

# Load the saved model weights
checkpoint = torch.load("gpt_language_model.pth", map_location=device)
model.load_state_dict(checkpoint)
model.eval()  # Set the model to evaluation mode

# Define the text generation function
def generate_response(prompt, max_length=100, temperature=1.0):
    generated_text = generate_text(model, prompt, max_length=max_length, temperature=temperature)
    return generated_text

# Gradio interface
def gradio_interface(prompt, max_length=100, temperature=1.0):
    return generate_response(prompt, max_length, temperature)

# Set up Gradio UI using gr.components (for Gradio 3.x or later)
interface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Prompt", value="Once upon a time"),
        gr.Slider(50, 240, step=1, value=75, label="Max Length"),
    ],
    outputs="text",
    title="Odeyssey Rhyme Generator",
    description="Enter a prompt to generate text."
)

# Launch the Gradio interface
if __name__ == "__main__":
    interface.launch(share=True)  # Set share=False if you don't want to share publicly