Update multiple retriever related stuff

Author: JKusio

docker-compose.yml

@@ -12,6 +12,7 @@ services:
     environment:
       - discovery.type=single-node
       - xpack.security.enabled=false
+      - ES_JAVA_OPTS=-Xms2g -Xmx2g
     ports:
       - 9200:9200
     volumes:
@@ -31,14 +32,18 @@ services:
       - kibana_data:/usr/share/kibana/data
     depends_on:
       - elasticsearch
-  redis:
-    image: redis:latest
+  mongo:
+    image: mongo:latest
     ports:
-      - 6379:6379
+      - 27017:27017
     volumes:
-      - redis_data:/data  
+      - mongo_data:/data/db
+    deploy:
+      resources:
+        limits:
+          memory: 2g
 volumes:
   qdrant_data:
   elasticsearch_data:
   kibana_data:
-  redis_data:
+  mongo_data:

src/cache/cache.py

@@ -1,14 +1,26 @@
-import redis
+from pymongo import MongoClient
 
 SIX_MONTHS = 60 * 60 * 24 * 30 * 6
 
 
 class Cache:
     def __init__(self):
-        self.redis = redis.Redis(host="localhost", port=6379, db=0)
+        mongo_client = MongoClient("mongodb://localhost:27017/")
+        db = mongo_client["polish-nl-qa"]
+        self.key_value_collection = db["key_value"]
 
     def get(self, key):
-        return self.redis.get(key)
+        maybe_cached_value = self.key_value_collection.find_one({"key": key})
+
+        if maybe_cached_value is None:
+            return None
+
+        return maybe_cached_value["value"]
 
     def set(self, key, value):
-        return self.redis.set(key, value, SIX_MONTHS)
+        self.key_value_collection.delete_many({"key": key})
+        item = {"key": key, "value": value}
+        self.key_value_collection.insert_one(item)
+
+    def unset(self, key):
+        return self.key_value_collection.delete_one({"key": key})

src/clear_cache.py

@@ -3,13 +3,7 @@
 r = redis.StrictRedis(host="localhost", port=6379, db=0)
 
 # comment out the keys you want to keep
-prefixes = [
-    "count:*",
-    "vectorizer:*",
-    "prompt:*",
-    "reranker:*",
-    "query:*",
-]
+prefixes = ["count:*", "vectorizer:*", "prompt:*", "reranker:*", "query:*", "score:*"]
 
 for prefix in prefixes:
     print(f"Clearing keys with prefix: {prefix}")

src/common/names.py

@@ -10,6 +10,22 @@
     "BAAI/bge-m3",
 ]
 
+QUERY_PREFIX_MAP = {
+    "sdadas/mmlw-retrieval-roberta-large": "zapytanie: ",
+    "ipipan/silver-retriever-base-v1": "Pytanie: ",
+    "intfloat/multilingual-e5-large": "query: ",
+    "sdadas/mmlw-roberta-large": "zapytanie: ",
+    "BAAI/bge-m3": "",
+}
+
+PASSAGE_PREFIX_MAP = {
+    "sdadas/mmlw-retrieval-roberta-large": "",
+    "ipipan/silver-retriever-base-v1": "",
+    "intfloat/multilingual-e5-large": "passage: ",
+    "sdadas/mmlw-roberta-large": "",
+    "BAAI/bge-m3": "",
+}
+
 DISTANCES = [Distance.COSINE, Distance.EUCLID]
 
 INDEX_NAMES = [
@@ -22,4 +38,10 @@
     "morfologik_stopwords_index",
 ]
 
+RERANKER_MODEL_NAMES = [
+    "sdadas/polish-reranker-large-ranknet",
+    "BAAI/bge-reranker-v2-gemma",
+    "unicamp-dl/mt5-13b-mmarco-100k",
+]
+
 CHUNK_SIZES = [(500, 100), (1000, 200), (2000, 500), (100000, 0)]

src/common/qdrant_data_importer.py

@@ -1,5 +1,5 @@
-from typing import List
 from common.passage_factory import PassageFactory
+from common.utils import get_query_with_prefix
 from repository.qdrant_repository import QdrantRepository
 from vectorizer.vectorizer import Vectorizer
 
@@ -10,17 +10,24 @@ def __init__(
         repository: QdrantRepository,
         passage_factory: PassageFactory,
         vectorizer: Vectorizer,
+        prefix: str = "",
     ):
         self.repository = repository
         self.passage_factory = passage_factory
         self.vectorizer = vectorizer
