Upload retrievers_setup.py

rag_chain/retrievers_setup.py

@@ -3,10 +3,13 @@ from functools import cache
 
 import qdrant_client
 import torch
+from langchain.prompts import PromptTemplate
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.document_compressors import EmbeddingsFilter
+from langchain.retrievers.multi_query import MultiQueryRetriever
 from langchain_community.retrievers import QdrantSparseVectorRetriever
 from langchain_community.vectorstores import Qdrant
+from langchain_openai import ChatOpenAI
 from langchain_openai.embeddings import OpenAIEmbeddings
 from transformers import AutoModelForMaskedLM, AutoTokenizer
 
@@ -28,6 +31,7 @@ class DenseRetrieverClient:
         self.client = qdrant_client.QdrantClient(
             url=os.getenv("QDRANT_URL"),
             api_key=os.getenv("QDRANT_API_KEY"),
+            prefer_grpc=True,
         )
         self.qdrant_collection = self.load_qdrant_collection()
 
@@ -91,6 +95,7 @@ class SparseRetrieverClient:
         self.client = qdrant_client.QdrantClient(url=os.getenv(
             "QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY"))
         self.model_id = splade_model_id
+        self.tokenizer, self.model = self.set_tokenizer_config()
         self.collection_name = collection_name
         self.vector_name = vector_name
         self.k = k
@@ -120,17 +125,23 @@ class SparseRetrieverClient:
         Returns:
             tuple[list[int], list[float]]: Indices and values of the sparse vector
         """
-        tokenizer, model = self.set_tokenizer_config()
-        tokens = tokenizer(text, return_tensors="pt",
-                           max_length=512, padding="max_length", truncation=True)
-        output = model(**tokens)
+        tokens = self.tokenizer(text, return_tensors="pt",
+                                max_length=512, padding="max_length", truncation=True)
+
+        with torch.no_grad():
+            output = self.model(**tokens)
+
         logits, attention_mask = output.logits, tokens.attention_mask
-        relu_log = torch.log(1 + torch.relu(logits))
+
+        relu_log = torch.log1p(torch.relu(logits))
         weighted_log = relu_log * attention_mask.unsqueeze(-1)
+
         max_val, _ = torch.max(weighted_log, dim=1)
         vec = max_val.squeeze()
-        indices = vec.nonzero().numpy().flatten()
-        values = vec.detach().numpy()[indices]
+
+        indices = torch.nonzero(vec, as_tuple=False).squeeze().cpu().numpy()
+        values = vec[indices].cpu().numpy()
+
         return indices.tolist(), values.tolist()
 
     def get_sparse_retriever(self):
@@ -172,3 +183,34 @@ def compression_retriever_setup(base_retriever, embeddings_model: str = "text-em
     )
 
     return compression_retriever
+
+
+def multi_query_retriever_setup(retriever) -> MultiQueryRetriever:
+    """ Configure a multi-query retriever using a base retriever and the LLM.
+
+    Args:
+        retriever:
+
+    Returns:
+        retriever: MultiQueryRetriever
+    """
+
+    QUERY_PROMPT = PromptTemplate(
+        input_variables=["question"],
+        template="""
+
+        Your task is to generate 3 different versions of the provided question, incorporating the user's location preference in each version. Each version must be separated by newlines. Ensure that no part of your response is enclosed in quotation marks. Do not modify any acronyms or unfamiliar terms. Keep your responses clear, concise, and limited to these alternatives. 
+        Note: The text provided are queries to Tall Tree Health Centre's AI virtual assistant.
+
+        Question:
+        {question}
+
+        """,
+    )
+
+    llm = ChatOpenAI(model='gpt-3.5-turbo', temperature=0)
+    multi_query_retriever = MultiQueryRetriever.from_llm(
+        retriever=retriever, llm=llm, prompt=QUERY_PROMPT, include_original=True,
+    )
+
+    return multi_query_retriever