MS-GPT / app.py
mahdee987's picture
Update app.py
e9bedb4 verified
Raw
History Blame Contribute Delete
1.63 kB
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
# Load pre-trained GPT-2 model and tokenizer
model_name = "gpt2"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# Function to generate text
def generate_text(prompt):
if not prompt.strip(): # Check if the prompt is empty or just whitespace
return "enter a valid prompt."
inputs = tokenizer.encode(prompt, return_tensors="pt")
# Check if inputs is empty
if inputs.shape[1] == 0:
return "Failed to process the promp, try again."
outputs = model.generate(
inputs,
max_length=200,
num_return_sequences=1,
no_repeat_ngram_size=2,
temperature=0.7,
top_k=50,
top_p=0.95,
do_sample=True
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_text
# Create Gradio interface with a Submit button
iface = gr.Interface(
fn=generate_text, # Function to call
inputs=gr.Textbox(label="Enter a prompt:", placeholder="Type something...", lines=3),
outputs=gr.Textbox(label="Generated text:"), # Output text box
live=False, # Disable live updates
title="GPT-2 Text Generator", # Optional title
description="This is a story maker. You feel betrayed?? Just drop few lines of your story and it will melt your heart by creativity. AI is coming for you (diabolical laugh)", # Optional description
allow_flagging="never" # Optionally disable flagging
)
# Launch the Gradio app
iface.launch(share=True)