Update llm/utils.py

llm/utils.py

@@ -3,7 +3,7 @@ from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 import gradio as gr
 import os
-#from CustomRetriever import CustomRetriever
+from llm.CustomRetriever import CustomRetriever
 
 from langchain.schema.retriever import BaseRetriever
 from langchain_core.documents import Document
@@ -14,41 +14,14 @@ from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain_core.documents import Document
 from langchain_core.runnables import chain
 
-
 API_TOKEN=os.getenv("TOKEN")
 
 
-#Because of bugs in pydantic it is not possible to take it out retr_func and RetrieverWithScores into a separate neat class. 
-#It is necessary to use dirty implementation through global variables.
-VDB=None
-THOLD=0.7
-
-@chain
-def retr_func(query: str)-> List[Document]:   
-    
-    docs, scores = zip(*VDB.similarity_search_with_relevance_scores(query))#similarity_search_with_score(query))
-    result=[]
-    for doc, score in zip(docs, scores):
-        if score>THOLD: 
-            doc.metadata["score"] = score
-            result.append(doc)
-        if len(result)==0:
-            result.append(Document(metadata={}, page_content='No data'))
-
-    return result
-    
-
-class RetrieverWithScores(BaseRetriever):                
-    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun)-> List[Document]:
-        return retr_func.invoke(query) 
-    
-
-
 # Initialize langchain LLM chain
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vdb, 
                         thold=0.8, progress=gr.Progress()):
-    global VDB
-    global THOLD
+    #global VDB
+    #global THOLD
     
     llm = HuggingFaceEndpoint(
             huggingfacehub_api_token = API_TOKEN,
@@ -64,13 +37,11 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vdb,
         return_messages=True
     )
     
-    VDB=vdb
-    THOLD=thold
-    #retr=CustomRetriever(vdb, thold=thold)    
-    #retriever=retr.retriever 
+    #VDB=vdb
+    #THOLD=thold    
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
-        retriever=RetrieverWithScores(),#retriever,
+        retriever=CustomRetriever(vectorstore=vdb, thold=thold),#RetrieverWithScores(),
         chain_type="stuff",
         memory=memory,
         return_source_documents=True,