Update app.py

app.py

@@ -27,7 +27,7 @@
 #    - **Hugging Face Model**: "sentence-transformers/all-MiniLM-L6-v2"
 
 # Choose your model from the above options:
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1" # Change this to one of the other models based on your needs.
+MODEL_NAME = "google/flan-t5-base" # Change this to one of the other models based on your needs.
 
 # ---------------------------------------------------------------------------
 # Code Below to Load, Generate, and Save Functional Requirement Documents
@@ -38,120 +38,96 @@ from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
 import os
 
-# Load Hugging Face Token from Environment (Set this in Hugging Face Spaces Secrets)
+# Load Hugging Face Token (Ensure it's set in Env Variables)
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# 🏆 Recommended Model: Mistral-7B (Lightweight & Open-Source)
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
+# ✅ Optimized Model (Flan-T5 for Lower Memory Usage)
+MODEL_NAME = "google/flan-t5-base"
 
-# 📌 Load Model Efficiently (Caching to avoid reloading)
+# 📌 Load Model Efficiently (Avoid Reloading)
 @st.cache_resource
 def load_model():
     try:
-        return pipeline("text-generation", model=MODEL_NAME, token=HF_TOKEN)
+        return pipeline("text2text-generation", model=MODEL_NAME, token=HF_TOKEN)
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
         return None
 
-generator = load_model()  # Load once and reuse
+# Load once and reuse
+generator = load_model()
 
 # 📌 Function to Generate Functional Requirement Document
 def generate_functional_requirements(topic):
     if generator is None:
         return "Error: Model failed to load."
 
-    # Structured Prompt for Better Document Output
     prompt = f"""
-    Generate a detailed functional requirements document for {topic} in the banking sector. 
-    The document should follow this structured format:
+    Generate a **detailed functional requirements document** for **{topic}** in the banking sector.
+    Ensure it includes the following sections:
     
-    1️⃣ **Introduction**  
-    - Overview  
-    - Purpose  
-    - Intended Users  
-
-    2️⃣ **Scope**  
-    - System Description  
-    - Key Functionalities  
-
-    3️⃣ **Functional Specifications**  
-    - Core Features  
-    - User Roles & Permissions  
-    - Transaction Processing  
-
-    4️⃣ **System Features**  
-    - Security & Compliance  
-    - Performance Metrics  
-
-    5️⃣ **Regulatory & Compliance**  
-    - Central Bank Regulations  
-    - Data Protection & Privacy  
-
-    6️⃣ **Conclusion**  
-    - Summary  
-    - Future Enhancements  
+    1️⃣ **Introduction**: Overview, Purpose, Intended Users
+    2️⃣ **Scope**: System Description, Key Functionalities
+    3️⃣ **Functional Specifications**: Features, User Roles, Transactions
+    4️⃣ **Security & Compliance**: Regulations, Data Protection
+    5️⃣ **Conclusion**: Summary, Future Enhancements
     """
 
     document = ""  # Initialize empty document
-
-    for _ in range(4):  # Generate 4 chunks of text
+    
+    for _ in range(3):  # Generate in chunks (Avoids memory crash)
         output = generator(prompt, max_length=1024, do_sample=True, temperature=0.7)
-        if output and len(output) > 0 and "generated_text" in output[0]:
+        
+        if output and "generated_text" in output[0]:
             document += output[0]['generated_text'] + "\n\n"
         else:
             return "Error: Model failed to generate text."
-
-    return document  # Return final document text
+    
+    return document
 
 # 📌 Function to Save Generated Content as PDF
 def save_to_pdf(content, filename):
     c = canvas.Canvas(filename, pagesize=letter)
     c.setFont("Helvetica", 10)
     
-    text = c.beginText(40, 750)  # Set starting position for text
-    text.setLeading(14)  # Line spacing
+    text = c.beginText(40, 750)
+    text.setLeading(14)
     
     for line in content.split("\n"):
-        text.textLine(line)  # Add text line by line
-        if text.getY() < 50:  # Start a new page when space runs out
+        text.textLine(line)
+        if text.getY() < 50:  # Handle page overflow
             c.drawText(text)
             c.showPage()
             text = c.beginText(40, 750)
             text.setLeading(14)
-
+    
     c.drawText(text)
     c.save()
 
 # 📌 Streamlit UI
 def main():
     st.title("📄 AI-Powered Functional Requirement Generator for Banking")
-
-    # Dropdown menu for selecting banking topics
+    
     banking_topics = [
-        "Core Banking System",
-        "Loan Management System",
-        "Payment Processing Gateway",
-        "Risk and Fraud Detection",
-        "Regulatory Compliance Management",
-        "Digital Banking APIs",
-        "Customer Onboarding & KYC",
-        "Treasury Management",
-        "Wealth & Portfolio Management"
+        "Core Banking System", "Loan Management System", "Payment Processing Gateway",
+        "Risk and Fraud Detection", "Regulatory Compliance Management", "Digital Banking APIs",
+        "Customer Onboarding & KYC", "Treasury Management", "Wealth & Portfolio Management"
     ]
     
     topic = st.selectbox("Select a Banking Functional Requirement Topic", banking_topics)
     
     if st.button("Generate Functional Requirement Document"):
         with st.spinner("Generating... This may take a while."):
-            content = generate_functional_requirements(topic)  # Generate content
+            content = generate_functional_requirements(topic)
             
             if "Error" in content:
                 st.error(content)
             else:
                 filename = "functional_requirement.pdf"
-                save_to_pdf(content, filename)  # Save to PDF
-                st.success("✅ Functional Requirement Document Generated!")
-                st.download_button(label="📥 Download PDF", data=open(filename, "rb"), file_name=filename, mime="application/pdf")
+                save_to_pdf(content, filename)
+                
+                st.success("✅ Document Generated!")
+                st.download_button("📥 Download PDF", data=open(filename, "rb"), file_name=filename, mime="application/pdf")
+                os.remove(filename)  # Cleanup file after download
 
 if __name__ == "__main__":
-    main()
+    main()