Update app.py

app.py

@@ -1,41 +1,75 @@
 import gradio as gr
 from transformers import pipeline
+import re
 
-# Function to clean the output by truncating at the last full sentence
+# Function to clean the output by truncating at the last full sentence and formatting paragraphs
 def clean_output(text):
+    # Remove unwanted symbols and replace them with appropriate punctuation or space
+    text = re.sub(r'\.{2,}', '.', text)  # Replace sequences of more than one period with a single period
+    text = re.sub(r'[:\-]+', '', text)  # Remove colons and dashes
+    text = re.sub(r'[()]+', '', text)  # Remove parentheses
+    text = re.sub(r'\s+', ' ', text)  # Replace excessive spaces
+    text = re.sub(r'[^\S\n]+', ' ', text)  # Remove non-visible spaces like tabs
+
+    # Ensure the text ends with a full sentence
     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
+        text = text[:text.rfind('.')+1]  # Truncate at the last full sentence
+    
+    # Add paragraph breaks by splitting sentences into paragraphs
+    sentences = re.split(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s', text)  # Split by sentence-ending punctuation
+    
+    # Create paragraphs by grouping sentences
+    paragraph_size = len(sentences) // 4  # Split into approximately 4 paragraphs
+    paragraphs = [' '.join(sentences[i:i + paragraph_size]) for i in range(0, len(sentences), paragraph_size)]
+    
+    # Limit the number of paragraphs to 3-4
+    paragraphs = paragraphs[:4]
+    
+    # Join paragraphs with double line breaks
+    formatted_text = '\n\n'.join(paragraphs)
+    
+    return formatted_text.strip()  # Return trimmed and formatted text
 
 # Function to generate the story
-def generate_story(title, model_name):
+def generate_story(title, model_name="gpt2"):
     # 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=300,  # 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
+                      max_length=500,  # Set the maximum length for the generated text (story)
+                      no_repeat_ngram_size=3,  # Avoid repeating any sequence of 3 words (to prevent repetitive text)
+                      temperature=0.8,  # Introduce some randomness
+                      top_p=0.95  # Use nucleus sampling for coherent output
                       )[0]['generated_text']
     
-    # Clean the generated story to ensure it ends with a full sentence
+    # Clean the generated story to ensure it ends with a full sentence and trim it into paragraphs
     cleaned_story = clean_output(story)
     
-    # Return the cleaned story
+    # Return the cleaned and formatted story
     return cleaned_story
-
-# Gradio interface setup
+# Create the Gradio interface using gr.Interface
 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','maldv/badger-writer-llama-3-8b', 'GOAT-AI/GOAT-70B-Storytelling'], value='gpt2', label="Select Model")  # Model selection
+        gr.Dropdown(choices=[
+            'gpt2',
+            'gpt2-large',
+            'EleutherAI/gpt-neo-2.7B',
+            'EleutherAI/gpt-j-6B',
+            'maldv/badger-writer-llama-3-8b',
+            'gpt-neo-2.7B'
+            'TheBloke/Iambe-Storyteller-20B-GPTQ
+        ], value='gpt2', label="Choose Model")  # Model selection
     ],
     outputs="text",
     title="AI Story Generator",
-    description="Generate a creative story using different AI models."
+    description="Generate a creative story using different AI models.",
+    examples=[
+        ["Sara burst into her friend's house, only to find it plunged into darkness. A strange, pulsing glow flickered from the corner, casting eerie shadows on the walls. Her heart raced as she called out, but there was no answer. Something wasn’t right. On the table sat an unfamiliar, glowing device—humming with energy. With a deep breath, Sara stepped closer, knowing that once she touched it, there would be no turning back."]
+    ]
 )
 
-demo.launch(share=True)
+# Launch the interface with sharing enabled
+demo.launch(share=True)