Spaces:
Sleeping
Sleeping
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() | |