mygen / app.py
ZenithBrew's picture
Update app.py
8a736aa verified
import gradio as gr
from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
# Pre-download the model and tokenizer
model_name = "gpt2"
GPT2Tokenizer.from_pretrained(model_name)
GPT2LMHeadModel.from_pretrained(model_name)
# Load the pipeline with additional parameters
generator = pipeline(
'text-generation',
model=model_name,
do_sample=True, # Enables sampling for more creative output
temperature=0.7, # Controls randomness (0.7 is balanced for creativity and coherence)
top_k=50, # Limits the vocabulary to the top 50 most likely tokens
max_length=150 # Increased length for longer responses
)
# Function to generate text
def generate_text(prompt):
output = generator(prompt, num_return_sequences=1)
return output[0]['generated_text']
# Build the interface
iface = gr.Interface(
fn=generate_text,
inputs=gr.Textbox(lines=2, placeholder="Enter a prompt here..."),
outputs="text",
title="AI Text Generator",
description="Generate text with AI!"
)
# Launch it
iface.launch()