import gradio as gr
from transformers import pipeline

# Function to clean the output by truncating at the last full sentence
def clean_output(text):
    if '.' in text:
        return text[:text.rfind('.')+1]  # Truncate at the last full sentence
    return text  # Return the text as is if no period is found

# Function to generate the story
def generate_story(title, model_name):
    # Use text-generation pipeline from Hugging Face
    generator = pipeline('text-generation', model=model_name)
    
    # Generate the story based on the input title
    story = generator(title,
                      max_length=230,  # Set the maximum length for the generated text (story) to 230 tokens
                      no_repeat_ngram_size=3,  # Avoid repeating sequences of 3 words
                      temperature=0.8,  # Introduce some randomness for diversity
                      top_p=0.95  # Nucleus sampling for more coherent text
                      )[0]['generated_text']
    
    # Clean the generated story to ensure it ends with a full sentence
    cleaned_story = clean_output(story)
    
    # Return the cleaned story
    return cleaned_story

# Gradio interface setup
demo = gr.Interface(
    fn=generate_story,
    inputs=[
        gr.Textbox(label="Enter Story Title", placeholder="Type a title here..."),  # Title input
        gr.Dropdown(choices=['gpt2', 'gpt2-large', 'EleutherAI/gpt-neo-2.7B'], value='gpt2', label="Select Model")  # Model selection
    ],
    outputs="text",
    title="AI Story Generator",
    description="Generate a creative story using different AI models."
)

demo.launch(share=True)