import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load the MARIOGPT model and tokenizer
model_name = "shyamsn97/Mario-GPT2-700-context-length"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# Define the function to generate Mario level
def generate_mario_level(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=200)
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_mario_level,
    inputs=gr.Textbox(lines=2, placeholder="Enter level description here..."),
    outputs="text",
    title="MARIOGPT Level Generator",
    description="Generate Mario levels using MARIOGPT by entering a level description.",
    examples=[["simple level"], ["difficult level with many enemies"], ["water level with lots of coins"]]
)

# Launch the Gradio interface
interface.launch()