Update app.py

app.py

@@ -4,106 +4,80 @@ import faiss
 import pickle
 from groq import Groq
 from datasets import load_dataset
-from transformers import AutoTokenizer, pipeline
-import subprocess  # For downloading if needed
+from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Initialize Groq API
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
-# 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()
+model_name = "rajkumarrrk/dialogpt-fine-tuned-on-daily-dialog"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
+chat_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 
-# 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()
+# Load datasets
+healthcare_ds = load_dataset("harishnair04/mtsamples")
+education_ds = load_dataset("ehovy/race", "all")
+finance_ds = load_dataset("warwickai/financial_phrasebank_mirror")
 
-# FAISS Index Setup (Simplified)
-index = faiss.IndexFlatL2(768)  # Adjust dimension if needed
+# Load chat model
+chat_pipe = pipeline("text-generation", model="rajkumarrrk/dialogpt-fine-tuned-on-daily-dialog")
+
+# FAISS Index Setup
+index = faiss.IndexFlatL2(768)
 chat_history = []
 
 # Streamlit UI Setup
 st.set_page_config(page_title="AI Chatbot", layout="wide")
 st.title("🤖 AI Chatbot (Healthcare, Education & Finance)")
 
-# ... (rest of your Streamlit UI code - sidebar, input, buttons)
+# 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!")
 
 # Chat Interface
 user_input = st.text_input("💬 Ask me anything:", placeholder="Type your query here...")
 if st.button("Send"):
     if user_input:
-        # 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. Please use keywords like 'health', 'education', or 'finance'.")
-            st.stop()
-
-        # 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"
-            )
-            response = chat_completion.choices[0].message.content
-        except Exception as e:
-            st.error(f"Error generating response: {e}")
-            response = "Error generating response."
-
-        # Save and display
+        # Determine dataset based on user query (Basic CAG Implementation)
+        dataset = healthcare_ds if "health" in user_input.lower() else \
+                  education_ds if "education" in user_input.lower() else \
+                  finance_ds
+        
+        # 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
         chat_history.append(f"User: {user_input}\nBot: {response}")
         st.text_area("🤖 AI Response:", value=response, height=200)
 
-# ... (rest of your Streamlit code - chat history display, save/load)
+# Display past chats
+st.sidebar.write("\n".join(chat_history))
 
-# Persistence functions (pickle)
+# Save chat history using pickle for persistence
 def save_chat_history():
-    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}")
+    with open("chat_history.pkl", "wb") as file:
+        pickle.dump(chat_history, file)
 
 def load_chat_history():
     global chat_history
-    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}")
+    if os.path.exists("chat_history.pkl"):
+        with open("chat_history.pkl", "rb") as file:
+            chat_history = pickle.load(file)
 
 load_chat_history()
 if st.sidebar.button("Save Chat History"):
-    save_chat_history()
+    save_chat_history()
+    st.sidebar.success("Chat history saved permanently!")