Update app.py

app.py

@@ -5,25 +5,38 @@ import pickle
 from groq import Groq
 from datasets import load_dataset
 from transformers import AutoTokenizer, pipeline
+import subprocess  # For downloading if needed
 
 # Initialize Groq API
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
-# Load datasets
-try:  # Handle potential dataset loading errors
+# Download model (if necessary - try requirements.txt first)
+try:
+    # Try loading directly (after requirements.txt)
+    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")
+    print("Model loaded successfully (direct load).") # Check in logs
+except Exception as e:
+    try:
+        # Fallback: Download using subprocess (less preferred)
+        print("Trying to download model...") # Check in logs
+        subprocess.run(["transformers-cli", "download", "rajkumarrrk/dialogpt-fine-tuned-on-daily-dialog"], check=True) # Updated download command
+        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")
+        print("Model downloaded and loaded successfully (subprocess).") # Check in logs
+    except Exception as download_e:
+        st.error(f"Error loading/downloading chat model: {e}. Download error: {download_e}")
+        st.stop()
+
+
+
+# Load datasets (with error handling)
+try:
     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 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 (Simplified)
@@ -40,28 +53,32 @@ st.title("🤖 AI Chatbot (Healthcare, Education & Finance)")
 user_input = st.text_input("💬 Ask me anything:", placeholder="Type your query here...")
 if st.button("Send"):
     if user_input:
-        # 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 if "finance" in user_input.lower() else None #Handle no dataset match
-        
+        # Dataset Selection (Improved)
+        dataset = None
+        if "health" in user_input.lower():
+            dataset = healthcare_ds
+        elif "education" in user_input.lower():
+            dataset = education_ds
+        elif "finance" in user_input.lower():
+            dataset = finance_ds
+
         if dataset is None:
-            st.warning("No relevant dataset found for your query.")
+            st.warning("No relevant dataset found for your query. Please use keywords like 'health', 'education', or 'finance'.")
             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
+        # RAG: Retrieve (Simplified and safer)
+        retrieved_data = dataset['train'][0]['text'] if dataset and len(dataset['train']) > 0 and 'text' in dataset['train'][0] else "No relevant data retrieved."
 
         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
+                model="llama-3.3-70b-versatile"
             )
             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
+            response = "Error generating response."
 
         # Save and display
         chat_history.append(f"User: {user_input}\nBot: {response}")
@@ -87,7 +104,6 @@ def load_chat_history():
     except Exception as e:
         st.sidebar.warning(f"Error loading chat history (may be corrupted): {e}")
 
-
-load_chat_history()  # Load on startup
+load_chat_history()
 if st.sidebar.button("Save Chat History"):
     save_chat_history()