Update llm/utils.py

llm/utils.py

@@ -1,74 +1,74 @@
-from langchain_community.llms import HuggingFaceEndpoint
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
-import gradio as gr
-import os
-from CustomRetriever import CustomRetriever
-
-
-API_TOKEN=os.getenv("TOKEN")
-
-# Initialize langchain LLM chain
-def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vdb, 
-                        thold=0.8, progress=gr.Progress()):
-    llm = HuggingFaceEndpoint(
-            huggingfacehub_api_token = API_TOKEN,
-            repo_id=llm_model,
-            temperature = temperature,
-            max_new_tokens = max_tokens,
-            top_k = top_k,
-        )
-
-    memory = ConversationBufferMemory(
-        memory_key="chat_history",
-        output_key='answer',
-        return_messages=True
-    )
-
-    retr=CustomRetriever(vdb, thold=thold)
-    retriever=retr.retriever #vector_db.as_retriever()
-    qa_chain = ConversationalRetrievalChain.from_llm(
-        llm,
-        retriever=retriever,
-        chain_type="stuff",
-        memory=memory,
-        return_source_documents=True,
-        verbose=False,
-    )
-    return qa_chain
-
-
-
-# Initialize LLM
-def initialize_LLM(llm_temperature, max_tokens, top_k, vector_db, thold, progress=gr.Progress()):
-    # print("llm_option",llm_option)
-    llm_name = "mistralai/Mistral-7B-Instruct-v0.2" #list_llm[llm_option]
-    #print("llm_name: ",llm_name)
-    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, thold)
-    return qa_chain #, "QA chain initialized. Chatbot is ready!"
-
-
-
-def format_chat_history(chat_history):#message, chat_history): #no need message
-    formatted_chat_history = []
-    for user_message, bot_message in chat_history:
-        formatted_chat_history.append(f"User: {user_message}")
-        formatted_chat_history.append(f"Assistant: {bot_message}")
-    return formatted_chat_history
-
-
-
-def postprocess(response):
-    try:
-        result=response["answer"]
-        for doc in response['source_documents']:
-            file_doc="\n\nFile: " + doc.metadata["source"]
-            page="\nPage: " + str(doc.metadata["page"])
-            content="\nFragment: " + doc.page_content.strip()
-            result+=file_doc+page+content
-        return result
-    except:
-        return response["answer"]
-
-
-
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+import gradio as gr
+import os
+from CustomRetriever import RetrieverWithScores #CustomRetriever
+
+
+API_TOKEN=os.getenv("TOKEN")
+
+# Initialize langchain LLM chain
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vdb, 
+                        thold=0.8, progress=gr.Progress()):
+    llm = HuggingFaceEndpoint(
+            huggingfacehub_api_token = API_TOKEN,
+            repo_id=llm_model,
+            temperature = temperature,
+            max_new_tokens = max_tokens,
+            top_k = top_k,
+        )
+
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
+    )
+
+    #retr=CustomRetriever(vdb, thold=thold)
+    retriever=RetrieverWithScores(vdb, thold=thold)  #retr.retriever #vector_db.as_retriever()
+    qa_chain = ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=retriever,
+        chain_type="stuff",
+        memory=memory,
+        return_source_documents=True,
+        verbose=False,
+    )
+    return qa_chain
+
+
+
+# Initialize LLM
+def initialize_LLM(llm_temperature, max_tokens, top_k, vector_db, thold, progress=gr.Progress()):
+    # print("llm_option",llm_option)
+    llm_name = "mistralai/Mistral-7B-Instruct-v0.2" #list_llm[llm_option]
+    #print("llm_name: ",llm_name)
+    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, thold)
+    return qa_chain #, "QA chain initialized. Chatbot is ready!"
+
+
+
+def format_chat_history(chat_history):#message, chat_history): #no need message
+    formatted_chat_history = []
+    for user_message, bot_message in chat_history:
+        formatted_chat_history.append(f"User: {user_message}")
+        formatted_chat_history.append(f"Assistant: {bot_message}")
+    return formatted_chat_history
+
+
+
+def postprocess(response):
+    try:
+        result=response["answer"]
+        for doc in response['source_documents']:
+            file_doc="\n\nFile: " + doc.metadata["source"]
+            page="\nPage: " + str(doc.metadata["page"])
+            content="\nFragment: " + doc.page_content.strip()
+            result+=file_doc+page+content
+        return result
+    except:
+        return response["answer"]
+
+
+