Update app.py

app.py

@@ -5,86 +5,52 @@ 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_user_stories(text, prompt):
+def summarize_text(text):
     try:
-        # Combine prompt with the text to guide the summarization
-        combined_input = f"Prompt: {prompt}\n\nText: {text}"
-        
         # Tokenize input with truncation to fit model requirements
-        inputs = tokenizer([combined_input], max_length=1024, return_tensors='pt', truncation=True)
+        inputs = tokenizer([text], max_length=1024, return_tensors='pt', truncation=True)
         
         # Generate summary
-        summary_ids = model.generate(inputs['input_ids'], max_new_tokens=150, do_sample=False)
+        summary_ids = model.generate(inputs['input_ids'], max_length=150, num_beams=4, early_stopping=True)
         summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
         
-        # Post-process to format as user stories
-        user_stories = format_as_user_stories(summary)
-        return user_stories
+        return summary
     except Exception as e:
         st.error(f"An error occurred: {e}")
         return ""
 
-def format_as_user_stories(summary):
-    user_stories = []
-    lines = summary.split('. ')
-    
-    for line in lines:
-        line = line.strip()
-        if "as a" in line.lower() and "i want" in line.lower():
-            # Extract the parts of the user story
-            parts = line.split("so that")
-            if len(parts) == 2:
-                story = {
-                    "As a": parts[0].strip().capitalize(),
-                    "I want to": parts[1].strip().capitalize(),
-                    "So that": parts[1].strip().capitalize()
-                }
-                user_stories.append(story)
-    
-    # Format user stories
-    formatted_stories = ""
-    for story in user_stories:
-        formatted_stories += f"**User Story:**\n\n" \
-                             f"**As a:** {story['As a']}\n" \
-                             f"**I want to:** {story['I want to']}\n" \
-                             f"**So that:** {story['So that']}\n\n"
-    
-    return formatted_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("User Story Generator")
+st.title("Text Summarizer")
 
 # User text inputs
 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 User Stories"):
-    if bulk_text and prompt:
-        with st.spinner("Generating user stories..."):
-            user_stories = generate_user_stories(bulk_text, prompt)
+if st.button("Summarize Text"):
+    if bulk_text:
+        with st.spinner("Generating summary..."):
+            summary = summarize_text(bulk_text)
             
-            if user_stories:
-                # 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 User Stories:")
-                st.write(user_stories)
+            if summary:
+                # Save the input and summary to the session state history
+                st.session_state['input_history'].append({"text": bulk_text, "summary": summary})
+                st.subheader("Summary:")
+                st.write(summary)
             else:
-                st.warning("No user stories were generated. Please check the input and try again.")
+                st.warning("No summary was generated. Please check the input and try again.")
     else:
-        st.warning("Please enter both the bulk text and the prompt.")
+        st.warning("Please enter the bulk text.")
 
-# Display the history of inputs and user stories
+# Display the history of inputs and summaries
 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"**User Stories {i+1}:** {entry['user_stories']}")
+        st.write(f"**Summary {i+1}:** {entry['summary']}")
         st.write("---")
 
 # Instructions for using the app
-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.")
+st.write("Enter your bulk text and click 'Summarize Text' to get a summary of the text.")