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 = "EleutherAI/gpt-neo-2.7B" # Change this to one of the other models based on your needs.
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1" # Change this to one of the other models based on your needs.
 
 # ---------------------------------------------------------------------------
 # Code Below to Load, Generate, and Save Functional Requirement Documents
@@ -39,82 +39,94 @@ from reportlab.pdfgen import canvas
 import os
 
 # Load Hugging Face Token from Environment (Set this in Hugging Face Spaces Secrets)
-HF_TOKEN = os.getenv("HF_TOKEN")  # Token from your Hugging Face account for accessing models.
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Model Selection: Use GPT-Neo-2.7B for generating functional requirements
-MODEL_NAME = "EleutherAI/gpt-neo-2.7B"  # Open-source GPT-Neo-2.7B model for longer text generation
+# 🏆 Recommended Model: Mistral-7B (Lightweight & Open-Source)
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
 
-# Load Model - Using Streamlit's caching to improve performance and load model only once
+# 📌 Load Model Efficiently (Caching to avoid reloading)
 @st.cache_resource
 def load_model():
     try:
-        # Initialize the Hugging Face pipeline with the selected model for text generation.
         return pipeline("text-generation", model=MODEL_NAME, token=HF_TOKEN)
     except Exception as e:
-        st.error(f"❌ Error loading model: {str(e)}")  # Error handling for failed model loading
-        return None  # Return None if model loading fails
+        st.error(f"❌ Error loading model: {str(e)}")
+        return None
 
-# Initialize the model generator
-generator = load_model()
+generator = load_model()  # Load once and reuse
 
-# Function to generate functional requirement document based on the selected topic
+# 📌 Function to Generate Functional Requirement Document
 def generate_functional_requirements(topic):
     if generator is None:
-        return "Error: Model not loaded properly."  # Return error if the model is not loaded
+        return "Error: Model failed to load."
 
-    # Define the structured prompt for generating functional requirement document
+    # Structured Prompt for Better Document Output
     prompt = f"""
-    Generate a comprehensive functional requirement document for a {topic} in the banking sector.
-    The document should be structured as follows:
+    Generate a detailed functional requirements document for {topic} in the banking sector. 
+    The document should follow this structured format:
     
-    1. **Introduction**: Provide an overview of the {topic} and its importance in the banking sector.
-    2. **Functional Requirements**: List and describe the key functional requirements that the system should fulfill.
-    3. **Non-Functional Requirements**: Discuss the performance, security, scalability, and other non-functional aspects.
-    4. **Use Cases**: Provide detailed use cases or scenarios that illustrate how the system will function.
-    5. **Data Flow Diagrams (DFD)**: Provide explanations of the system's data flows and processes (optional).
-    6. **User Interface (UI) Requirements**: Describe the UI/UX expectations for the system.
-    7. **System Architecture**: Explain the high-level architecture of the system and integration with other systems.
-    8. **Regulatory Compliance**: Detail the compliance requirements for the system (e.g., GDPR, PCI-DSS).
-    9. **Conclusion**: Summarize the document with key takeaways.
-    
-    The document should be detailed and cover all sections thoroughly, producing a functional document that is suitable for a technical audience in a banking institution.
+    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  
     """
-    
-    # Generate content using the text generation model
-    output = generator(prompt, max_length=3000, do_sample=True, temperature=0.7)  # Generate long-form content
-    generated_text = output[0]['generated_text']
-    
-    if not generated_text.strip():
-        return "No content generated. Please try again."  # Handle the case where the model generates an empty response
-    
-    return generated_text  # Return the generated content
 
-# Function to save generated content as a PDF
+    document = ""  # Initialize empty document
+
+    for _ in range(4):  # Generate 4 chunks of text
+        output = generator(prompt, max_length=1024, do_sample=True, temperature=0.7)
+        if output and len(output) > 0 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
+
+# 📌 Function to Save Generated Content as PDF
 def save_to_pdf(content, filename):
-    c = canvas.Canvas(filename, pagesize=letter)  # Initialize canvas to create a PDF
-    c.setFont("Helvetica", 10)  # Set the font for the PDF
+    c = canvas.Canvas(filename, pagesize=letter)
+    c.setFont("Helvetica", 10)
     
-    text = c.beginText(40, 750)  # Set the starting position for the text on the PDF
-    text.setLeading(14)  # Set line height for text
+    text = c.beginText(40, 750)  # Set starting position for text
+    text.setLeading(14)  # Line spacing
     
-    # Ensure content is not too large to fit in one page, splitting into lines
-    lines = content.split("\n")
-    for line in lines:
-        text.textLine(line)  # Add the line to the PDF
-        if text.getY() < 50:  # Check if the text has gone below the bottom margin
-            c.drawText(text)  # Draw the current text
-            c.showPage()  # Create a new page in the PDF
-            text = c.beginText(40, 750)  # Reset the text position for the new page
-            text.setLeading(14)  # Reset the line height
-
-    c.drawText(text)  # Draw any remaining text
-    c.save()  # Save the PDF file
-
-# Streamlit UI - User Interface for interacting with the app
+    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
+            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")  # Title for the app
+    st.title("📄 AI-Powered Functional Requirement Generator for Banking")
 
-    # Define a list of banking topics that users can select from
+    # Dropdown menu for selecting banking topics
     banking_topics = [
         "Core Banking System",
         "Loan Management System",
@@ -127,21 +139,19 @@ def main():
         "Wealth & Portfolio Management"
     ]
     
-    # Dropdown menu to select a topic
     topic = st.selectbox("Select a Banking Functional Requirement Topic", banking_topics)
     
-    # Button to trigger the document generation
     if st.button("Generate Functional Requirement Document"):
-        with st.spinner("Generating... This may take a while."):  # Show a loading spinner while generating
-            content = generate_functional_requirements(topic)  # Generate the content based on the selected topic
+        with st.spinner("Generating... This may take a while."):
+            content = generate_functional_requirements(topic)  # Generate content
+            
             if "Error" in content:
-                st.error(content)  # Show an error message if the content generation fails
+                st.error(content)
             else:
-                st.write(content[:2000])  # Display first 2000 characters for debugging
-                filename = "functional_requirement.pdf"  # Set the filename for the generated PDF
-                save_to_pdf(content, filename)  # Save the content as a PDF file
-                st.success("✅ Functional Requirement Document Generated!")  # Show success message
-                st.download_button(label="📥 Download PDF", data=open(filename, "rb"), file_name=filename, mime="application/pdf")  # Provide a download link for the PDF
+                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")
 
 if __name__ == "__main__":
-    main()  # Run the main function to start the app
+    main()