debeir.core.executor

  1from typing import Dict, Optional, Union
  2
  3import loguru
  4from debeir.engines.elasticsearch.executor import ElasticsearchExecutor
  5from debeir.core.config import GenericConfig, NIRConfig
  6from debeir.core.query import GenericElasticsearchQuery
  7from debeir.rankers.transformer_sent_encoder import Encoder
  8from debeir.utils.scaler import unpack_elasticsearch_scores
  9from elasticsearch import AsyncElasticsearch as Elasticsearch
 10
 11
 12class GenericElasticsearchExecutor(ElasticsearchExecutor):
 13    """
 14    Generic Executor class for Elasticsearch
 15    """
 16    query: GenericElasticsearchQuery
 17
 18    def __init__(
 19            self,
 20            topics: Dict[Union[str, int], Dict[str, str]],
 21            client: Elasticsearch,
 22            index_name: str,
 23            output_file: str,
 24            query: GenericElasticsearchQuery,
 25            encoder: Optional[Encoder] = None,
 26            config=None,
 27            *args,
 28            **kwargs,
 29    ):
 30        super().__init__(
 31            topics,
 32            client,
 33            index_name,
 34            output_file,
 35            query,
 36            encoder,
 37            config=config,
 38            *args,
 39            **kwargs,
 40        )
 41
 42        self.query_fns = {
 43            "query": self.generate_query,
 44            "embedding": self.generate_embedding_query,
 45        }
 46
 47    def generate_query(self, topic_num, best_fields=True, **kwargs):
 48        """
 49        Generates a standard BM25 query given the topic number
 50
 51        :param topic_num: Query topic number to generate
 52        :param best_fields: Whether to use a curated list of fields
 53        :param kwargs:
 54        :return:
 55        """
 56        return self.query.generate_query(topic_num, **kwargs)
 57
 58    # def generate_query_ablation(self, topic_num, **kwargs):
 59    #    return self.query.generate_query_ablation(topic_num)
 60
 61    def generate_embedding_query(
 62            self,
 63            topic_num,
 64            cosine_weights=None,
 65            query_weights=None,
 66            norm_weight=2.15,
 67            automatic_scores=None,
 68            **kwargs,
 69    ):
 70        """
 71        Executes an NIR-style query with combined scoring.
 72
 73        :param topic_num:
 74        :param cosine_weights:
 75        :param query_weights:
 76        :param norm_weight:
 77        :param automatic_scores:
 78        :param kwargs:
 79        :return:
 80        """
 81        assert self.encoder is not None or self.config.encoder is not None
 82
 83        if "encoder" not in kwargs:
 84            kwargs["encoder"] = self.encoder
 85
 86        return self.query.generate_query_embedding(
 87            topic_num,
 88            cosine_weights=cosine_weights,
 89            query_weight=query_weights,
 90            norm_weight=norm_weight,
 91            automatic_scores=automatic_scores,
 92            **kwargs,
 93        )
 94
 95    # @apply_config
 96    async def execute_query(
 97            self, query=None, return_size: int = None, return_id_only: bool = None,
 98            topic_num=None, ablation=False, query_type=None,
 99            **kwargs
100    ):
101        """
102        Executes a query using the underlying elasticsearch client.
103
104        :param query:
105        :param topic_num:
106        :param ablation:
107        :param query_type:
108        :param return_size:
109        :param return_id_only:
110        :param kwargs:
111        :return:
112        """
113
114        if ablation:
115            query_type = "ablation"
116
117        assert query is not None or topic_num is not None
118
119        if query:
120            if return_id_only:
121                # query["fields"] = [self.query.id_mapping]
122                # query["_source"] = False
123                query["_source"] = [self.query.id_mapping]
124            res = await self.client.search(
125                index=self.index_name, body=query, size=return_size
126            )
127
128            return [query, res]
129
130        if topic_num:
131            loguru.logger.debug(query_type)
132            body = self.query_fns[query_type](topic_num=topic_num, **kwargs)
133            if return_id_only:
134                loguru.logger.debug("Skip")
135                body["_source"] = [self.query.id_mapping]
136
137            loguru.logger.debug(body)
138            res = await self.client.search(
139                index=self.index_name, body=body, size=return_size
140            )
141
142            return [topic_num, res]
143
144    async def run_automatic_adjustment(self, return_results=False):
145        """
146        Get the normalization constant to be used in NIR-style queries for all topics given an initial
147        run of BM25 results.
148        """
149        loguru.logger.info("Running automatic BM25 weight adjustment")
150
151        # Backup variables temporarily
152        # size = self.return_size
153        # self.return_size = 1
154        # self.return_id_only = True
155        # prev_qt = self.config.query_type
156        # self.config.query_type = "query"
157
158        results = await self.run_all_queries(query_type="query",
159                                             return_results=True,
160                                             return_size=1,
161                                             return_id_only=True)
162
163        res = unpack_elasticsearch_scores(results)
164        self.query.set_bm25_scores(res)
165
166        if return_results:
167            return results
168
169    @classmethod
170    def build_from_config(cls, topics: Dict, query_obj: GenericElasticsearchQuery, client,
171                          config: GenericConfig, nir_config: NIRConfig):
172        """
173        Build an query executor engine from a config file.
174        """
175
176        return cls(
177            topics=topics,
178            client=client,
179            config=config,
180            index_name=config.index,
181            output_file="",
182            return_size=nir_config.return_size,
183            query=query_obj
184        )