Return count of relevant docs

Author: JKusio

src/common/names.py

@@ -22,13 +22,4 @@
     "morfologik_stopwords_index",
 ]
 
-SEMANTIC_TYPES = ["interquartile", "standard_deviation", "percentile"]
-
 CHUNK_SIZES = [(500, 100), (1000, 200), (2000, 500), (100000, 0)]
-
-CHARACTER_SPLITTING_FUNCTION = [
-    "character-500",
-    "character-1000",
-    "character-2000",
-    "character-100000",
-]

src/common/utils.py

@@ -1,11 +1,5 @@
 import uuid
-from common.names import (
-    CHUNK_SIZES,
-    DATASET_NAMES,
-    DISTANCES,
-    MODEL_NAMES,
-    SEMANTIC_TYPES,
-)
+from common.names import CHUNK_SIZES, DATASET_NAMES, DISTANCES, INDEX_NAMES, MODEL_NAMES
 from common.passage import Passage
 import hashlib
 
@@ -50,25 +44,42 @@ def get_reranker_hash(model: str, query: str, passage_ids: list, count: int):
     return "reranker:" + hashed
 
 
-def get_all_qdrant_collection_names():
-    names = []
-    for dataset_name in DATASET_NAMES:
-        for model_name in MODEL_NAMES:
-            for distance in DISTANCES:
-                for chunk_size, _ in CHUNK_SIZES:
-                    name = get_qdrant_collection_name(
-                        dataset_name, model_name, "character", chunk_size, distance
-                    )
-                    names.append(name)
-
-                for semantic_type in SEMANTIC_TYPES:
-                    name = get_qdrant_collection_name(
-                        dataset_name, model_name, semantic_type, 1.5, distance
-                    )
-                    names.append(name)
-
-    return names
+def get_relevant_document_count_hash(id: str, dataset_key: str):
+    hashed = hashlib.sha256((id + dataset_key).encode()).hexdigest()
+    return "count:" + hashed
 
 
 def get_dataset_key(dataset_name: str, split: str):
     return replace_slash_with_dash(f"{dataset_name}-{split}")
+
+
+def get_all_es_index_combinations():
+    dataset_keys = [
+        get_dataset_key(dataset_name, split)
+        for dataset_name in DATASET_NAMES
+        for split, _ in CHUNK_SIZES
+    ]
+
+    return [
+        (index, dataset_key) for index in INDEX_NAMES for dataset_key in dataset_keys
+    ]
+
+
+def get_all_qdrant_model_combinations():
+    dataset_keys = [
+        get_dataset_key(dataset_name, split)
+        for dataset_name in DATASET_NAMES
+        for split, _ in CHUNK_SIZES
+    ]
+
+    qdrant_collection_names = [
+        get_qdrant_collection_name(model, distance)
+        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
+    ]

src/evaluation/retriever_evaluator.py

@@ -7,10 +7,9 @@
 
 class RetrieverEvaluator:
     # Calculate NDCG for top 10 results
-    def calculate_ndcg(self, result: Result, correct_passage_title: str) -> float:
+    def calculate_ndcg(self, result: Result, correct_passage_id: str) -> float:
         relevances = [
-            1 if passage.title == correct_passage_title else 0
-            for passage in result.passages
+            1 if passage.id == correct_passage_id else 0 for passage in result.passages
         ]
 
         sorted_relevances = sorted(relevances, reverse=True)
@@ -23,22 +22,22 @@ def calculate_ndcg(self, result: Result, correct_passage_title: str) -> float:
         return dcg / idcg if idcg != 0 else 0
 
     # Calculate MRR for top 10 results
-    def calculate_mrr(self, result: Result, correct_passage_title: str) -> float:
+    def calculate_mrr(self, result: Result, correct_passage_id: str) -> float:
         for i, passage in enumerate(result.passages):
-            if passage.title == correct_passage_title:
+            if passage.id == correct_passage_id:
                 return 1 / (i + 1)
         return 0
 
     # Calculate recall for top 10 results
     def calculate_recall(
-        self, result: Result, correct_passage_title: str, relevant_documents_count: int
+        self, result: Result, correct_passage_id: str, relevant_documents_count: int
     ) -> float:
         relevant_documents = sum(
-            1 for passage in result.passages if passage.title == correct_passage_title
+            1 for passage in result.passages if passage.id == correct_passage_id
         )
 
         return relevant_documents / relevant_documents_count
 
     # Calculate accuracy for top 1 result
