File size: 2,869 Bytes
e92d49a d35bd88 e92d49a d35bd88 e92d49a 8017059 e92d49a 6f1e05b e92d49a | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from __future__ import annotations
import logging
from typing import Optional, Iterable, List
from langchain_core.embeddings import Embeddings
from langchain_community.vectorstores import FAISS
from .config import settings
from .splitter import documents
logger = logging.getLogger(__name__)
# Directory where FAISS index is persisted
_VECTORSTORE_DIR = settings.vectorstore_dir
_INDEX_NAME = settings.vectorstore_index_name
_VECTORSTORE_DIR.mkdir(parents=True, exist_ok=True)
class _LazyHuggingFaceEmbeddings(Embeddings):
def __init__(self, model_name: str, encode_kwargs: Optional[dict] = None) -> None:
self.model_name = model_name
self.encode_kwargs = encode_kwargs or {}
self._real = None
def _init(self) -> None:
if self._real is None:
from langchain_huggingface import HuggingFaceEmbeddings
self._real = HuggingFaceEmbeddings(
model_name=self.model_name, encode_kwargs=self.encode_kwargs
)
async def aembed_documents(self, texts: Iterable[str]) -> List[List[float]]:
self._init()
return await self._real.aembed_documents(texts) # type: ignore[attr-defined]
def embed_documents(self, texts: Iterable[str]) -> List[List[float]]:
self._init()
return self._real.embed_documents(texts) # type: ignore[attr-defined]
def embed_query(self, text: str) -> List[float]:
self._init()
return self._real.embed_query(text) # type: ignore[attr-defined]
async def aembed_query(self, text: str) -> List[float]:
self._init()
return await self._real.aembed_query(text) # type: ignore[attr-defined]
_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
_embeddings = _LazyHuggingFaceEmbeddings(_EMBEDDING_MODEL, encode_kwargs={"normalize_embeddings": True})
_vectorstore: Optional[FAISS] = None
def get_vectorstore() -> FAISS:
"""Return a persisted FAISS vectorstore, building it on first run only.
This attempts to load an existing index from disk. If the index is not
available it computes embeddings for the documents and saves the index.
"""
global _vectorstore
if _vectorstore is not None:
return _vectorstore
try:
_vectorstore = FAISS.load_local(
str(_VECTORSTORE_DIR),
_embeddings,
index_name=_INDEX_NAME,
allow_dangerous_deserialization=True,
)
return _vectorstore
except (FileNotFoundError, OSError, ValueError, RuntimeError) as error:
logger.info(
"FAISS index missing or unreadable; rebuilding index. %s",
error,
)
_vectorstore = FAISS.from_documents(documents, _embeddings) # will call embed_documents
_vectorstore.save_local(str(_VECTORSTORE_DIR), index_name=_INDEX_NAME)
return _vectorstore
|