Update app.py

app.py

@@ -6,26 +6,35 @@ tokenizer = AutoTokenizer.from_pretrained("suriya7/bart-finetuned-text-summariza
 model = AutoModelForSeq2SeqLM.from_pretrained("suriya7/bart-finetuned-text-summarization")
 
 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=150, do_sample=False)
-    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
+    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)
+        
+        # Generate summary
+        summary_ids = model.generate(inputs['input_ids'], max_new_tokens=150, do_sample=False)
+        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
+    except Exception as e:
+        st.error(f"An error occurred: {e}")
+        return ""
 
 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('. ')
+    # Enhanced formatting logic to extract user stories
     user_stories = []
+    lines = summary.split('. ')
+    
     for line in lines:
-        # Example of simple pattern matching (can be customized)
-        if 'as a' in line.lower() and 'i want' in line.lower():
+        # Improved pattern matching for user stories
+        line = line.strip()
+        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
@@ -44,11 +53,13 @@ if st.button("Generate User Stories"):
         with st.spinner("Generating user stories..."):
             user_stories = generate_user_stories(bulk_text, prompt)
             
-        # 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 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)
+            else:
+                st.warning("No user stories were generated. Please check the input and try again.")
     else:
         st.warning("Please enter both the bulk text and the prompt.")