File size: 1,453 Bytes
0fe75fc |
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 |
import os
from pathlib import Path
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
# Text splitter
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50
)
# Embedding model
embeddings = HuggingFaceBgeEmbeddings(
model_name="BAAI/bge-small-en",
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": True}
)
VECTORSTORE_DIR = "user_vectorstores"
os.makedirs(VECTORSTORE_DIR, exist_ok=True)
def ingest_report(user_id: str, report_text: str):
# Split into documents
documents = text_splitter.create_documents([report_text])
# Create FAISS vectorstore
vectorstore = FAISS.from_documents(documents, embeddings)
# Save to disk
user_path = Path(VECTORSTORE_DIR) / f"{user_id}_faiss"
vectorstore.save_local(str(user_path))
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
return vectorstore, retriever
def get_user_retriever(user_id: str):
user_path = Path(VECTORSTORE_DIR) / f"{user_id}_faiss"
if not user_path.exists():
raise FileNotFoundError(f"Vectorstore for user {user_id} not found.")
vectorstore = FAISS.load_local(str(user_path), embeddings, allow_dangerous_deserialization=True)
return vectorstore.as_retriever(search_kwargs={"k": 3})
|