Update app.py

app.py

@@ -1,4 +1,3 @@
-# Import libraries
 from sentence_transformers import SentenceTransformer
 from PyPDF2 import PdfReader
 import tiktoken
@@ -13,7 +12,7 @@ import pickle
 # == Buat folder models ==
 os.makedirs("models", exist_ok=True)
 
-# == Load API Key dari File (Hindari Hardcoded Key) ==
+# == Load API Key dari File ==
 def load_api_key():
     with open("config.json", "r") as f:
         config = json.load(f)
@@ -23,7 +22,6 @@ GROQ_API_KEY = load_api_key()
 
 # == Ekstraksi Teks dari PDF ==
 def extract_text_from_pdf(pdf_file: str) -> str:
-    """Ekstrak teks dari PDF dan gabungkan menjadi satu string."""
     with open(pdf_file, 'rb') as pdf:
         reader = PdfReader(pdf)
         text = " ".join(page.extract_text() or "" for page in reader.pages)
@@ -31,8 +29,7 @@ def extract_text_from_pdf(pdf_file: str) -> str:
 
 # == Chunking Teks ==
 def chunk_text(text: str, max_tokens: int = 512) -> list:
-    """Membagi teks menjadi chunk berdasarkan token menggunakan tokenizer OpenAI."""
-    tokenizer = tiktoken.get_encoding("cl100k_base")  # Gunakan tokenizer OpenAI
+    tokenizer = tiktoken.get_encoding("cl100k_base")
     tokens = tokenizer.encode(text)
 
     chunks = []
@@ -43,41 +40,38 @@ def chunk_text(text: str, max_tokens: int = 512) -> list:
 
     return chunks
 
-# == Embedding dengan Ollama ==
+# == Embedding dengan SentenceTransformer ==
+model = SentenceTransformer('all-MiniLM-L6-v2')  # Global model
+
 def get_embedding(text: str):
-    """Mendapatkan embedding dari teks menggunakan Ollama."""
-    model = SentenceTransformer('all-MiniLM-L6-v2')
-    embedding = model.encode("This is a test sentence.")
-    return np.array(embedding["embeddings"][0], dtype=np.float32)  # Pastikan mengambil list pertama
-
-# == Simpan Embedding ke FAISS ==
-d = 1024  # Dimensi embedding dari model `mxbai-embed-large`
-index = faiss.IndexFlatL2(d)  # Inisialisasi FAISS Index
+    return np.array(model.encode(text), dtype=np.float32)
+
+# == Setup FAISS ==
+d = 384  # Dimensi embedding sesuai dengan model
+index = faiss.IndexFlatL2(d)
 text_chunks = []
 
 def add_to_db(text_chunks_local):
-    """Menambahkan embedding ke FAISS."""
     global text_chunks
-    text_chunks = text_chunks_local  # Simpan chunk ke global var
-    embeddings = np.array([get_embedding(text) for text in text_chunks], dtype=np.float32)
+    text_chunks = text_chunks_local
+    embeddings = np.array([get_embedding(text) for text in text_chunks], dtype=np.float32).reshape(-1, d)
     index.add(embeddings)
 
 def search_db(query, k=5):
-    """Melakukan pencarian query dalam FAISS Index."""
+    if index.ntotal == 0:
+        return ["Database masih kosong, silakan tambahkan data."]
+    
     query_embedding = np.array([get_embedding(query)], dtype=np.float32).reshape(1, -1)
     distances, indices = index.search(query_embedding, k)
-    return [text_chunks[i] for i in indices[0]]  # Ambil teks chunk yang relevan
+    return [text_chunks[i] for i in indices[0] if i < len(text_chunks)]
 
 def save_to_faiss(index_path="vector_index.faiss"):
-    """Menyimpan FAISS index ke file."""
     faiss.write_index(index, index_path)
 
 def load_faiss(index_path="vector_index.faiss"):
-    """Memuat kembali FAISS index dari file."""
     global index
     index = faiss.read_index(index_path)
 
-# == Simpan dan Load Model Embedding ==
 def save_embeddings(embeddings_path="models/embeddings.pkl"):
     with open(embeddings_path, "wb") as f:
         pickle.dump(index, f)
@@ -91,7 +85,6 @@ def load_embeddings(embeddings_path="models/embeddings.pkl"):
 client = groq.Client(api_key=GROQ_API_KEY)
 
 def query_llama(prompt):
-    """Menggunakan LLaMA untuk menjawab pertanyaan dengan konteks yang diberikan."""
     response = client.chat.completions.create(
         model="llama3-8b-8192",
         messages=[{"role": "user", "content": prompt}],
@@ -102,14 +95,12 @@ def query_llama(prompt):
 # == Main Workflow ==
 if __name__ == '__main__':
     pdf_text = extract_text_from_pdf('dini_anggriyani_synthetic_data.pdf')
-    text_chunks = chunk_text(pdf_text, max_tokens=1024)  # Sesuaikan dengan LLaMA
+    text_chunks = chunk_text(pdf_text, max_tokens=1024)
 
-    # Tambahkan ke database FAISS
     add_to_db(text_chunks)
-    save_to_faiss()  # Simpan FAISS index
-    save_embeddings()  
+    save_to_faiss()
+    save_embeddings()
 
-    # Tes pencarian RAG
     retrieved_chunks = search_db("Apa isi dokumen ini?")
     context = "\n".join(retrieved_chunks)
     
@@ -117,9 +108,9 @@ if __name__ == '__main__':
     answer = query_llama(prompt)
     print(answer)
 
-# == Buat Chatbot Interface ==
+# == Chatbot Interface ==
 def chatbot_interface(user_query):
-    retrieved_chunks = search_db(user_query)  # Sudah berupa teks
+    retrieved_chunks = search_db(user_query)
     context = "\n".join(retrieved_chunks)
     
     prompt = f"Gunakan informasi berikut untuk menjawab:\n{context}\n\nPertanyaan: {user_query}"