Update create_retriever.py

create_retriever.py

@@ -75,7 +75,7 @@ def get_text_splitter(splitter_type='character', chunk_size=500, chunk_overlap=3
         raise ValueError("Unsupported splitter type. Choose from 'character', 'recursive', or 'token'.")
 
 # Retriever using Chroma and HuggingFace embeddings
-def retriever_chroma(flag, model_name="BAAI/bge-large-en-v1.5", splitter_type='character', chunk_size=500, chunk_overlap=30, separator="\n", max_tokens=1000):
+def retriever_chroma(flag, model_name="sentence-transformers/all-mpnet-base-v2", splitter_type='character', chunk_size=500, chunk_overlap=30, separator="\n", max_tokens=1000):
     # Load or download the embedding model
     embeddings = get_embedding_model(model_name)
 
@@ -84,18 +84,16 @@ def retriever_chroma(flag, model_name="BAAI/bge-large-en-v1.5", splitter_type='c
         all_doc = docs_return(0)
 
         # Use the splitter parameters
-        text_splitter = get_text_splitter(splitter_type=splitter_type, chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator, max_tokens=max_tokens)
+        # text_splitter = get_text_splitter(splitter_type=splitter_type, chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator, max_tokens=max_tokens)
 
         # Split the documents using the text splitter
-        docs = text_splitter.split_documents(documents=all_doc)
+        # docs = text_splitter.split_documents(documents=all_doc)
 
         # Create a Chroma vector database
-        vectordb = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_db")
-
+        vectordb = Chroma.from_documents(all_doc, embeddings, persist_directory="./chroma_db")
+        # print(all_doc)
         # Create the retriever
-        chroma_retriever = vectordb.as_retriever(
-            search_type="mmr", search_kwargs={"k": 4, "fetch_k": 10}
-        )
+        chroma_retriever = vectordb.as_retriever()
         return chroma_retriever
     else:
         # Load a local Chroma vectorstore