Spaces:
Running
Running
File size: 1,492 Bytes
6a8cbae |
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 |
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
|