hojichar.core.models

  1import time
  2from dataclasses import dataclass, fields
  3from typing import Any, Dict, List, Optional
  4
  5from hojichar.utils.warn_deprecation import deprecated_since
  6
  7
  8@deprecated_since("0.1.0", "Document")
  9class Token:
 10    def __init__(self, text: str, is_rejected: bool = False) -> None:
 11        self.text = text
 12        self.__original = text
 13        self.is_rejected = is_rejected
 14
 15    @property
 16    def original(self) -> str:
 17        return self.__original
 18
 19    def __str__(self) -> str:
 20        return self.text
 21
 22
 23class Document:
 24    """
 25    Document class represents a text document with metadata.
 26    It contains the text of the document, a flag indicating whether it is rejected,
 27     and additional metadata stored in the `extras` dictionary.
 28
 29    The `tokens` attribute will be deprecated in future versions,
 30    and users are encouraged to use the `extras` dictionary to store token-related information.
 31
 32    Attributes:
 33        text (str): The text content of the document.
 34        is_rejected (bool): A flag indicating whether the document is rejected.
 35        extras (Dict[str, Any]): A dictionary to store additional metadata about the document.
 36        reject_reason (Dict[str, Any]): A dictionary to store the reason for rejection. The
 37          filter class and the member name and value will logged at the filter is logged here.
 38        initial_stats (Optional[Dict[str, Any]]): Internal copy of the document statistics
 39          captured before the pipeline mutates the document.
 40
 41    Next attributes will be deprecated in future versions:
 42        tokens (List[Token]): A list of tokens extracted from the document.
 43    """
 44
 45    def __init__(
 46        self,
 47        text: str,
 48        is_rejected: bool = False,
 49        tokens: Optional[List[Token]] = None,
 50        extras: Optional[Dict[str, Any]] = None,
 51    ) -> None:
 52        self.text = text
 53        self.__original = text
 54        self.is_rejected = is_rejected
 55        if tokens is None:
 56            self.tokens: List[Token] = []
 57        else:
 58            self.tokens = tokens
 59
 60        if extras is None:
 61            self.extras: Dict[str, Any] = {}
 62        else:
 63            self.extras = extras
 64
 65        self.reject_reason: Dict[str, Any] = {}
 66        self._initial_stats: Optional[dict[str, Any]] = None
 67        if "__init_stats" in self.extras:
 68            self._initial_stats = self.extras.pop("__init_stats")
 69
 70    @property
 71    def original(self) -> str:
 72        return self.__original
 73
 74    def _set_initial_stats(self, stats: dict[str, Any]) -> None:
 75        """
 76        Store the document statistics captured before the pipeline modifies the document.
 77        Internal API: not intended for filter implementations.
 78        """
 79        self._initial_stats = stats
 80
 81    def _get_initial_stats(self) -> Optional[dict[str, Any]]:
 82        """
 83        Retrieve the stored initial statistics of the document, if available.
 84        Internal API: not intended for filter implementations.
 85        """
 86        return self._initial_stats
 87
 88    def _clear_initial_stats(self) -> None:
 89        """
 90        Remove the stored initial statistics. Useful after the stats are consumed.
 91        Internal API: not intended for filter implementations.
 92        """
 93        self._initial_stats = None
 94
 95    @deprecated_since("1.0.0")
 96    def set_tokens(self, tokens: List[str]) -> None:
 97        self.tokens = [Token(token) for token in tokens]
 98
 99    @deprecated_since("1.0.0")
