feat: Get answers using preferred number of chunks

This commit is contained in:
impulsivus 2023-05-24 21:12:42 +03:00
parent 573c4363c4
commit cf709a6b7a
No known key found for this signature in database
GPG key ID: B0D9B2893A7ADC1A
3 changed files with 5 additions and 2 deletions

View file

@ -16,6 +16,7 @@ persist_directory = os.environ.get('PERSIST_DIRECTORY')
model_type = os.environ.get('MODEL_TYPE')
model_path = os.environ.get('MODEL_PATH')
model_n_ctx = os.environ.get('MODEL_N_CTX')
target_source_chunks = int(os.environ.get('TARGET_SOURCE_CHUNKS',4))
from constants import CHROMA_SETTINGS
@ -24,7 +25,7 @@ def main():
args = parse_arguments()
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
retriever = db.as_retriever()
retriever = db.as_retriever(search_kwargs={"k": target_source_chunks})
# activate/deactivate the streaming StdOut callback for LLMs
callbacks = [] if args.mute_stream else [StreamingStdOutCallbackHandler()]
# Prepare the LLM