Update model.py

model.py

@@ -1,102 +1,102 @@
-#from langchain import PromptTemplate
-from langchain_core.prompts import PromptTemplate
-
-from langchain_community.embeddings import HuggingFaceBgeEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain_community.llms.ctransformers import CTransformers
-#from langchain.chains import RetrievalQA
-from langchain.chains.retrieval_qa.base import RetrievalQA
-import chainlit as cl
-from transformers import AutoModel
-
-DB_FAISS_PATH = 'vectorstores/'
-
-custom_prompt_template = '''
-use the following pieces of information to answer the user's questions/
-If you don't know the answer, please just say that don't know the answer, don't try to make uo an answer.
-
-Content : {}
-Question : {question}
-
-only return the helpful answer below and nothing else.
-'''
-
-def set_custom_prompt():
-    """
-    Prompt template for QA retrieval for vector stores
-    """
-    prompt = PromptTemplate(template = custom_prompt_template,
-                            input_variables = ['context','question'])
-    
-    return prompt
-
-def load_llm():
-    llm = CTransformers(
-        model = 'TheBloke/Llama-2-7B-Chat-GGML',
-        #model = AutoModel.from_pretrained("TheBloke/Llama-2-7B-Chat-GGML"),
-        model_type = 'llama',
-        max_new_token = 512,
-        temperature = 0.5
-    )
-    return llm
-
-def retrieval_qa_chain(llm,prompt,db):
-    qa_chain = RetrievalQA.from_chain_type(
-        llm = llm,
-        chain_type = 'stuff',
-        retriever = db.as_retriever(search_kwargs= {'k': 2}),
-        return_source_documents = True,
-        chain_type_kwargs = {'prompt': prompt}
-    )
-
-    return qa_chain
-
-def qa_bot():
-    embeddings = HuggingFaceBgeEmbeddings(model_name = 'sentence-transformers/all-MiniLM-L6-v2',
-                                          model_kwargs = {'device':'cpu'})
-    
-    db = FAISS.load_local(DB_FAISS_PATH,embeddings)
-    llm = load_llm()
-    qa_prompt = set_custom_prompt()
-    qa = retrieval_qa_chain(llm,qa_prompt, db)
-
-    return qa
-
-def final_result(query):
-    qa_result = qa_bot()
-    response = qa_result({'quert' : query})
-
-    return response
-
-
-## Chainlit
-@cl.on_chat_start
-async def start():
-    chain = qa_bot()
-    msg = cl.Message(content = 'Starting the bot...')
-    await msg.send()
-
-    msg.conteny = "Hi Welcome to the medical Bot. What is your query?"
-    await msg.update()
-    cl.user_session.set('chain', chain)
-    
-@cl.on_message
-async def main(message):
-    chain = cl.user_session.set('chain')
-    cb = cl.AsyncLangchainCallbackHandler(
-        stream_final_answer= True,
-        answer_prefix_tokens= ['FINAL','ANSWER']
-    )
-    cb.answer_reached = True
-    res = await chain.acall(message,callbacks = [cb])
-    answer = res['result']
-    sources = res['sources_documents']
-
-    if sources :
-        answer += f"\nSources :" + str(sources)
-    
-    else :
-        answer += f"\nNo Rources Found"
-    
-    await cl.Message(content=answer).send()
-
+#from langchain import PromptTemplate
+from langchain_core.prompts import PromptTemplate
+
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_community.llms.ctransformers import CTransformers
+#from langchain.chains import RetrievalQA
+from langchain.chains.retrieval_qa.base import RetrievalQA
+import chainlit as cl
+from transformers import AutoModel
+
+DB_FAISS_PATH = 'vectorstores/'
+
+custom_prompt_template = '''
+use the following pieces of information to answer the user's questions.
+If you don't know the answer, please just say that don't know the answer, don't try to make uo an answer.
+
+Context : {}
+Question : {question}
+
+only return the helpful answer below and nothing else.
+'''
+
+def set_custom_prompt():
+    """
+    Prompt template for QA retrieval for vector stores
+    """
+    prompt = PromptTemplate(template = custom_prompt_template,
+                            input_variables = ['context','question'])
+    
+    return prompt
+
+def load_llm():
+    llm = CTransformers(
+        model = 'TheBloke/Llama-2-7B-Chat-GGML',
+        #model = AutoModel.from_pretrained("TheBloke/Llama-2-7B-Chat-GGML"),
+        model_type = 'llama',
+        max_new_token = 512,
+        temperature = 0.5
+    )
+    return llm
+
+def retrieval_qa_chain(llm,prompt,db):
+    qa_chain = RetrievalQA.from_chain_type(
+        llm = llm,
+        chain_type = 'stuff',
+        retriever = db.as_retriever(search_kwargs= {'k': 2}),
+        return_source_documents = True,
+        chain_type_kwargs = {'prompt': prompt}
+    )
+
+    return qa_chain
+
+def qa_bot():
+    embeddings = HuggingFaceBgeEmbeddings(model_name = 'sentence-transformers/all-MiniLM-L6-v2',
+                                          model_kwargs = {'device':'cpu'})
+    
+    db = FAISS.load_local(DB_FAISS_PATH,embeddings)
+    llm = load_llm()
+    qa_prompt = set_custom_prompt()
+    qa = retrieval_qa_chain(llm,qa_prompt, db)
+
+    return qa
+
+def final_result(query):
+    qa_result = qa_bot()
+    response = qa_result({'query' : query})
+
+    return response
+
+
+## Chainlit
+@cl.on_chat_start
+async def start():
+    chain = qa_bot()
+    msg = cl.Message(content = 'Starting the bot...')
+    await msg.send()
+
+    msg.conteny = "Hi Welcome to the medical Bot. What is your query?"
+    await msg.update()
+    cl.user_session.set('chain', chain)
+    
+@cl.on_message
+async def main(message):
+    chain = cl.user_session.set('chain')
+    cb = cl.AsyncLangchainCallbackHandler(
+        stream_final_answer= True,
+        answer_prefix_tokens= ['FINAL','ANSWER']
+    )
+    cb.answer_reached = True
+    res = await chain.acall(message,callbacks = [cb])
+    answer = res['result']
+    sources = res['sources_documents']
+
+    if sources :
+        answer += f"\nSources :" + str(sources)
+    
+    else :
+        answer += f"\nNo Rources Found"
+    
+    await cl.Message(content=answer).send()
+