100    def get_tokens(self) -> List[str]:
101        return [token.text for token in self.tokens]
102
103    def __str__(self) -> str:
104        return self.text
105
106    def __repr__(self) -> str:
107        return (
108            f"Document(text={self.text!r}, is_rejected={self.is_rejected}, extras={self.extras})"  # noqa
109        )
110
111
112@dataclass
113class Statistics:
114    """
115    Statistics class to track the performance of the document processing pipeline.
116    """
117
118    name: Optional[str] = None
119    input_num: int = 0
120    input_bytes: int = 0
121    input_chars: int = 0
122    output_num: int = 0
123    output_bytes: int = 0
124    output_chars: int = 0
125    discard_num: int = 0
126    diff_bytes: int = 0
127    diff_chars: int = 0
128    cumulative_time_ns: int = 0
129    errors: int = 0
130
131    def to_dict(self) -> Dict[str, Any]:
132        """
133        Convert the Statistics object to a dictionary.
134        """
135        return {f.name: getattr(self, f.name) for f in fields(self)}
136
137    def update(self, other: "Statistics") -> None:
138        """
139        Update the statistics by adding another Statistics object.
140        """
141        self.input_num += other.input_num
142        self.input_bytes += other.input_bytes
143        self.input_chars += other.input_chars
144        self.output_num += other.output_num
145        self.output_bytes += other.output_bytes
146        self.output_chars += other.output_chars
147        self.discard_num += other.discard_num
148        self.diff_bytes += other.diff_bytes
149        self.diff_chars += other.diff_chars
150        self.cumulative_time_ns += other.cumulative_time_ns
151        self.errors += other.errors
152
153    def reset(self) -> "Statistics":
154        """
155        Reset the statistics to their initial values.
156        """
157        self.input_num = 0
158        self.input_bytes = 0
159        self.input_chars = 0
160        self.output_num = 0
161        self.output_bytes = 0
162        self.output_chars = 0
163        self.discard_num = 0
164        self.diff_bytes = 0
165        self.diff_chars = 0
166        self.cumulative_time_ns = 0
167        self.errors = 0
168        return self
169
170    def update_by_diff(
171        self,
172        before: dict[str, Any],
173        after: dict[str, Any],
174    ) -> None:
175        """
176        Update the statistics by calculating the differences between two Doc-info mappings.
177        This method is used to update the statistics after a filter is applied.
178        """
179        if not before["is_rejected"] and after["is_rejected"]:
180            # Document is rejected after the filter is applied
181            self.input_num += 1
182            self.input_bytes += before["bytes"]
183            self.input_chars += before["chars"]
184            self.discard_num += 1
185            self.diff_bytes -= before["bytes"]
186            self.diff_chars -= before["chars"]
187            self.cumulative_time_ns += after["time_ns"] - before["time_ns"]
188        else:
189            # Document is not rejected or still not rejected after the filter is applied
190            self.input_num += 1
191            self.input_bytes += before["bytes"]
192            self.input_chars += before["chars"]
193            self.output_num += 1
194            self.output_bytes += after["bytes"]
195            self.output_chars += after["chars"]
196            self.diff_bytes += after["bytes"] - before["bytes"]
197            self.diff_chars += after["chars"] - before["chars"]
198            self.cumulative_time_ns += after["time_ns"] - before["time_ns"]
199
200    @staticmethod
201    def add(x: "Statistics", y: "Statistics") -> "Statistics":
202        """
203        Add two Statistics objects together.
204        This method assumes that the names of the two Statistics objects match.
205        If they do not match, it will raise an AssertionError."""
206        assert x.name == y.name, "Layer names must match"
207        return Statistics(
208            name=x.name,
209            input_num=x.input_num + y.input_num,
210            input_bytes=x.input_bytes + y.input_bytes,
211            input_chars=x.input_chars + y.input_chars,
212            output_num=x.output_num + y.output_num,
213            output_bytes=x.output_bytes + y.output_bytes,
214            output_chars=x.output_chars + y.output_chars,
215            discard_num=x.discard_num + y.discard_num,
216            diff_bytes=x.diff_bytes + y.diff_bytes,
217            diff_chars=x.diff_chars + y.diff_chars,
218            cumulative_time_ns=x.cumulative_time_ns + y.cumulative_time_ns,
219            errors=x.errors + y.errors,
220        )
221
222    @staticmethod
223    def add_list_of_stats(x: List["Statistics"], y: List["Statistics"]) -> List["Statistics"]:
224        """
225        Add FilterStatistics objects from two lists by matching their names.
226        This method assumes that both lists contain FilterStatistics objects
227        with the same names, and it will raise a ValueError if the sets of names
228        in the two lists do not match.
229        """
230        # check if the names in both lists match
231        names_x = {stat.name for stat in x}
232        names_y = {stat.name for stat in y}
233        if names_x != names_y:
234            raise ValueError(f"name の集合が一致しません: {names_x} vs {names_y}")
235
236        y_map = {stat.name: stat for stat in y}
237
238        # keep the order of x and add corresponding y
239        result: List[Statistics] = []
240        for stat_x in x:
241            stat_y = y_map[stat_x.name]
242            result.append(Statistics.add(stat_x, stat_y))
243
244        return result
245
246    @staticmethod
247    def get_filter(name: str, stats: List["Statistics"]) -> "Statistics":
248        """
249        Get a Statistics object by its name from a list of statistics.
250        If the name is not found, return None.
251        """
252        for stat in stats:
253            if stat.name == name:
254                return stat
255        raise KeyError(f"Statistics with name '{name}' not found in the list.")
256
257
258def get_doc_info(document: Document) -> dict[str, Any]:
259    """
260    Create a document-info mapping from a Document instance.
261    This function is used to extract metadata from the Document for statistics tracking.
262    """
263    return {
264        "is_rejected": document.is_rejected,
265        "bytes": len(document.text.encode("utf-8")),
266        "chars": len(document.text),
267        "time_ns": time.perf_counter_ns(),
268    }