Update app.py

app.py

@@ -4,72 +4,90 @@ import faiss
 import pickle
 from groq import Groq
 from datasets import load_dataset
-from transformers import pipeline
+from transformers import AutoTokenizer, pipeline
 
 # Initialize Groq API
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
 # Load datasets
-healthcare_ds = load_dataset("harishnair04/mtsamples")
-education_ds = load_dataset("ehovy/race", "all")
-finance_ds = load_dataset("warwickai/financial_phrasebank_mirror")
+try:  # Handle potential dataset loading errors
+    healthcare_ds = load_dataset("harishnair04/mtsamples")
+    education_ds = load_dataset("ehovy/race", "all")
+    finance_ds = load_dataset("warwickai/financial_phrasebank_mirror")
+except Exception as e:
+    st.error(f"Error loading datasets: {e}")
+    st.stop()  # Stop execution if datasets fail to load
 
-# Load chat model
-chat_pipe = pipeline("text-generation", model="rajkumarrrk/dialogpt-fine-tuned-on-daily-dialog")
+# Load chat model and tokenizer (with error handling and cache)
+try:
+    tokenizer = AutoTokenizer.from_pretrained("rajkumarrrk/dialogpt-fine-tuned-on-daily-dialog", cache_dir="./.cache")
+    chat_pipe = pipeline("text-generation", model="rajkumarrrk/dialogpt-fine-tuned-on-daily-dialog", tokenizer=tokenizer, cache_dir="./.cache")
+except Exception as e:
+    st.error(f"Error loading chat model: {e}")
+    st.stop()
 
-# FAISS Index Setup
-index = faiss.IndexFlatL2(768)
+# FAISS Index Setup (Simplified)
+index = faiss.IndexFlatL2(768)  # Adjust dimension if needed
 chat_history = []
 
 # Streamlit UI Setup
 st.set_page_config(page_title="AI Chatbot", layout="wide")
 st.title("🤖 AI Chatbot (Healthcare, Education & Finance)")
 
-# Sidebar for chat history
-st.sidebar.title("📜 Chat History")
-if st.sidebar.button("Download Chat History"):
-    with open("chat_history.txt", "w") as file:
-        file.write("\n".join(chat_history))
-    st.sidebar.success("Chat history saved!")
+# ... (rest of your Streamlit UI code - sidebar, input, buttons)
 
 # Chat Interface
 user_input = st.text_input("💬 Ask me anything:", placeholder="Type your query here...")
 if st.button("Send"):
     if user_input:
-        # Determine dataset based on user query (Basic CAG Implementation)
+        # Determine dataset (Basic CAG)
         dataset = healthcare_ds if "health" in user_input.lower() else \
                   education_ds if "education" in user_input.lower() else \
-                  finance_ds
+                  finance_ds if "finance" in user_input.lower() else None #Handle no dataset match
         
-        # RAG: Retrieve relevant data
-        retrieved_data = dataset['train'][0]  # Simplified retrieval
-        
-        # Generate response using Llama via Groq API
-        chat_completion = client.chat.completions.create(
-            messages=[{"role": "user", "content": f"{user_input} {retrieved_data}"}],
-            model="llama-3.3-70b-versatile"
-        )
-        response = chat_completion.choices[0].message.content
-        
-        # Save chat to FAISS and display
+        if dataset is None:
+            st.warning("No relevant dataset found for your query.")
+            st.stop()
+
+        # RAG: Retrieve (Simplified)
+        retrieved_data = dataset['train'][0] if dataset and len(dataset['train']) > 0 else "No relevant data retrieved." #Check dataset is not empty
+
+        try:
+            # Generate response (Groq)
+            chat_completion = client.chat.completions.create(
+                messages=[{"role": "user", "content": f"{user_input} {retrieved_data}"}],
+                model="llama-3.3-70b-versatile" #Ensure model name is correct
+            )
+            response = chat_completion.choices[0].message.content
+        except Exception as e:
+            st.error(f"Error generating response: {e}")
+            response = "Error generating response." #Provide default response in case of error
+
+        # Save and display
         chat_history.append(f"User: {user_input}\nBot: {response}")
         st.text_area("🤖 AI Response:", value=response, height=200)
 
-# Display past chats
-st.sidebar.write("\n".join(chat_history))
+# ... (rest of your Streamlit code - chat history display, save/load)
 
-# Save chat history using pickle for persistence
+# Persistence functions (pickle)
 def save_chat_history():
-    with open("chat_history.pkl", "wb") as file:
-        pickle.dump(chat_history, file)
+    try:
+        with open("chat_history.pkl", "wb") as file:
+            pickle.dump(chat_history, file)
+        st.sidebar.success("Chat history saved permanently!")
+    except Exception as e:
+        st.sidebar.error(f"Error saving chat history: {e}")
 
 def load_chat_history():
     global chat_history
-    if os.path.exists("chat_history.pkl"):
-        with open("chat_history.pkl", "rb") as file:
-            chat_history = pickle.load(file)
+    try:
+        if os.path.exists("chat_history.pkl"):
+            with open("chat_history.pkl", "rb") as file:
+                chat_history = pickle.load(file)
+    except Exception as e:
+        st.sidebar.warning(f"Error loading chat history (may be corrupted): {e}")
+
 
-load_chat_history()
+load_chat_history()  # Load on startup
 if st.sidebar.button("Save Chat History"):
-    save_chat_history()
-    st.sidebar.success("Chat history saved permanently!")
+    save_chat_history()