+        self.prefix = prefix
 
     def import_data(self):
         passages = self.passage_factory.get_passages()
 
         for i in range(0, len(passages), 10):
             passages_and_vectors = [
-                (passage, self.vectorizer.get_vector(passage.context))
+                (
+                    passage,
+                    self.vectorizer.get_vector(
+                        get_query_with_prefix(passage.context, self.prefix)
+                    ),
+                )
                 for passage in passages[i : i + 10]
             ]
 

src/common/utils.py

@@ -32,8 +32,10 @@ def get_vectorizer_hash(model: str, prompt: str):
     return "vectorizer:" + hashed
 
 
-def get_prompt_hash(model: str, prompt: str):
-    hashed = hashlib.sha256((model + prompt).encode()).hexdigest()
+def get_prompt_hash(model: str, dataset_key: str, prompt: str, distance: str):
+    hashed = hashlib.sha256(
+        (model + dataset_key + prompt + distance).encode()
+    ).hexdigest()
     return "prompt:" + hashed
 
 
@@ -42,9 +44,11 @@ def get_es_query_hash(index_name: str, dataset_key: str, query: str):
     return "query:" + hashed
 
 
-def get_reranker_hash(model: str, query: str, passage_ids: list, count: int):
+def get_reranker_hash(
+    model: str, query: str, passage_ids: list, dataset_key: str, count: int
+):
     hashed = hashlib.sha256(
-        (model + query + str(passage_ids) + str(count)).encode()
+        (model + query + str(passage_ids) + dataset_key + str(count)).encode()
     ).hexdigest()
     return "reranker:" + hashed
 
@@ -77,14 +81,13 @@ def get_all_qdrant_model_combinations():
         for split, _ in CHUNK_SIZES
     ]
 
-    qdrant_collection_names = [
-        get_qdrant_collection_name(model, distance)
+    return [
+        (model, distance, dataset_key)
+        for dataset_key in dataset_keys
         for model in MODEL_NAMES
         for distance in DISTANCES
     ]
 
-    return [
-        (collection_name, dataset_key)
-        for collection_name in qdrant_collection_names
-        for dataset_key in dataset_keys
-    ]
+
+def get_query_with_prefix(query: str, prefix: str):
+    return f"{prefix}{query}"

src/load_qdrant_data.py

@@ -6,6 +6,8 @@
     DATASET_NAMES,
     DISTANCES,
     MODEL_NAMES,
+    PASSAGE_PREFIX_MAP,
+    QUERY_PREFIX_MAP,
 )
 from common.passage_factory import PassageFactory
 from common.qdrant_data_importer import QdrantDataImporter
@@ -22,7 +24,7 @@ def main():
     client = QdrantClient(host="localhost", port=6333)
     cache = Cache()
 
-    for model_name in MODEL_NAMES:
+    for model_name in ["intfloat/multilingual-e5-large"]:
         vectorizer = HFVectorizer(model_name, cache)
         for distance in DISTANCES:
             insert_passage_data(client, model_name, distance, cache, vectorizer)
@@ -56,16 +58,23 @@ def insert_passage_data(
                 chunk_size, chunk_overlap, dataset_name
             )
 
+            passage_prefix = PASSAGE_PREFIX_MAP[model_name]
+            query_prefix = QUERY_PREFIX_MAP[model_name]
+
             repository = QdrantRepository(
                 client,
                 collection_name,
                 model_name,
                 VectorParams(size=MODEL_DIMENSIONS_MAP[model_name], distance=distance),
                 vectorizer,
                 cache,
+                passage_prefix,
+                query_prefix,
             )
 
-            data_importer = QdrantDataImporter(repository, passage_factory, vectorizer)
+            data_importer = QdrantDataImporter(
+                repository, passage_factory, vectorizer, passage_prefix
+            )
 
             data_importer.import_data()
 

src/repository/es_repository.py

@@ -86,8 +86,8 @@ def count_relevant_documents(self, passage_id: str, dataset_key: str) -> int:
 
         cached_value = self.cache.get(hash_key)
 
