import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load pretrained GPT-2 model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")

# Function to generate script
def generate_script(title):
    # Add a context prompt to the title
    prompt = "Title: " + title + "\n\nScript:"

    # Tokenize the prompt
    input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=1024, truncation=True)

    # Generate text based on the prompt
    output = model.generate(input_ids, max_length=200, num_return_sequences=1, temperature=0.7)

    # Decode generated text and return
    script = tokenizer.decode(output[0], skip_special_tokens=True)
    return script

# Create Gradio interface
title_input = gr.inputs.Textbox(lines=2, label="Enter Title")
script_output = gr.outputs.Textbox(label="Generated Script")

gr.Interface(generate_script, title_input, script_output, title="Script Generator", description="Generate a script based on the provided title.").launch()