debeir.core.query

  1import dataclasses
  2from typing import Dict, Optional, Union
  3
  4import loguru
  5from debeir.engines.elasticsearch.generate_script_score import generate_script
  6from debeir.core.config import GenericConfig, apply_config
  7from debeir.utils.scaler import get_z_value
  8
  9
 10@dataclasses.dataclass(init=True)
 11class Query:
 12    """
 13    A query interface class
 14    :param topics: Topics that the query will be composed of
 15    :param config: Config object that contains the settings for querying
 16    """
 17    topics: Dict[int, Dict[str, str]]
 18    config: GenericConfig
 19
 20
 21class GenericElasticsearchQuery(Query):
 22    """
 23    A generic elasticsearch query. Contains methods for NIR-style (embedding) queries and normal BM25 queries.
 24    Requires topics, configs to be included
 25    """
 26    id_mapping: str = "Id"
 27
 28    def __init__(self, topics, config, top_bm25_scores=None, mappings=None, id_mapping=None, *args, **kwargs):
 29        super().__init__(topics, config)
 30
 31        if id_mapping is None:
 32            self.id_mapping = "id"
 33
 34        if mappings is None:
 35            self.mappings = ["Text"]
 36        else:
 37            self.mappings = mappings
 38
 39        self.topics = topics
 40        self.config = config
 41        self.query_type = self.config.query_type
 42
 43        self.embed_mappings = ["Text_Embedding"]
 44
 45        self.query_funcs = {
 46            "query": self.generate_query,
 47            "embedding": self.generate_query_embedding,
 48        }
 49
 50        self.top_bm25_scores = top_bm25_scores
 51
 52    def _generate_base_query(self, topic_num):
 53        qfield = list(self.topics[topic_num].keys())[0]
 54        query = self.topics[topic_num][qfield]
 55        should = {"should": []}
 56
 57        for i, field in enumerate(self.mappings):
 58            should["should"].append(
 59                {
 60                    "match": {
 61                        f"{field}": {
 62                            "query": query,
 63                        }
 64                    }
 65                }
 66            )
 67
 68        return qfield, query, should
 69
 70    def generate_query(self, topic_num, *args, **kwargs):
 71        """
 72        Generates a simple BM25 query based off the query facets. Searches over all the document facets.
 73        :param topic_num:
 74        :param args:
 75        :param kwargs:
 76        :return:
 77        """
 78        _, _, should = self._generate_base_query(topic_num)
 79
 80        query = {
 81            "query": {
 82                "bool": should,
 83            }
 84        }
 85
 86        return query
 87
 88    def set_bm25_scores(self, scores: Dict[Union[str, int], Union[int, float]]):
 89        """
 90        Sets BM25 scores that are used for NIR-style scoring. The top BM25 score for each topic is used
 91        for log normalization.
 92
 93        Score = log(bm25)/log(z) + embed_score
 94        :param scores: Top BM25 Scores of the form {topic_num: top_bm25_score}
 95        """
 96        self.top_bm25_scores = scores
 97
 98    def has_bm25_scores(self):
 99        """
100        Checks if BM25 scores have been set
101        :return:
102        """
103        return self.top_bm25_scores is not None
104
105    @apply_config
106    def generate_query_embedding(
107            self, topic_num, encoder, *args, norm_weight=2.15, ablations=False, cosine_ceiling=Optional[float],
108            cosine_offset: float = 1.0, **kwargs):
109        """
110        Generates an embedding script score query for Elasticsearch as part of the NIR scoring function.
111
112        :param topic_num: The topic number to search for
113        :param encoder: The encoder that will be used for encoding the topics
114        :param norm_weight: The BM25 log normalization constant
115        :param ablations: Whether to execute ablation style queries (i.e. one query facet
116                          or one document facet at a time)
117        :param cosine_ceiling: Cosine ceiling used for automatic z-log normalization parameter calculation
118        :param args:
119        :param kwargs: Pass disable_cache to disable encoder caching
120        :return:
121            An elasticsearch script_score query
122        """
123
124        qfields = list(self.topics[topic_num].keys())
125        should = {"should": []}
126
127        if self.has_bm25_scores():
128            cosine_ceiling = len(self.embed_mappings) * len(qfields) if cosine_ceiling is None else cosine_ceiling
129            norm_weight = get_z_value(
130                cosine_ceiling=cosine_ceiling,
131                bm25_ceiling=self.top_bm25_scores[topic_num],
132            )
133            loguru.logger.debug(f"Automatic norm_weight: {norm_weight}")
134
135        params = {
136            "weights": [1] * (len(self.embed_mappings) * len(self.mappings)),
137            "offset": cosine_offset,
138            "norm_weight": norm_weight,
139            "disable_bm25": ablations,
140        }
141
142        embed_fields = []
143
144        for qfield in qfields:
145            for field in self.mappings:
146                should["should"].append(
147                    {
148                        "match": {
149                            f"{field}": {
150                                "query": self.topics[topic_num][qfield],
151                            }
152                        }
153                    }
154                )
155
156            params[f"{qfield}_eb"] = encoder.encode(topic=self.topics[topic_num][qfield])
157            embed_fields.append(f"{qfield}_eb")
158
159        query = {
160            "query": {
161                "script_score": {
162                    "query": {
163                        "bool": should,
164                    },
165                    "script": generate_script(
166                        self.embed_mappings, params, qfields=embed_fields
167                    ),
168                }
169            }
170        }
171
172        loguru.logger.debug(query)
173        return query
174
175    @classmethod
176    def get_id_mapping(cls, hit):
177        """
178        Get the document ID
179
180        :param hit: The raw document result
181        :return:
182            The document's ID
183        """
184        return hit[cls.id_mapping]