Update output_gen.py

Irash-Perera · Oct 17, 2024 · 040d4db · 040d4db
1 parent e388f04
commit 040d4db
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/output_gen.py b/output_gen.py
@@ -78,7 +78,7 @@ def db_search(question, llm, embeddings, persist_directory):
     answer:
     """
 
-    retriever = vectorstore.as_retriever()
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
     prompt = PromptTemplate.from_template(template)
     combined_doc_chain = create_stuff_documents_chain(llm, prompt)
     retrieval_chain = create_retrieval_chain(retriever, combined_doc_chain)