Update app.py

app.py

@@ -1,33 +1,83 @@
 import streamlit as st
 import requests
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+import numpy as np
+import faiss
 
-st.set_page_config(
-        page_title="IOGPT",
-        page_icon="🤖",
-        menu_items={}  # This helps hide the menu
-    )
-
-# Hide Streamlit menu and footer
-hide_menu_style = """
-        <style>
-        #MainMenu {visibility: hidden;}
-        footer {visibility: hidden;}
-        </style>
-        """
-st.markdown(hide_menu_style, unsafe_allow_html=True)
+class CompanyKnowledgeBase:
+    def __init__(self, dataset_name="JustKiddo/IODataset"):
+        # Load dataset from Hugging Face
+        try:
+            self.dataset = load_dataset(dataset_name)['train']
+            
+            # Initialize semantic search
+            self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
+            
+            # Prepare embeddings for all questions
+            self.embeddings = self.model.encode([
+                q for entry in self.dataset 
+                for q in entry['questions']
+            ])
+            
+            # Create FAISS index for efficient similarity search
+            self.index = faiss.IndexFlatL2(self.embeddings.shape[1])
+            self.index.add(self.embeddings)
+            
+            # Prepare a mapping of embeddings to answers
+            self.question_to_answer = {}
+            for entry in self.dataset:
+                for question in entry['questions']:
+                    self.question_to_answer[question] = entry['answer']
+        
+        except Exception as e:
+            st.error(f"Error loading knowledge base: {e}")
+            self.dataset = None
+            self.embeddings = None
+            self.index = None
+            self.question_to_answer = {}
+    
+    def find_answer(self, query, threshold=0.8):
+        if not self.dataset:
+            return None
+        
+        try:
+            # Embed the query
+            query_embedding = self.model.encode([query])
+            
+            # Search for similar questions
+            D, I = self.index.search(query_embedding, 1)
+            
+            # If similarity is high enough, return the corresponding answer
+            if D[0][0] < threshold:
+                # Find the matched question
+                matched_question = list(self.question_to_answer.keys())[I[0][0]]
+                return self.question_to_answer[matched_question]
+        
+        except Exception as e:
+            st.error(f"Error in semantic search: {e}")
+        
+        return None
 
 class VietnameseChatbot:
     def __init__(self):
-        self.api_key = st.secrets["GROQ_API_KEY"]  # Store your API key in Huggingface Secrets
+        self.api_key = st.secrets["GROQ_API_KEY"]
         self.api_url = "https://api.groq.com/openai/v1/chat/completions"
         self.headers = {
-            "Content-Type": "application/json",
+            "Content-Type": "application/json", 
             "Authorization": f"Bearer {self.api_key}"
         }
-
+        # Initialize company knowledge base
+        self.company_kb = CompanyKnowledgeBase()
+    
     def get_response(self, user_query):
+        # First, check company knowledge base
+        company_answer = self.company_kb.find_answer(user_query)
+        if company_answer:
+            return company_answer
+        
+        # If no company-specific answer, proceed with original API call
         try:
-            # Add a system message to guide the model's response
             payload = {
                 "model": "llama-3.2-3b-preview",
                 "messages": [
@@ -38,19 +88,25 @@ class VietnameseChatbot:
                     {"role": "user", "content": user_query}
                 ]
             }
+            
             response = requests.post(
-                self.api_url, headers=self.headers, json=payload
+                self.api_url, 
+                headers=self.headers, 
+                json=payload
             )
+            
             if response.status_code == 200:
                 return response.json()['choices'][0]['message']['content']
             else:
                 print(f"API Error: {response.status_code}")
                 print(f"Response: {response.text}")
                 return "Đã xảy ra lỗi khi kết nối với API. Xin vui lòng thử lại."
+        
         except Exception as e:
             print(f"Response generation error: {e}")
             return "Đã xảy ra lỗi. Xin vui lòng thử lại."
 
+# Cached initialization of chatbot
 @st.cache_resource
 def initialize_chatbot():
     return VietnameseChatbot()
@@ -61,32 +117,32 @@ def main():
 
     # Initialize chatbot using cached initialization
     chatbot = initialize_chatbot()
-    
+
     # Chat history in session state
     if 'messages' not in st.session_state:
         st.session_state.messages = []
-    
+
     # Display chat messages
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
-    
+
     # User input
     if prompt := st.chat_input("Hãy nói gì đó..."):
         # Add user message to chat history
         st.session_state.messages.append({"role": "user", "content": prompt})
-        
+
         # Display user message
         with st.chat_message("user"):
             st.markdown(prompt)
-        
+
         # Get chatbot response
         response = chatbot.get_response(prompt)
-        
+
         # Display chatbot response
         with st.chat_message("assistant"):
             st.markdown(response)
-        
+
         # Add assistant message to chat history
         st.session_state.messages.append({"role": "assistant", "content": response})