-        # if cached_value:
-        #     return int(cached_value)
+        if cached_value:
+            return int(cached_value)
 
         body = {
             "query": {

src/repository/qdrant_repository.py

@@ -6,6 +6,7 @@
 from common.utils import (
     get_prompt_hash,
     get_qdrant_collection_name,
+    get_query_with_prefix,
     get_relevant_document_count_hash,
 )
 from repository.repository import Repository
@@ -27,12 +28,21 @@ def __init__(
         vectors_config: VectorParams,
         vectorizer: Vectorizer,
         cache: Cache,
+        passage_prefix: str = "",
+        query_prefix: str = "",
     ):
         self.qdrant = client
         self.collection_name = collection_name
         self.model_name = model_name
         self.vectorizer = vectorizer
         self.cache = cache
+        self.passage_prefix = passage_prefix
+        self.query_prefix = query_prefix
+        self.distance = (
+            Distance.COSINE
+            if Distance.COSINE.lower() in collection_name.lower()
+            else Distance.EUCLID
+        )
 
         collections = self.qdrant.get_collections()
         if collection_name not in [
@@ -53,7 +63,9 @@ def insert_one(self, passage: Passage):
             points=[
                 PointStruct(
                     id=str(uuid.uuid4()),
-                    vector=self.vectorizer.get_vector(passage.text),
+                    vector=self.vectorizer.get_vector(
+                        get_query_with_prefix(passage.context, self.passage_prefix)
+                    ),
                     payload=passage.dict(),
                 )
             ],
@@ -76,7 +88,9 @@ def insert_many(self, passages: List[Passage]):
         points = [
             PointStruct(
                 id=str(uuid.uuid4()),
-                vector=self.vectorizer.get_vector(passage.text),
+                vector=self.vectorizer.get_vector(
+                    get_query_with_prefix(passage.context, self.passage_prefix)
+                ),
                 payload=passage.dict(),
             )
             for passage in passages
@@ -86,7 +100,7 @@ def insert_many(self, passages: List[Passage]):
             collection_name=self.collection_name, wait=True, points=points
         )
 
-    def insert_many_with_vectors(self, passages: List[Tuple]):
+    def insert_many_with_vectors(self, passages: List[Tuple[Passage, List[float]]]):
         points = [
             PointStruct(
                 id=str(uuid.uuid4()),
@@ -101,7 +115,10 @@ def insert_many_with_vectors(self, passages: List[Tuple]):
         )
 
     def find(self, query: str, dataset_key: str) -> Result:
-        hash_key = get_prompt_hash(self.model_name, query)
+        full_query = get_query_with_prefix(query, self.query_prefix)
+        hash_key = get_prompt_hash(
+            self.model_name, dataset_key, full_query, self.distance
+        )
 
         cached_value = self.cache.get(hash_key)
 
@@ -110,7 +127,7 @@ def find(self, query: str, dataset_key: str) -> Result:
             passages = [(Passage.from_dict(d["passage"]), d["score"]) for d in dicts]
             return Result(query, passages)
 
-        vector = self.vectorizer.get_vector(query)
+        vector = self.vectorizer.get_vector(full_query)
 
         data = self.qdrant.search(
             collection_name=self.collection_name,
@@ -164,6 +181,8 @@ def get_repository(
         model_name: str,
         distance: Distance,
         cache: Cache,
+        passage_prefix: str = "",
+        query_prefix: str = "",
     ):
         collection_name = get_qdrant_collection_name(model_name, distance)
         vectorizer = HFVectorizer(model_name, cache)
@@ -175,6 +194,8 @@ def get_repository(
             VectorParams(size=MODEL_DIMENSIONS_MAP[model_name], distance=distance),
             vectorizer,
             cache,
+            passage_prefix,
+            query_prefix,
         )
 
     def count_relevant_documents(self, passage_id, dataset_key) -> int:

src/rerankers/hf_reranker.py

@@ -1,7 +1,9 @@
 import json
 from typing import List
+
 from cache.cache import Cache
 from common.passage import Passage
+from common.result import Result
 from common.utils import get_reranker_hash
 from rerankers.reranker import Reranker
 from sentence_transformers import CrossEncoder
@@ -15,37 +17,49 @@ def __init__(self, model_name: str, cache: Cache):
 
         print(f"Vectorizer with model {model_name} initialized")
 
-    def get_relevant_passages(
-        self, query: str, passages: List[Passage], count: int
-    ) -> List[Passage]:
-        passages_ids = list(map(lambda passage: passage.id, passages))
+    def rerank(self, result: Result, count: int, dataset_key: str) -> Result:
+        if (len(result.passages)) == 0:
+            return result
+
+        passages_ids = [passage[0].id for passage in result.passages]
         sorted_passages_ids = sorted(passages_ids)
 
-        reranker_hash = get_reranker_hash(
-            self.model_name, query, sorted_passages_ids, count
+        hash_key = get_reranker_hash(
+            self.model_name, result.query, sorted_passages_ids, dataset_key, count
         )
 
-        maybe_cached_result = self.cache.get(reranker_hash)
+        maybe_cached_result = self.cache.get(hash_key)
 
         if maybe_cached_result:
-            json_result = json.loads(maybe_cached_result)
-            return list(map(lambda passage: Passage.from_dict(passage), json_result))
+            dicts = json.loads(maybe_cached_result)
+            passages = [(Passage.from_dict(d["passage"]), d["score"]) for d in dicts]
+            return Result(result.query, passages)
 
-        pairs = [
-            [query, passage]
-            for passage in list(map(lambda passage: passage.context, passages))
-        ]
+        pairs = [[result.query, passage[0].context] for passage in result.passages]
 
         results = self.model.predict(pairs)
         results_list = results.tolist()
 
-        scored_passages = list(zip(results_list, passages))
+        scored_passages = list(zip(results_list, result.passages))
         sorted_passages = sorted(scored_passages, key=lambda x: x[0], reverse=True)
         top_n_passages = sorted_passages[:count]
-        top_n_passages = [passage for _, passage in top_n_passages]
+        top_n_passages = [(passage, score) for score, (passage, _) in top_n_passages]
+
+        max_score = top_n_passages[0][1]
+        min_score = top_n_passages[-1][1]
 
-        top_n_passages_dict = list(map(lambda passage: passage.dict(), top_n_passages))
+        score_diff = max_score - min_score
 
-        self.cache.set(reranker_hash, json.dumps(top_n_passages_dict))
+        normalized_passages = [
+            (p, 1 if score_diff == 0 else (s - min_score) / score_diff)
+            for (p, s) in top_n_passages
+        ]
+
+        reranked_result = Result(result.query, normalized_passages)
+
+        result_json = json.dumps(
+            [{"passage": p.dict(), "score": s} for (p, s) in normalized_passages]
+        )
+        self.cache.set(hash_key, result_json)
 
-        return top_n_passages
+        return reranked_result

src/rerankers/reranker.py

@@ -1,11 +1,8 @@
 from abc import ABC, abstractmethod
-from typing import List
-from common.passage import Passage
+from common.result import Result
 
 
 class Reranker(ABC):
     @abstractmethod
-    def get_relevant_passages(
-        self, query: str, passages: List[Passage], count: int
-    ) -> List[Passage]:
+    def rerank(self, result: Result, count: int) -> Result:
         pass

src/retrievers/es_retriever.py

@@ -1,15 +1,22 @@
 from typing import List
 from common.result import Result
 from repository.es_repository import ESRepository
+from rerankers.hf_reranker import HFReranker
 from retrievers.retriever import Retriever
 
 
 class ESRetriever(Retriever):
-    def __init__(self, repository: ESRepository, dataset_key: str):
+    def __init__(
+        self, repository: ESRepository, dataset_key: str, reranker: HFReranker = None
+    ):
         self.repository = repository
         self.dataset_key = dataset_key
+        self.reranker = reranker
 
     def get_relevant_passages(self, query: str) -> Result:
         result = self.repository.find(query, self.dataset_key)
 
+        if self.reranker:
+            result = self.reranker.rerank(result, 10, self.dataset_key)
+
         return result

src/retrievers/hybrid_retriever.py

@@ -3,6 +3,7 @@
 from common.result import Result
 from repository.es_repository import ESRepository
 from repository.qdrant_repository import QdrantRepository
+from rerankers.hf_reranker import HFReranker
 from retrievers.retriever import Retriever
 
 
@@ -13,11 +14,13 @@ def __init__(
         qdrant_repository: QdrantRepository,
         dataset_key: str,
         alpha: float = 0.5,  # weight for ES
+        reranker: HFReranker = None,
     ):
         self.es_repository = es_repository
         self.qdrant_repository = qdrant_repository
         self.dataset_key = dataset_key
         self.alpha = alpha
+        self.reranker = reranker
 
     def get_relevant_passages(self, query: str) -> List[str]:
         es_result = self.es_repository.find(query, self.dataset_key)
@@ -39,4 +42,9 @@ def get_relevant_passages(self, query: str) -> List[str]:
         final_results = list(combined_scores.items())
         final_results.sort(key=lambda x: x[1], reverse=True)
 
-        return Result(query, final_results[:10])
+        result = Result(query, final_results[:10])
+
+        if self.reranker:
+            result = self.reranker.rerank(result, 10, self.dataset_key)
+
+        return result

src/retrievers/qdrant_retriever.py

@@ -1,15 +1,25 @@
 from typing import List
 from common.result import Result
 from repository.qdrant_repository import QdrantRepository
+from rerankers.hf_reranker import HFReranker
 from retrievers.retriever import Retriever
 
 
 class QdrantRetriever(Retriever):
-    def __init__(self, repository: QdrantRepository, dataset_key: str):
+    def __init__(
+        self,
+        repository: QdrantRepository,
+        dataset_key: str,
+        reranker: HFReranker = None,
+    ):
         self.repository = repository
         self.dataset_key = dataset_key
+        self.reranker = reranker
 
     def get_relevant_passages(self, query: str) -> Result:
         result = self.repository.find(query, self.dataset_key)
 
+        if self.reranker:
+            result = self.reranker.rerank(result, 10, self.dataset_key)
+
         return result