-    def calculate_accuracy(self, result: Result, correct_passage_title: str) -> float:
-        return 1 if result.passages[0].title == correct_passage_title else 0
+    def calculate_accuracy(self, result: Result, correct_passage_id: str) -> float:
+        return 1 if result.passages[0].id == correct_passage_id else 0

src/main.py

@@ -1,29 +1,8 @@
-from ast import Dict
-from tkinter import Place
-from xml.etree.ElementInclude import include
-from elasticsearch import Elasticsearch
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from common.names import DATASET_NAMES, INDEX_NAMES
-from common.passage import Passage
+from common.names import DATASET_NAMES
 from common.passage_factory import PassageFactory
-from common.utils import replace_slash_with_dash
 from dataset.poquad_dataset_getter import PoquadDatasetGetter
-from repository.es_repository import ESRepository
-from repository.qdrant_repository import QdrantRepository
-from qdrant_client import QdrantClient
-from qdrant_client.models import Distance, VectorParams
-from vectorizer.hf_vectorizer import HFVectorizer
-from langchain_experimental.text_splitter import SemanticChunker
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from dataset.polqa_dataset_getter import PolqaDatasetGetter
-from elasticsearch import Elasticsearch
-from qdrant_client import QdrantClient
-from cache.cache import Cache
-from common.models_dimensions import MODEL_DIMENSIONS_MAP
-from common.names import DISTANCES, MODEL_NAMES
-from common.utils import get_all_qdrant_collection_names
-from repository.qdrant_repository import QdrantRepository
-from vectorizer.hf_vectorizer import HFVectorizer
 
 
 def get_passage_factory(

src/repository/es_repository.py

@@ -3,7 +3,10 @@
 from cache.cache import Cache
 from common.passage import Passage
 from common.result import Result
-from common.utils import get_es_query_hash
+from common.utils import (
+    get_es_query_hash,
+    get_relevant_document_count_hash,
+)
 from repository.repository import Repository
 import json
 
@@ -74,3 +77,29 @@ def find(self, query: str, dataset_key: str) -> Result:
     def delete(self, query: str):
         body = {"query": {"match": {"text": query}}}
         return self.client.delete_by_query(index=self.index_name, body=body)
+
+    def count_relevant_documents(self, id, dataset_key) -> int:
+        hash_key = get_relevant_document_count_hash(id, dataset_key)
+
+        cached_value = self.cache.get(hash_key)
+
+        if cached_value:
+            return int(cached_value)
+
+        if cached_value:
+            return int(cached_value)
+
+        body = {
+            "query": {
+                "bool": {
+                    "must": [
+                        {"match": {"id": id}},
+                        {"match": {"dataset_key": dataset_key}},
+                    ]
+                }
+            },
+        }
+
+        response = self.client.count(index=self.index_name, body=body)
+
+        return response["count"]

src/repository/qdrant_repository.py

@@ -3,7 +3,11 @@
 from common.models_dimensions import MODEL_DIMENSIONS_MAP
 from common.passage import Passage
 from common.result import Result
-from common.utils import get_prompt_hash, get_qdrant_collection_name
+from common.utils import (
+    get_prompt_hash,
+    get_qdrant_collection_name,
+    get_relevant_document_count_hash,
+)
 from repository.repository import Repository
 from qdrant_client import QdrantClient, models
 from qdrant_client.models import VectorParams, PointStruct, Distance
@@ -169,3 +173,32 @@ def get_repository(
             vectorizer,
             cache,
         )
+
+    def count_relevant_documents(self, id, dataset_key) -> int:
+        hash_key = get_relevant_document_count_hash(id, dataset_key)
+
+        cached_value = self.cache.get(hash_key)
+
+        if cached_value:
+            return int(cached_value)
+
+        result = self.qdrant.count(
+            collection_name=self.collection_name,
+            count_filter=models.Filter(
+                must=[
+                    models.FieldCondition(
+                        key="id",
+                        match=models.MatchValue(value=int(id)),
+                    ),
+                    models.FieldCondition(
+                        key="dataset_key",
+                        match=models.MatchValue(value=dataset_key),
+                    ),
+                ]
+            ),
+            exact=True,
+        )
+
+        self.cache.set(hash_key, str(result.count))
+
+        return result.count

src/repository/repository.py

@@ -20,3 +20,7 @@ def find(self, query, dataset_key) -> Result:
     @abstractmethod
     def delete(self, query):
         pass
+
+    @abstractmethod
+    def count_relevant_documents(self, id, dataset_key) -> int:
+        pass