debeir.core.document

  1import abc
  2import dataclasses
  3from collections import defaultdict
  4from typing import Dict, List, Union
  5
  6from debeir.utils.utils import flatten
  7
  8
  9@dataclasses.dataclass
 10class Document:
 11    """
 12    Generic Document class.
 13    Used as an interface for interacting across multiple indexes with different mappings.
 14    """
 15    doc_id: Union[int, float, str]
 16    topic_num: Union[int, str, float] = None
 17    facets: Dict = None
 18    score: Union[float, int] = 0.0  # Primay Score
 19    scores: Dict[str, Union[float, int]] = dataclasses.field(
 20        default_factory=lambda: {})  # Include other scores if needed
 21
 22    @classmethod
 23    @abc.abstractmethod
 24    def from_results(cls, results, *args, **kwargs) -> Dict[Union[int, float], 'Document']:
 25        """
 26        Produces a list of Document objects from raw results returned from the index
 27
 28        In the format {topic_num: [Document, ..., Document]}
 29        """
 30        pass
 31
 32    def get_document_id(self):
 33        """
 34        :return:
 35            self.doc_id
 36        """
 37        return self.doc_id
 38
 39    def flatten_facets(self, *args, **kwargs):
 40        """
 41        Flattens multi-level internal document facets into a single level
 42            e.g. Doc['Upper']['Lower'] -> Doc['Upper_Lower']
 43        :param args:
 44        :param kwargs:
 45        """
 46        self.facets = flatten(self.facets, *args, **kwargs)
 47
 48    @classmethod
 49    def _get_document_facet(cls, intermediate_repr, key):
 50        return intermediate_repr[key]
 51
 52    def get_document_facet(self, key, sep="_"):
 53        """
 54        Retrieve a document facet
 55        Works for multidimensional keys or single
 56        :param key: Facet to retrieve
 57        :param sep: The seperator for multidimensional key
 58        :return:
 59            Returns the document facet given the key (field)
 60        """
 61        if sep in key:
 62            keys = key.split(sep)
 63
 64            intermediate_repr = self.facets
 65            for k in keys:
 66                intermediate_repr = self._get_document_facet(intermediate_repr, k)
 67
 68            return intermediate_repr
 69
 70        return self.facets[key]
 71
 72    def set(self, doc_id=None, facets=None, score=None, facet=None, facet_value=None) -> 'Document':
 73        """
 74        Set attributes of the object. Use keyword arguments to do so. Works as a builder class.
 75        doc.set(doc_id="123").set(facets={"title": "my title"})
 76        :param doc_id:
 77        :param facets:
 78        :param score:
 79        :param facet:
 80        :param facet_value:
 81
 82        :return:
 83            Returns document object
 84        """
 85        if doc_id is not None:
 86            self.doc_id = doc_id
 87
 88        if facets is not None:
 89            self.facets = facets
 90
 91        if score is not None:
 92            self.score = score
 93
 94        if facet is not None and facet_value is not None:
 95            self.facets[facet] = facet_value
 96
 97        return self
 98
 99    def to_trec_format(self, rank, run_name) -> str:
100        """
101        Returns TREC format for the document
102        :return:
103            A trec formatted string
104        """
105
106        return f"{self.topic_num}\t" \
107               f"Q0\t" \
108               f"{self.doc_id}\t" \
109               f"{rank}\t" \
110               f"{self.score}\t" \
111               f"{run_name}\n"
112
113    @classmethod
114    def get_trec_format(cls, ranked_list: List['Document'], run_name="NO_RUN_NAME", sort=True, sorting_func=None):
115        """
116        Get the trec format of a list of ranked documents. This function is a generator.
117
118        :param ranked_list: A list of Document-type objects
119        :param run_name: Run name to print in the TREC formatted string
120        :param sort: Whether to sort the input list in descending order of score.
121        :param sorting_func: Custom sorting function will be used if provided
122        """
123
124        if sort:
125            if sorting_func:
126                ranked_list = sorting_func(ranked_list)
127            else:
128                ranked_list.sort(key=lambda doc: doc.score, reverse=True)
129
130        for rank, document in enumerate(ranked_list, start=1):
131            yield document.to_trec_format(rank, run_name)
132
133
134class ElasticsearchDocument(Document):
135    @classmethod
136    def from_results(cls, results, query_cls, ignore_facets=True,
137                     *args, **kwargs) -> Dict[Union[int, float], 'Document']:
138
139        documents = defaultdict(lambda: [])
140
141        for (topic_num, res) in results:
142            for rank, result in enumerate(res["hits"]["hits"], start=1):
143                doc_id = query_cls.get_id_mapping(result["_source"])
144                facets = {}
145
146                if not ignore_facets:
147                    facets = {k: v for (k, v) in result['_source'].items() if not k.startswith("_")}
148
149                documents[topic_num].append(ElasticsearchDocument(doc_id,
150                                                                  topic_num,
151                                                                  facets=facets,
152                                                                  score=float(result['_score'])))
153
154                documents[topic_num][-1].scores['rank'] = rank
155
156        return dict(documents)
157
158
159document_factory = {
160    "elasticsearch": ElasticsearchDocument
161}