Spaces:
Runtime error
Runtime error
File size: 1,977 Bytes
ec99d5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import os
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
import json
# Paths for persistence
INDEX_PATH = os.path.join(os.path.dirname(__file__), "faiss_index.bin")
STORE_PATH = os.path.join(os.path.dirname(__file__), "doc_store.json")
# Initialize model
# all-MiniLM-L6-v2 is small, fast, and great for sentence embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')
dimension = 384
# Map vector ID to text snippet
document_store = {}
current_id = 0
if os.path.exists(INDEX_PATH):
index = faiss.read_index(INDEX_PATH)
else:
index = faiss.IndexFlatL2(dimension)
if os.path.exists(STORE_PATH):
with open(STORE_PATH, "r", encoding="utf-8") as f:
stored_data = json.load(f)
# Convert string keys back to int
document_store = {int(k): v for k, v in stored_data.items()}
current_id = max(document_store.keys()) + 1 if document_store else 0
def save_state():
faiss.write_index(index, INDEX_PATH)
with open(STORE_PATH, "w", encoding="utf-8") as f:
json.dump(document_store, f, ensure_ascii=False, indent=2)
def add_documents(texts: list[str]):
global current_id
if not texts:
return
embeddings = model.encode(texts)
index.add(np.array(embeddings).astype('float32'))
for i, text in enumerate(texts):
document_store[current_id + i] = text
current_id += len(texts)
# Save to disk after adding
save_state()
def search_library(query: str, top_k: int = 3) -> list[str]:
if index.ntotal == 0:
return ["Perpustakaan Anda masih kosong. Tidak ada data jurnal yang bisa dicari."]
query_vector = model.encode([query])
distances, indices = index.search(np.array(query_vector).astype('float32'), top_k)
results = []
for idx in indices[0]:
if idx != -1 and idx in document_store:
results.append(document_store[idx])
return results
|