File size: 1,321 Bytes
8cf4b8e 36cac86 8cf4b8e 36cac86 8cf4b8e 36cac86 |
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 |
import os
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
# Hàm khởi tạo retriever
def load_retriever():
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore_path = "vectorstore"
# Kiểm tra nếu tồn tại index.faiss mới load
if os.path.exists(vectorstore_path) and os.path.exists(os.path.join(vectorstore_path, "index.faiss")):
retriever = FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True).as_retriever(search_kwargs={"k": 5})
print("✅ Đã load vectorstore thành công.")
else:
retriever = None
print("⚠️ Vectorstore chưa tồn tại. Hãy upload tài liệu trước khi tìm kiếm.")
return retriever
# Lần đầu load retriever
retriever = load_retriever()
# Hàm reload retriever khi thêm tài liệu
def reload_retriever():
global retriever
retriever = load_retriever()
# Hàm retrieve_docs để lấy tài liệu
def retrieve_docs(query):
if retriever:
return retriever.get_relevant_documents(query)
else:
print("⚠️ Chưa có vectorstore để tìm kiếm. Bạn cần tải tài liệu trước.")
return []
|