Update app.py

app.py

@@ -11,21 +11,25 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+# Store conversation history
+conversation_history = []
+
 def generate_full_story(excerpt: str) -> str:
     """
-    Given an incomplete story excerpt (without header details), this function calls the model 
-    to generate the complete story that includes Parv, Key Event, Section and the story continuation.
+    Given an incomplete story excerpt, generate the complete story including Parv, Key Event, Section, and continuation.
     """
+    if not excerpt.strip():
+        return "Please enter a valid story excerpt."
+    
     # Tokenize the user-provided excerpt.
     encoded_input = tokenizer(excerpt, return_tensors = "pt")
-    # Move tensors to the appropriate device.
     encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
     
-    # Generate tokens. Here, we set parameters to control length and creativity.
+    # Generate text with controlled parameters.
     output = model.generate(
         encoded_input["input_ids"],
         attention_mask = encoded_input["attention_mask"],
-        max_new_tokens = 200,         # Generate 200 new tokens on top of the input.
+        max_new_tokens = 200,
         do_sample = True,
         temperature = 0.8,
         top_p = 0.95,
@@ -33,26 +37,46 @@ def generate_full_story(excerpt: str) -> str:
         return_dict_in_generate = True
     )
     
-    # Decode the generated sequence.
+    # Decode generated text.
     generated_text = tokenizer.decode(output.sequences[0], skip_special_tokens = True)
     
+    # Append to conversation history
+    conversation_history.append((excerpt, generated_text))
+    
     return generated_text
 
+def get_conversation_history():
+    """Displays conversation history."""
+    if not conversation_history:
+        return "No conversations started."
+    return "\n\n".join([f"**User:** {inp}\n**AI:** {out}" for inp, out in conversation_history])
+
+def clear_conversation():
+    """Clears the conversation history."""
+    conversation_history.clear()
+    return "No conversations started."
+
 # Build the Gradio interface.
-interface = gr.Interface(
-    fn = generate_full_story,
-    inputs = gr.Textbox(
-        lines = 5,
-        label = "Incomplete story excerpt",
-        placeholder = "Enter an excerpt from the Mahabharata here..."
-    ),
-    outputs = gr.Textbox(label = "Chapter summary"),
-    title = "🏺 Mythology Storyteller",
-    description = (
-        "Enter a phrase from a chapter of your choice (if possible please enter the Parv, Key Event, & Section for an accurate answer). "
-        "The model will generate the summary of the respective chapter."
-    )
-)
+with gr.Blocks() as interface:
+    gr.Markdown("# 🏺 Mythology Storyteller")
+    gr.Markdown("Enter a phrase from a chapter of your choice. The model will generate the summary of the respective chapter.")
+    
+    with gr.Row():
+        user_input = gr.Textbox(lines = 5, label = "Incomplete story excerpt", placeholder = "Enter an excerpt from the Mahabharata here...")
+    
+    output_text = gr.Textbox(label = "Chapter summary")
+    
+    generate_btn = gr.Button("Generate Story")
+    generate_btn.click(fn = generate_full_story, inputs = user_input, outputs = output_text)
+    
+    with gr.Row():
+        history_display = gr.Textbox(label = "Conversation History", interactive = False)
+    
+    show_history_btn = gr.Button("Show Conversation History")
+    show_history_btn.click(fn = get_conversation_history, outputs = history_display)
+    
+    clear_btn = gr.Button("Clear Conversation")
+    clear_btn.click(fn = clear_conversation, outputs = history_display)
 
 # Launch the Gradio app.
-interface.launch()
+interface.launch()