Update app.py

app.py

@@ -9,6 +9,22 @@ from torch.nn import functional as F
 import tiktoken
 import gradio as gr
 import asyncio 
+
+# Add the post-processing function here
+def post_process_text(text):
+    # Ensure the text starts with a capital letter
+    text = text.capitalize()
+    
+    # Remove any incomplete sentences at the end
+    sentences = text.split('.')
+    complete_sentences = sentences[:-1] if len(sentences) > 1 else sentences
+    
+    # Rejoin sentences and add a period if missing
+    processed_text = '. '.join(complete_sentences)
+    if not processed_text.endswith('.'):
+        processed_text += '.'
+    
+    return processed_text
 # Define the model architecture
 class GPTConfig:
     def __init__(self):
@@ -134,7 +150,7 @@ import gradio as gr
 # [Your existing model code remains unchanged]
 
 # Modify the generate_text function to be asynchronous
-async def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
+async def generate_text(prompt, max_length=432, temperature=0.8, top_k=40):
     input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
     generated = []
     
@@ -151,13 +167,16 @@ async def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
             input_ids = torch.cat([input_ids, next_token], dim=-1)
             generated.append(next_token.item())
             
-            yield enc.decode([next_token.item()])
+            next_token_str = enc.decode([next_token.item()])
+            yield next_token_str
             
-            if next_token.item() == enc.encode('\n')[0] and len(generated) > 20:
+            if next_token.item() == enc.encode('\n')[0] and len(generated) > 100:
                 break
             
-            await asyncio.sleep(0.05)  # Use asyncio.sleep instead of time.sleep
+            await asyncio.sleep(0.02)  # Slightly faster typing effect
 
+    if len(generated) == max_length:
+        yield "... (output truncated due to length)"
 # Modify the gradio_generate function to be asynchronous
 async def gradio_generate(prompt, max_length, temperature, top_k):
     output = ""
@@ -178,20 +197,23 @@ css = """
 </style>
 """
 
-# 6. Gradio App Definition
 with gr.Blocks(css=css) as demo:
-    gr.HTML("<div class='header'><h1>🌟 GPT-2 Text Generator</h1></div>")
+    gr.HTML("<div class='header'><h1>🌟 GPT-2 Storyteller</h1></div>")
     
     with gr.Row():
         with gr.Column(scale=3):
-            prompt = gr.Textbox(placeholder="Enter your prompt here...", label="Prompt", elem_classes="user-input")
+            prompt = gr.Textbox(
+                placeholder="Start your story here (e.g., 'Once upon a time in a magical forest...')",
+                label="Story Prompt",
+                elem_classes="user-input"
+            )
         with gr.Column(scale=1):
-            generate_btn = gr.Button("Generate", elem_classes="generate-btn")
+            generate_btn = gr.Button("Generate Story", elem_classes="generate-btn")
     
     with gr.Row():
-        max_length = gr.Slider(minimum=20, maximum=500, value=100, step=1, label="Max Length")
-        temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
-        top_k = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k")
+        max_length = gr.Slider(minimum=50, maximum=500, value=432, step=1, label="Max Length")
+        temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Temperature")
+        top_k = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-k")
     
     output = gr.Markdown(elem_classes="output-box")
     
@@ -201,6 +223,6 @@ with gr.Blocks(css=css) as demo:
         outputs=output
     )
 
-# 7. Launch the app
+
 if __name__ == "__main__":
     demo.launch()