Update app/load_model.py

app/load_model.py

@@ -114,9 +114,9 @@ def load_vectorstore(model_name, collection, metadata):
         )
         return vectorstore
 
-def create_chain(_llm, collection, model_name, metadata):
+def create_chain(_llm, collection, model_name, metadata,k=4):
     vectorstore = load_vectorstore(model_name, collection, metadata=metadata)
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
+    retriever = vectorstore.as_retriever(search_kwargs={"k": k})
     chain = RetrievalQA.from_chain_type(llm=_llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
     return chain
 # %%