Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,26 +1,22 @@
 # personal_finance_chatbot.py
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import json
 from datetime import datetime
-import random
-import pandas as pd
-import numpy as np
 
 # Configuration
-MODEL_NAME = "ibm/granite-7b-base"  # Using IBM Granite model via HuggingFace
+MODEL_NAME = "ibm-granite/granite-7b-base"  # Correct Granite HF name
 USER_TYPES = ["student", "professional"]
-FINANCE_CATEGORIES = ["savings", "taxes", "investments", "budget", "spending"]
 
 # Initialize NLP pipeline
 @st.cache_resource
 def load_model():
-    """Load and cache the language model"""
+    """Load and cache Granite model for text generation"""
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-    return pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# User profile management
+# ------------------ USER PROFILE ------------------
 class UserProfile:
     def __init__(self, user_type, financial_goals=None, income=0, expenses=None):
         self.user_type = user_type
@@ -40,20 +36,17 @@ class UserProfile:
 
     def get_budget_summary(self):
         """Generate a budget summary"""
-        total_expenses = sum(t["amount"] for t in self.transaction_history
-                           if t["amount"] < 0)
-        total_income = sum(t["amount"] for t in self.transaction_history
-                         if t["amount"] > 0)
+        total_expenses = sum(t["amount"] for t in self.transaction_history if t["amount"] < 0)
+        total_income = sum(t["amount"] for t in self.transaction_history if t["amount"] > 0)
 
         return {
             "total_income": total_income,
             "total_expenses": abs(total_expenses),
-            "net_savings": total_income + total_expenses,
+            "net_savings": total_income + total_expenses,  # total_expenses is negative already
             "category_breakdown": self._get_category_breakdown()
         }
 
     def _get_category_breakdown(self):
-        """Get expense breakdown by category"""
         breakdown = {}
         for t in self.transaction_history:
             if t["amount"] < 0:
@@ -61,7 +54,7 @@ class UserProfile:
                 breakdown[cat] = breakdown.get(cat, 0) + abs(t["amount"])
         return breakdown
 
-# Chatbot core functionality
+# ------------------ CHATBOT CORE ------------------
 class FinanceChatbot:
     def __init__(self):
         self.nlp = load_model()
@@ -69,96 +62,83 @@ class FinanceChatbot:
         self.current_user = None
 
     def set_user(self, user_id, user_type):
-        """Initialize or switch user profile"""
         if user_id not in self.user_profiles:
             self.user_profiles[user_id] = UserProfile(user_type=user_type)
         self.current_user = user_id
 
     def generate_response(self, query):
-        """Generate context-aware response to user query"""
         if not self.current_user:
-            return "Please identify yourself first. Are you a student or professional?"
+            return "⚠️ Please set up your profile first (student or professional)."
 
         profile = self.user_profiles[self.current_user]
         context = self._build_context(profile)
 
-        # Adjust tone based on user type
         tone_instruction = (
-            "Use simple, encouraging language suitable for a student."
+            "Use simple, encouraging language for a student."
             if profile.user_type == "student" else
-            "Use professional, concise language suitable for a financial professional."
+            "Use concise, professional language for a working professional."
         )
 
         prompt = f"""
-        You are a helpful financial assistant. {tone_instruction}
-        User profile context: {context}
-        Current query: {query}
-
-        Provide a helpful response that:
-        1. Directly answers the financial question
-        2. Offers 1-2 specific actionable suggestions when appropriate
-        3. Maintains a {profile.user_type}-appropriate tone
-        4. Is concise (3 sentences or less unless more detail is requested)
+        You are an AI-powered financial assistant.
+        User profile: {context}
+        Instruction: {tone_instruction}
+
+        User asked: "{query}"
+
+        Respond with:
+        1. Direct and clear answer
+        2. 1-2 actionable suggestions
+        3. Keep it under 3 sentences unless more detail is needed.
         """
 
         try:
-            result = self.nlp(prompt, max_length=200, num_return_sequences=1)
-            response = result[0]['generated_text'].strip()
-
-            # Post-process to extract clean response
-            if "Response:" in response:
-                response = response.split("Response:")[-1].strip()
+            result = self.nlp(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
+            response = result[0]['generated_text'].replace(prompt, "").strip()
             return response
         except Exception as e:
-            return f"Sorry, I encountered an error processing your request. {str(e)}"
+            return f"❌ Error: {str(e)}"
 
     def _build_context(self, profile):
-        """Build context string from user profile"""
         budget = profile.get_budget_summary()
         return json.dumps({
             "user_type": profile.user_type,
             "income": profile.income,
             "net_savings": budget["net_savings"],
             "top_expenses": sorted(budget["category_breakdown"].items(),
-                                 key=lambda x: x[1], reverse=True)[:3],
-            "recent_transactions": profile.transaction_history[-3:] if profile.transaction_history else []
+                                   key=lambda x: x[1], reverse=True)[:3],
+            "recent_transactions": profile.transaction_history[-3:]
         })
 
     def analyze_spending(self):
-        """Generate spending insights"""
         if not self.current_user:
-            return "No user profile selected"
+            return "⚠️ No user profile selected."
 
         profile = self.user_profiles[self.current_user]
         budget = profile.get_budget_summary()
-
         if not budget["category_breakdown"]:
