Update main.py

main.py

@@ -49,8 +49,31 @@ qa_chain =  RetrievalQA.from_llm(llm=local_llm, retriever=retriever)
 
 
 def gradinterface(query,history):
-    result = qa_chain({'query': query})
-    return result.split(' : ').strip()
+    while True:
+  query = input("\nQuery: ")
+  if query == "exit":
+      break
+  if query.strip() == "":
+      continue
+
+  template = """Use the following pieces of context to answer the question at the end.
+  If you don't know the answer, just say that you don't know, don't try to make up an answer.
+  Use three sentences maximum and keep the answer as concise as possible.
+  {context}
+  Question: {question}
+  Helpful Answer:"""
+
+  QA_CHAIN_PROMPT = PromptTemplate(
+      input_variables=["context", "question"],
+      template=template,
+  )
+  qa_chain = RetrievalQA.from_chain_type(
+        local_llm,
+        retriever=vector.as_retriever(),
+        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT},
+    )
+  result = qa_chain({"query": query})
+  return result['result'].split(':')[-1].strip()
 
 
 demo = gr.ChatInterface(fn=gradinterface, title='OUR_OWN_BOT')