Spaces:
Sleeping
Sleeping
def create_retriever_from_chroma(vectorstore_path="./docs/chroma/", search_type='mmr', k=7, chunk_size=300, chunk_overlap=30,lambda_mult= 0.7): | |
model_name = "Alibaba-NLP/gte-large-en-v1.5" | |
model_kwargs = {'device': 'cpu', | |
"trust_remote_code" : 'False'} | |
encode_kwargs = {'normalize_embeddings': True} | |
embeddings = HuggingFaceEmbeddings( | |
model_name=model_name, | |
model_kwargs=model_kwargs, | |
encode_kwargs=encode_kwargs | |
) | |
if os.path.exists(vectorstore_path) and os.listdir(vectorstore_path): | |
vectorstore = Chroma(persist_directory=vectorstore_path,embedding_function=embeddings) | |
else: | |
st.write("Vector store doesnt exist and will be created now") | |
loader = DirectoryLoader('./data/', glob="./*.txt", loader_cls=TextLoader) | |
docs = loader.load() | |
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder( | |
chunk_size=chunk_size, chunk_overlap=chunk_overlap, | |
separators=["\n\n \n\n","\n\n\n", "\n\n", r"In \[[0-9]+\]", r"\n+", r"\s+"], | |
is_separator_regex = True | |
) | |
split_docs = text_splitter.split_documents(docs) | |
vectorstore = Chroma.from_documents( | |
documents=split_docs, embedding=embeddings, persist_directory=vectorstore_path | |
) | |
retriever=vectorstore.as_retriever(search_type = search_type, search_kwargs={"k": k}) | |
return retriever | |