-            return "No spending data available for analysis"
+            return "ℹ️ No spending data yet."
 
         prompt = f"""
-        Analyze this spending data and provide insights:
-        {json.dumps(budget['category_breakdown'])}
-
-        User type: {profile.user_type}
-        Income: {profile.income}
+        Analyze the spending breakdown: {json.dumps(budget['category_breakdown'])}.
+        User type: {profile.user_type}, Income: {profile.income}.
 
         Provide:
-        1. 1 key insight about spending patterns
-        2. 1 specific suggestion for optimization
-        3. Formatted for {profile.user_type} audience
+        1. One key spending insight
+        2. One actionable saving tip
+        3. Tone adapted for {profile.user_type}
         """
 
         try:
-            result = self.nlp(prompt, max_length=150)
-            return result[0]['generated_text'].strip()
+            result = self.nlp(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)
+            return result[0]['generated_text'].replace(prompt, "").strip()
         except Exception as e:
-            return f"Error generating insights: {str(e)}"
+            return f"❌ Error: {str(e)}"
 
-# Streamlit UI
+# ------------------ STREAMLIT UI ------------------
 def main():
     st.set_page_config(page_title="Personal Finance Chatbot", layout="wide")
 
-    # Initialize chatbot
     if 'chatbot' not in st.session_state:
         st.session_state.chatbot = FinanceChatbot()
     if 'user_id' not in st.session_state:
@@ -166,10 +146,10 @@ def main():
     if 'messages' not in st.session_state:
         st.session_state.messages = []
 
-    # Sidebar for user setup
+    # Sidebar
     with st.sidebar:
-        st.title("User Profile")
-        user_id = st.text_input("Your ID (any name)", value=st.session_state.get('user_id', ''))
+        st.title("👤 User Profile")
+        user_id = st.text_input("Your ID", value=st.session_state.get('user_id', ''))
         user_type = st.selectbox("I am a:", USER_TYPES)
 
         if st.button("Save Profile"):
@@ -178,72 +158,47 @@ def main():
             st.success(f"Profile saved as {user_type}")
 
         st.divider()
-        st.subheader("Quick Actions")
+        st.subheader("📊 Quick Actions")
 
         if st.session_state.user_id:
-            col1, col2 = st.columns(2)
-            with col1:
-                if st.button("Budget Summary"):
-                    profile = st.session_state.chatbot.user_profiles[st.session_state.user_id]
-                    summary = profile.get_budget_summary()
-                    st.session_state.messages.append(
-                        {"role": "assistant", "content": f"### Your Budget Summary\n\n"
-                                                        f"- **Income**: ${summary['total_income']:.2f}\n"
-                                                        f"- **Expenses**: ${summary['total_expenses']:.2f}\n"
-                                                        f"- **Net Savings**: ${summary['net_savings']:.2f}\n\n"
-                                                        f"Top Expenses: {', '.join([f'{k}(${v:.2f})' for k,v in summary['category_breakdown'].items()])}"
-                        }
-                    )
-            with col2:
-                if st.button("Spending Insights"):
-                    insights = st.session_state.chatbot.analyze_spending()
-                    st.session_state.messages.append(
-                        {"role": "assistant", "content": f"### Spending Insights\n\n{insights}"}
-                    )
-
-    # Main chat interface
+            if st.button("View Budget Summary"):
+                profile = st.session_state.chatbot.user_profiles[st.session_state.user_id]
+                summary = profile.get_budget_summary()
+                st.session_state.messages.append(
+                    {"role": "assistant", "content":
+                        f"### Budget Summary\n- Income: ${summary['total_income']:.2f}\n"
+                        f"- Expenses: ${summary['total_expenses']:.2f}\n"
+                        f"- Net Savings: ${summary['net_savings']:.2f}\n"
+                        f"- Top Expenses: {summary['category_breakdown']}"}
+                )
+
+            if st.button("Get Spending Insights"):
+                insights = st.session_state.chatbot.analyze_spending()
+                st.session_state.messages.append(
+                    {"role": "assistant", "content": f"### Spending Insights\n{insights}"}
+                )
+
+    # Main chat
     st.title("💰 Personal Finance Chatbot")
-    st.write("Ask me about savings, taxes, investments, or budgeting!")
+    st.write("Ask me about **savings, taxes, investments, or budgeting!**")
 
-    # Display chat messages
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-    # Chat input
-    if prompt := st.chat_input("What would you like to know about your finances?"):
+    if prompt := st.chat_input("Type your financial question here..."):
         if not st.session_state.user_id:
-            st.error("Please set up your profile first!")
+            st.error("⚠️ Please set up your profile first!")
         else:
-            # Add user message to chat history
             st.session_state.messages.append({"role": "user", "content": prompt})
             with st.chat_message("user"):
                 st.markdown(prompt)
 
-            # Generate and display assistant response
-            with st.spinner("Thinking..."):
+            with st.spinner("💡 Thinking..."):
                 response = st.session_state.chatbot.generate_response(prompt)
                 with st.chat_message("assistant"):
                     st.markdown(response)
                 st.session_state.messages.append({"role": "assistant", "content": response})
 
-    # Sample transactions for demo
-    if st.session_state.user_id and st.checkbox("Load sample data (demo only)"):
-        profile = st.session_state.chatbot.user_profiles[st.session_state.user_id]
-        sample_data = [
-            (1200, "income", "Monthly salary"),
-            (-400, "housing", "Rent payment"),
-            (-200, "food", "Groceries"),
-            (-150, "transportation", "Public transport"),
-            (-100, "entertainment", "Movies"),
-            (50, "investments", "Stock purchase"),
-            (-80, "utilities", "Electricity bill")
-        ]
-
-        for amount, category, desc in sample_data:
-            profile.add_transaction(amount, category, desc)
-
-        st.success("Loaded sample transactions! Try asking about your budget or spending.")
-
 if __name__ == "__main__":
     main()