Update app.py

app.py

@@ -5,47 +5,61 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 tokenizer = AutoTokenizer.from_pretrained("suriya7/bart-finetuned-text-summarization")
 model = AutoModelForSeq2SeqLM.from_pretrained("suriya7/bart-finetuned-text-summarization")
 
-def generate_summary(text, prompt):
-    # Combine text and prompt into one input
-    combined_input = f"{prompt}: {text}"
+def generate_user_stories(text, prompt):
+    # Combine prompt with the text to guide the summarization
+    combined_input = f"Prompt: {prompt}\n\nText: {text}"
     inputs = tokenizer([combined_input], max_length=1024, return_tensors='pt', truncation=True)
-    summary_ids = model.generate(inputs['input_ids'], max_new_tokens=100, do_sample=False)
+    summary_ids = model.generate(inputs['input_ids'], max_new_tokens=150, do_sample=False)
     summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     
-    return summary
+    # Post-process to format as user stories
+    user_stories = format_as_user_stories(summary)
+    
+    return user_stories
+
+def format_as_user_stories(summary):
+    # Placeholder for formatting logic to extract user stories
+    # Here you can add specific rules or patterns to convert summary into user stories
+    lines = summary.split('. ')
+    user_stories = []
+    for line in lines:
+        # Example of simple pattern matching (can be customized)
+        if 'as a' in line.lower() and 'i want' in line.lower():
+            user_stories.append(line)
+    return '. '.join(user_stories)
 
 # Initialize session state for input history if it doesn't exist
 if 'input_history' not in st.session_state:
     st.session_state['input_history'] = []
 
 # Streamlit interface
-st.title("Text Summarization App")
+st.title("User Story Generator")
 
 # User text inputs
-bulk_text = st.text_area("Enter the bulk text you want to summarize", height=200)
-prompt = st.text_input("Enter the prompt for the summary", "What are the key points?")
+bulk_text = st.text_area("Enter the bulk text (e.g., client calls, meeting transcripts)", height=300)
+prompt = st.text_input("Enter the prompt for the user stories", "Extract user stories from the following text.")
 
-if st.button("Generate Summary"):
+if st.button("Generate User Stories"):
     if bulk_text and prompt:
-        with st.spinner("Generating summary..."):
-            summary = generate_summary(bulk_text, prompt)
+        with st.spinner("Generating user stories..."):
+            user_stories = generate_user_stories(bulk_text, prompt)
             
-        # Save the input and summary to the session state history
-        st.session_state['input_history'].append({"text": bulk_text, "prompt": prompt, "summary": summary})
+        # Save the input and user stories to the session state history
+        st.session_state['input_history'].append({"text": bulk_text, "prompt": prompt, "user_stories": user_stories})
         
-        st.subheader("Generated Summary:")
-        st.write(summary)
+        st.subheader("Generated User Stories:")
+        st.write(user_stories)
     else:
         st.warning("Please enter both the bulk text and the prompt.")
 
-# Display the history of inputs and summaries
+# Display the history of inputs and user stories
 if st.session_state['input_history']:
     st.subheader("History")
     for i, entry in enumerate(st.session_state['input_history']):
         st.write(f"**Input {i+1} (Text):** {entry['text']}")
         st.write(f"**Prompt {i+1}:** {entry['prompt']}")
-        st.write(f"**Summary {i+1}:** {entry['summary']}")
+        st.write(f"**User Stories {i+1}:** {entry['user_stories']}")
         st.write("---")
 
 # Instructions for using the app
-st.write("Enter your bulk text and a prompt for summarization, then click 'Generate Summary' to get a summarized version based on your prompt.")
+st.write("Enter your bulk text and a prompt for user story extraction, then click 'Generate User Stories' to get user stories from the text.")