Update app.py

app.py

@@ -1,39 +1,97 @@
-from transformers import pipeline
+# ---------------------------------------------------------------------------
+# Model Recommendations and Descriptions:
+# ---------------------------------------------------------------------------
+
+# 1. **DistilGPT-2**: A distilled, lighter version of GPT-2.
+#    - **Size**: ~500 MB (smaller and faster than full GPT-2).
+#    - **Performance**: Great for text generation tasks with reduced inference time.
+#    - **Use case**: If you need a lightweight, efficient model with decent quality.
+#    - **Hugging Face Model**: "distilgpt2"
+#
+# 2. **GPT-Neo 1.3B**: A smaller alternative to GPT-3 with reasonable performance.
+#    - **Size**: ~5 GB.
+#    - **Performance**: Powerful text generation model with good results on a wide variety of tasks.
+#    - **Use case**: Ideal for slightly larger tasks where better quality is needed than DistilGPT-2.
+#    - **Hugging Face Model**: "EleutherAI/gpt-neo-1.3B"
+#
+# 3. **Mistral 7B**: Open-source model with smaller weights.
+#    - **Size**: ~4.5 GB.
+#    - **Performance**: Comparable to larger models, with good speed and quality.
+#    - **Use case**: Ideal for high-quality generation at a relatively smaller size.
+#    - **Hugging Face Model**: "mistralai/Mistral-7B-Instruct-v0.1"
+#
+# 4. **TinyBERT or MiniLM**: Excellent for lighter tasks, even smaller models.
+#    - **Size**: ~100 MB.
+#    - **Performance**: Great for tasks like classification or sentence embeddings but not suitable for long text generation.
+#    - **Use case**: Perfect for applications requiring minimal memory and fast processing but not for full-fledged generation.
+#    - **Hugging Face Model**: "sentence-transformers/all-MiniLM-L6-v2"
+
+# Choose your model from the above options:
+MODEL_NAME = "distilgpt2"  # Change this to one of the other models based on your needs.
+
+# ---------------------------------------------------------------------------
+# Code Below to Load, Generate, and Save Functional Requirement Documents
+# ---------------------------------------------------------------------------
+
 import streamlit as st
+from transformers import pipeline
+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)
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN")  # Token from your Hugging Face account for accessing models.
 
-# Select model: Use a quantized, smaller model like DistilGPT-2
-MODEL_NAME = "distilgpt2"  # Change to any of the small models mentioned above
-
-# Load Model
+# Load Model - Using Streamlit's caching to improve performance and load model only once
 @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)}")
-        return None
+        st.error(f"❌ Error loading model: {str(e)}")  # Error handling for failed model loading
+        return None  # Return None if model loading fails
 
+# Initialize the model generator
 generator = load_model()
 
-# Function to Generate Functional Requirement Document
+# Function to generate functional requirement document based on the selected topic
 def generate_functional_requirements(topic):
     if generator is None:
-        return "Error: Model not loaded properly."
+        return "Error: Model not loaded properly."  # Return error if the model is not loaded
 
+    # Define the prompt based on the selected topic for generating functional requirements
     prompt = f"Generate a comprehensive functional requirement document for {topic} in the banking sector."
     
-    # Generate content
-    output = generator(prompt, max_length=800, do_sample=True, temperature=0.8)
-    return output[0]['generated_text']
+    # Use the text generation model to generate content
+    output = generator(prompt, max_length=800, do_sample=True, temperature=0.8)  # Generate text with certain constraints
+    return output[0]['generated_text']  # Return the generated text
+
+# Function to save generated content as a 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
+    
+    text = c.beginText(40, 750)  # Set the starting position for the text on the PDF
+    text.setLeading(14)  # Set line height for text
+
+    # Iterate through each line of the content and add it to the PDF
+    for line in content.split("\n"):
+        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
+# Streamlit UI - User Interface for interacting with the app
 def main():
-    st.title("📄 AI-Powered Functional Requirement Generator for Banking")
+    st.title("📄 AI-Powered Functional Requirement Generator for Banking")  # Title for the app
 
+    # Define a list of banking topics that users can select from
     banking_topics = [
         "Core Banking System",
         "Loan Management System",
@@ -46,18 +104,20 @@ 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."):
-            content = generate_functional_requirements(topic)
+        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
             if "Error" in content:
-                st.error(content)
+                st.error(content)  # Show an error message if the content generation fails
             else:
-                filename = "functional_requirement.pdf"
-                save_to_pdf(content, filename)
-                st.success("✅ Functional Requirement Document Generated!")
-                st.download_button(label="📥 Download PDF", data=open(filename, "rb"), file_name=filename, mime="application/pdf")
+                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
 
 if __name__ == "__main__":
-    main()
+    main()  # Run the main function to start the app