import gradio as gr
import torch
from torch.nn import functional as F
from gpt_class import GPTConfig, GPT
import tiktoken

# Setup device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model
state_dict = torch.load('model_51999.pt', map_location=device)
config = state_dict['config']
model = GPT(config)
model.load_state_dict(state_dict['model'])
model.to(device)
model.eval()

# Set seed for reproducibility
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)

# Get tokenizer
tokenizer = tiktoken.get_encoding("gpt2")

def generate_text(example, num_return_sequences='4', max_length='64'):
    num_return_sequences = int(num_return_sequences) if num_return_sequences.isdigit() else 4
    max_length = int(max_length) if max_length.isdigit() else 64

    model.eval()
    tokens = tokenizer.encode(example)
    tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).repeat(num_return_sequences, 1)
    tokens = tokens.to(device)
    sample_rng = torch.Generator(device=device)

    xgen = tokens
    while xgen.size(1) < max_length:
        with torch.no_grad():
            with torch.autocast(device_type=device):
                logits, _ = model(xgen)  # Assumes model returns logits and optional loss
            logits = logits[:, -1, :]  # Get last token logits
            probs = F.softmax(logits, dim=-1)
            topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
            ix = torch.multinomial(topk_probs, 1, generator=sample_rng)
            xcol = torch.gather(topk_indices, -1, ix)
            xgen = torch.cat((xgen, xcol), dim=1)

    results = []
    for i in range(num_return_sequences):
        tokens = xgen[i, :max_length].tolist()
        decoded = tokenizer.decode(tokens)
        results.append(decoded)
    return "\n\n".join(results)

# Create Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.components.Textbox(label="Prompt"),
        gr.components.Textbox(label="Number of Sequences [1-4]"),
        gr.components.Textbox(label="Maximum Length [32-128]")
    ],
    outputs=gr.components.Textbox(label="Generated Text"),
    title="Text Generator",
    description="Enter a prompt to generate text using a GPT model. Adjust the number of sequences and the maximum length as needed.",
    examples=[
        ["It is raining and my family", "2", "64"],
        ["We entered into the forest and", "2", "64"],
        ["I sat for doing my homework", "2", "64"]
    ]
)

iface.launch(share=True)