IDK2 / app.py
Arpit1234's picture
Create app.py
9d4746b verified
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()