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)