Update app.py

app.py

@@ -3,43 +3,70 @@ import gspread
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from oauth2client.service_account import ServiceAccountCredentials
+from llama_index.core import VectorStoreIndex, Settings
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.schema import Document
 
-# =============== 1. Cache Google Sheets ===============
-cached_text_data = None
+# =============== 1. Cache dan Inisialisasi Index Google Sheets ===============
+cached_index = None
 
-def search_google_sheets(user_query):
+def read_google_sheets():
     try:
         scope = ["https://www.googleapis.com/auth/spreadsheets", "https://www.googleapis.com/auth/drive"]
         creds = ServiceAccountCredentials.from_json_keyfile_name("credentials.json", scope)
         client = gspread.authorize(creds)
-
+        
         SPREADSHEET_ID = "1e_cNMhwF-QYpyYUpqQh-XCw-OdhWS6EuYsoBUsVtdNg"
         sheet_names = ["datatarget", "datacuti", "dataabsen", "datalembur", "pkb"]
-
-        matched_data = []
+        
+        all_data = []
         spreadsheet = client.open_by_key(SPREADSHEET_ID)
-
+        
         for sheet_name in sheet_names:
-            sheet = spreadsheet.worksheet(sheet_name)
-            data = sheet.get_all_values()
-            for row in data:
-                row_text = " | ".join(row)
-                if user_query.lower() in row_text.lower():
-                    matched_data.append(row_text)
-
-        return "\n".join(matched_data) if matched_data else "Maaf, saya tidak menemukan informasi yang relevan."
-
+            try:
+                sheet = spreadsheet.worksheet(sheet_name)
+                data = sheet.get_all_values()
+                all_data.append(f"=== Data dari {sheet_name.upper()} ===")
+                all_data.extend([" | ".join(row) for row in data])
+                all_data.append("\n")
+            except gspread.exceptions.WorksheetNotFound:
+                all_data.append(f"❌ ERROR: Worksheet {sheet_name} tidak ditemukan.")
+        
+        return "\n".join(all_data).strip()
     except Exception as e:
         return f"❌ ERROR: {str(e)}"
 
+def initialize_index():
+    global cached_index
+    text_data = read_google_sheets()
+    document = Document(text=text_data)
+    parser = SentenceSplitter(chunk_size=100, chunk_overlap=30)
+    nodes = parser.get_nodes_from_documents([document])
+    
+    embedding = HuggingFaceEmbedding("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    Settings.embed_model = embedding
+    
+    cached_index = VectorStoreIndex(nodes)
+
+def search_google_sheets_vector(query):
+    if cached_index is None:
+        initialize_index()
+    
+    retriever = cached_index.as_retriever(similarity_top_k=3)
+    retrieved_nodes = retriever.retrieve(query)
+    
+    results = [node.text for node in retrieved_nodes]
+    return "\n".join(results) if results else "Maaf, saya tidak menemukan informasi yang relevan."
+
 # =============== 2. Load Model Transformers ===============
 def load_model():
-    model_id = "HuggingFaceH4/zephyr-7b-beta"  # Bisa ganti ke Zephyr juga
+    model_id = "HuggingFaceH4/zephyr-7b-beta"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        device_map="auto",  # ✅ Otomatis ke GPU
-        torch_dtype=torch.float16  # Ganti ke bfloat16 jika float16 bermasalah
+        device_map="auto",
+        torch_dtype=torch.float16  
     )
     pipe = pipeline(
         "text-generation",
@@ -52,7 +79,6 @@ def load_model():
     )
     return pipe
 
-
 # =============== 3. Buat Prompt dan Jawaban ===============
 def generate_prompt(user_message, context_data):
     prompt = f"""
@@ -66,10 +92,9 @@ Anda adalah chatbot HRD yang membantu karyawan memahami administrasi perusahaan.
 """
     return prompt.strip()
 
-
 # =============== 4. Generate Response ===============
 def generate_response(message, history, pipe):
-    context = search_google_sheets(message)  # 🔍 Cari hanya bagian yang relevan
+    context = search_google_sheets_vector(message)  # 🔍 Pencarian berbasis vektor
     full_prompt = generate_prompt(message, context)
     response = pipe(full_prompt)[0]["generated_text"]
 
@@ -81,6 +106,7 @@ def generate_response(message, history, pipe):
 # =============== 5. Jalankan Gradio ===============
 def main():
     pipe = load_model()
+    initialize_index()  # 🔹 Inisialisasi index sebelum chatbot berjalan
 
     def chatbot_response(message, history):
         return generate_response(message, history, pipe)