Spaces:
Build error
Build error
File size: 1,157 Bytes
e46711a |
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 |
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
import config
def get_embedding_model():
"""
Initialize the embedding model.
"""
return HuggingFaceEmbeddings(
model_name=config.EMBEDDING_MODEL_NAME,
model_kwargs={'device': 'cpu'}
)
def create_vector_store(chunks, embedding_model):
"""
Create and persist a Chroma vector store from document chunks.
"""
vectorstore = Chroma.from_documents(
documents=chunks,
embedding=embedding_model,
persist_directory=config.CHROMA_DB_DIR
)
return vectorstore
def get_vector_store(embedding_model):
"""
Load existing vector store.
"""
# Simply initializing with persist_directory attempts to load it
return Chroma(
persist_directory=config.CHROMA_DB_DIR,
embedding_function=embedding_model
)
def get_retriever(vectorstore):
"""
Get a retriever from the vector store.
"""
return vectorstore.as_retriever(
search_type="similarity",
search_kwargs={"k": config.RETRIEVER_K}
)
|