15class Compose(Filter):
 16    def __init__(
 17        self,
 18        filters: List[Union[Filter, TokenFilter]],
 19        random_state: Optional[Union[int, np.random.Generator]] = None,
 20        *args: Any,
 21        **kwargs: Any,
 22    ) -> None:
 23        """
 24        Compose a filter from pre-defined filter-objects.
 25        Filter which has `skip_rejected` flag ignores a document which has `is_rejected` flag.
 26        By doing so, Compose avoid applying filters that do not affect the output.
 27
 28        Parameters
 29        ----------
 30        filters : List[Union[Filter, TokenFilter]]
 31            Filter instances which apply to the corpus.
 32
 33        random_state : Union[None, int, np.random.Generator], optional
 34            Default = None
 35            Seed for applying filters randomly.
 36            `random_state` must be int or np.random.Generator instance.
 37        """
 38        super().__init__(random_state=random_state, *args, **kwargs)
 39        self.set_filters(filters)
 40        self.logger = logging.getLogger(f"{self.__module__}.{self.__class__.__name__}")
 41
 42        self._statistics.name = "Total"
 43
 44    def set_filters(self, filters: List[Union[Filter, TokenFilter]]) -> None:
 45        """
 46        Set the filter to a Compose object. The filter is expanded if the
 47        list of filters in the argument contains a filter bound by Compose.
 48
 49        Args:
 50            filters (List[Union[Filter, TokenFilter]]): Target filters
 51        """
 52        self.filters: List[Union[Filter, TokenFilter]] = []
 53
 54        filter_idx = 0
 55        for f in filters:
 56            if isinstance(f, Compose):
 57                for sub in f.filters:
 58                    sub._set_rng_if_not_initialized(self._rng)
 59                    name = f"{filter_idx}-{sub.__class__.__name__}"
 60                    sub.name = name
 61                    sub._statistics.name = name
 62                    self.filters.append(sub)
 63                    filter_idx += 1
 64            else:
 65                f._set_rng_if_not_initialized(self._rng)
 66                name = f"{filter_idx}-{f.__class__.__name__}"
 67                f.name = name
 68                f._statistics.name = name
 69                self.filters.append(f)
 70                filter_idx += 1
 71
 72    def __call__(self, text: str, **kwargs: Any) -> str:
 73        """
 74        Apply the composed filter to a text and return the processed text.
 75        If the document is rejected, return an empty string.
 76        """
 77        document = Document(text, **kwargs)
 78        document = self.apply(document)
 79        if document.is_rejected:
 80            return ""
 81        else:
 82            return document.text
 83
 84    def apply(self, document: Document) -> Document:
 85        """
 86        Apply the composed filter to a document and return the processed document.
 87        """
 88        stat = get_doc_info(document)
 89        for i, filt in enumerate(self.filters):
 90            document = filt._apply(document)
 91        new_stat = get_doc_info(document)
 92        self._statistics.update_by_diff(stat, new_stat)
 93        return document
 94
 95    def apply_batch(self, batch: Sequence[Document]) -> List[Document]:
 96        """
 97        Apply the composed filter to a batch of documents and return the processed documents.
 98        The `apply_batch` method implemented in sub-filters is called in order.
 99        """
100
101        stats = [get_doc_info(doc) for doc in batch]
102        for i, filt in enumerate(self.filters):
103            batch = filt._apply_batch(batch)
104        batch = self._finalize_batch(batch, stats)
105        return list(batch)
106
107    def apply_stream(self, stream: Iterable[Document]) -> Iterable[Document]:
108        """
109        Apply the composed filter to a stream of documents and return the processed documents.
110        The `apply_stream` method implemented in sub-filters is called in order.
111
112
113        In a sub-filter, if `apply_batch` is overridden and implemented, you need to set `use_batch`
114        to True at that filter to utilize that implementation. Otherwise, the
115        method implemented in `apply` will be applied to the stream.
116        """
117        stream = self._count_input_stats(stream)
118        for i, filt in enumerate(self.filters):
119            stream = filt.apply_stream(stream)
120
121        for doc in stream:
122            in_stat = doc._get_initial_stats()
123            if in_stat is None:
124                in_stat = get_doc_info(doc)
125                self.logger.debug(
126                    "Initial stats missing for document during stream aggregation; "
127                    "using current stats as fallback"
128                )
129            out_stat = get_doc_info(doc)
130
131            self._statistics.update_by_diff(in_stat, out_stat)
132            doc._clear_initial_stats()
133            yield doc
134
135    def _count_input_stats(self, stream: Iterable[Document]) -> Iterable[Document]:
136        for doc in stream:
137            doc._set_initial_stats(get_doc_info(doc))
138            yield doc
139
140    def get_total_statistics(self) -> List[Statistics]:
141        """
142        Get the statistics of the Compose object and sub filters.
143
144        The statistics of the Compose class are stored in an object with the name "Total",
145        and sub-filters's are stored with names in the format {filter_index}-{filter class name}.
146        """
147        stats = []
148        stats.append(self.get_statistics())
149        for i, filt in enumerate(self.filters):
150            stats.append(filt.get_statistics())
151        return stats
152
153    def get_total_statistics_map(self) -> List[Dict[str, Any]]:
154        """
155        Get the statistics of the Compose object and sub filters as a list of dictionaries.
156        """
157        stats = self.get_total_statistics()
158        return [stat.to_dict() for stat in stats]
159
160    def shutdown(self) -> None:
161        for f in self.filters:
162            f.shutdown()
163
164        super().shutdown()
165
166    @property
167    def statistics(self) -> dict:
168        """
169        Deprecated
170
171        Get the statistics of the Compose object and sub filters.
172
173        This property is retained for compatibility with previous versions.
174        Please use `get_total_statistics` or `get_total_statistics_map` instead.
175        """
176        return inspection.statistics_obj_adapter(  # type: ignore
177            self.get_total_statistics()
178        ).get_human_readable_values()
179
180    @property
181    def statistics_obj(self) -> inspection.StatsContainer:
182        """
183        Deprecated
184
185        Get the statistics of the Compose object and sub filters.
186        This method returns a StatsContainer object which contains the statistics
187        of the Compose object and sub filters.
188
189        This property is retained for compatibility with previous versions.
190        Please use `get_total_statistics` or `get_total_statistics_map` instead.
191        """
192        return inspection.statistics_obj_adapter(self.get_total_statistics())  # type: ignore
193
194    @deprecated_since("1.0.0", "get_total_statistics")
195    def summary(self, format: str = "print") -> None:
196        info = [
197            {
198                "layer": i,
199                "name": filt.name,
200                "doc": filt.__doc__,
201            }
202            for i, filt in enumerate(self.filters)
203        ]
204
205        def to_json(filter_info: dict) -> dict:
206            filter_info["doc"] = "".join(d.strip() for d in filter_info["doc"].split("\n"))
207            return filter_info
208
209        if format == "json":
210            print(json.dumps(list(map(to_json, info)), ensure_ascii=False, indent="\t"))
211        if format == "print":
212            for layer in info:
213                print(f"[{layer['layer']}] {layer['name']}")
214                pprint.pprint(layer["doc"])

 22class Filter(ABC):
 23    """
 24    Base class for all filters.
 25    Document-level filters must inherit from this class.
 26
 27    The definition of text processing is in `apply` method.
 28    If you define a new filter, override the method.
 29
 30    When this class is called, apply the filter from string to string.
 31
 32    With context manager, you can use the filter as follows:
 33    ```python
 34    with YourFilter(p=0.5) as filt:
 35        text = filt("This is a sample text.")
 36    ```
 37
 38    """
 39
 40    def __init__(
 41        self,
 42        p: float = 1.0,
 43        skip_rejected: bool = True,
 44        *args: Any,
 45        random_state: Optional[Union[int, np.random.Generator]] = None,
 46        use_batch: bool = False,
 47        batch_size: int = 128,
 48        **kwargs: Any,
 49    ) -> None:
 50        """
 51        Initialize the filter.
 52        Parameters
 53        ----------
 54        p : float
 55            The probability of applying the filter.
 56            If `p` is 1, the filter will always be applied.
 57        skip_rejected : bool
 58            If `True`, the filter will skip documents that are already rejected.
 59            If you want to apply the filter to all documents (e.g., postprocess), set this to `False`.
 60        random_state : Optional[Union[int, np.random.Generator]]
 61            Seed for the random number generator.
 62            If `None`, a new random number generator will be created.
 63            If `None`, and use in the `Compose` class, the random state is shared with the `Compose` object.
 64        use_batch : bool
 65            If `True`, the filter will process documents in batches in the `apply_stream` method.
 66        batch_size : int
 67            The size of the batch to process documents in the `apply_stream` method.
 68        kwargs : Any
 69            Additional keyword arguments to pass to the filter.
 70        """
 71        self.name = self.__class__.__name__
 72        self.logger = logging.getLogger(f"{self.__module__}.{self.__class__.__name__}")
 73        assert 0 <= p <= 1
 74        self.p = p
 75        self.__init_rng(random_state)
 76        self.skip_rejected = skip_rejected
 77        self.use_batch = use_batch
 78        self.batch_size = batch_size
 79
 80        self._statistics: Statistics = Statistics()
 81
 82    @abstractmethod
 83    def apply(self, document: Document) -> Document:
 84        """
 85        Definition of filter behavior.
 86
 87        The document must have a protocol `TextContent`,
 88        and mostly used hojichar.Document class.
 89
 90        In this method, the filter will modify `document.text` or
 91        `document.extras` and set `document.is_rejected = True` to discard the document.
 92
 93        Parameters
 94        ----------
 95        document : Document
 96            Input document
 97
 98        Returns
 99        -------
100        Document
101            Processed Document
102        """
103
104    @deprecated_since(version="1.0.0", alternative="apply")
105    def apply_filter(self, document: Document) -> Document:
106        document = self.apply(document)
107        return document
108
109    def _check_skip(self, document: Document) -> bool:
110        """
111        Check if the document should be skipped by this filter.
112        If `skip_rejected` is set to `True`, this method will return `True`
113        if the document is already rejected.
114        If `p` is less than 1, this method will return `True` with a probability of `1 - p`.
115        """
116        skip = self.skip_rejected and document.is_rejected
117        if skip:
118            return True
119        if self.p < 1:
120            if self._rng.random() > self.p:
121                return True
122        return False
123
124    def _apply(self, document: Document) -> Document:
125        """
126        Apply the filter to a single document.
127        This method
128          - checks if the document should be skipped
129          - counts and logging the statistics
130          - logging the reason for rejection if the document is rejected
131
132        This method may be used in `apply` method of `Compose` class.
133        """
134
135        stats = get_doc_info(document)
136
137        if not self._check_skip(document):
138            document = self.apply(document)
139
140        new_stats = get_doc_info(document)
141        self._statistics.update_by_diff(stats, new_stats)
142
143        if not stats["is_rejected"] and new_stats["is_rejected"]:
144            document.reject_reason = self.get_jsonable_vars()
145
146        return document
147
148    def apply_batch(self, batch: Sequence[Document]) -> List[Document]:
149        """
150        Apply the filter to a batch of documents.
151        You can override this method if you want to
152        apply the filter to a batch of documents at once.
153
154        This method may be used in `apply_batch` method of `Compose` class.
155
156        Parameters
157        ----------
158        documents : Sequence[Document]
159            List-like object of input documents
160
161        Returns
162        -------
163        list[Document]
164            List of processed documents
165        """
166        return [self.apply(document) for document in batch]
167
168    def _apply_batch(self, batch: Sequence[Document]) -> List[Document]:
169        """
170        Apply the filter to a batch of documents.
171        This method
172        - checks if the documents should be skipped
173        - counts and logs the statistics
174        - logs the reason for rejection if any document is rejected
175        """
176        skip = False
177        if self.p < 1:
178            skip = self._rng.random() > self.p
179
180        stats = [get_doc_info(document=doc) for doc in batch]
181        if not skip:
182            batch = self.apply_batch(batch)
183        batch = self._finalize_batch(batch, stats)
184        return batch
185
186    def apply_stream(self, stream: Iterable[Document]) -> Iterable[Document]:
187        """
188        Apply the filter to a stream of documents.
189        This method is used when you want to process documents one by one.
190        If `use_batch` is set to `True` in the constructor,
191        this method will process documents in batches using the `apply_batch` method.
192
193        Even if an exception occurs during processing, the process will continue, and the following actions will be taken:
194        - Set the `is_rejected` flag of the document to `True`
195        - Set the error details in `reject_reason`
196        - Increment the `errors` count in the statistics retrievable via `get_statistics`
197
198        Parameters
199        ----------
200        stream : Iterable[Document]
201            Stream of input documents
202
203        Returns
204        -------
205        Iterable[Document]
206            Stream of processed documents
207        """
208
209        if not self.use_batch:
210            for document in stream:
211                yield self._try_process(document, self._apply)
212        else:
213            batch: list[Document] = []
214            for document in stream:
215                if self._check_skip(document):
216                    yield document
217                    continue
218
219                batch.append(document)
220                if len(batch) >= self.batch_size:
221                    stats = [get_doc_info(doc) for doc in batch]
222                    batch = self._try_process(batch, self.apply_batch)
223                    batch = self._finalize_batch(batch, stats)
224                    yield from batch
225                    batch.clear()
226            if batch:
227                stats = [get_doc_info(doc) for doc in batch]
228                batch = self._try_process(batch, self.apply_batch)
229                batch = self._finalize_batch(batch, stats)
230                yield from batch
231
232    def _try_process(self, target: T, func: Callable[[T], T]) -> T:
233        try:
234            return func(target)
235        except Exception as e:
236            if isinstance(target, Document):
237                msg = f"{e!r} occurs while processing {self.name} with {target!r}"
238                target.is_rejected = True
239                target.reject_reason = {"error": msg}
240                self._statistics.errors += 1
241                self.logger.error(msg, exc_info=True)
242                return target
243            if isinstance(target, list):
244                msg = f"{e!r} occurs while batch processing {self.name}"
245                self.logger.error(msg, exc_info=True)
246                for doc in target:
247                    doc.is_rejected = True
248                    doc.reject_reason = {"error": msg}
249                self._statistics.errors += len(target)
250                return target
251            else:
252                raise e
253
254    def __call__(self, text: str, **kwargs: Any) -> str:
255        document = Document(text, **kwargs)
256        document = self._apply(document)
257        return document.text
258
259    def get_statistics(self) -> Statistics:
260        """
261        Get the statistics of this filter.
262        This method returns the statistics of the filter,
263        which includes the number of processed documents, discarded documents, and other statistics.
264        """
265        return self._statistics
266
267    def get_statistics_map(self) -> Dict[str, Any]:
268        """
269        Get the statistics of this filter as a dictionary.
270        """
271        return self._statistics.to_dict()
272
273    def shutdown(self) -> None:
274        """
275        This method is called when the filter is no longer needed.
276        You can override this method to release resources or perform cleanup tasks.
277        """
278        pass
279
280    def __enter__(self) -> "Filter":
281        return self
282
283    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
284        """
285        This method is called when the filter is used in a context manager.
286        It calls the `shutdown` method to release resources or perform cleanup tasks.
287        """
288        self.shutdown()
289
290    def get_jsonable_vars(self, exclude_keys: Optional[Set[str]] = None) -> Dict[str, Any]:
291        """
292        Get the member variable of this filter.
293        Eligible variables are primitive types; [bool, int, float, str, None],
294        and the name of the variable not starts with the underscore; `_`.
295        """
296        if exclude_keys is None:
297            exclude_keys = set()
298        return {
299            k: v
300            for k, v in vars(self).items()
301            if (_is_jsonable(v) and (k not in exclude_keys) and (not k.startswith("_")))
302        }
303
304    def _finalize_batch(
305        self: "Filter",
306        batch: Sequence[Document],
307        old_stats: List[Dict[str, Any]] = [],
308    ) -> List[Document]:
309        new_stats = [get_doc_info(doc) for doc in batch]
310        for old, new, doc in zip(old_stats, new_stats, batch):
311            self._statistics.update_by_diff(old, new)
312            if not old["is_rejected"] and new["is_rejected"]:
313                doc.reject_reason = self.get_jsonable_vars()
314        return list(batch)
315
316    def __init_rng(self, random_state: Optional[Union[int, np.random.Generator]]) -> None:
317        self._owns_rng = True
318        if random_state is None:
319            self._rng = np.random.default_rng()
320            self._owns_rng = False
321        elif isinstance(random_state, int):
322            self._rng = np.random.default_rng(random_state)
323        elif isinstance(random_state, np.random.Generator):
324            self._rng = random_state
325
326    def _set_rng_if_not_initialized(self, rng: np.random.Generator) -> None:
327        """
328        Set the random number generator for this filter if it is not already initialized.
329        This method is called by Compose class.
330        """
331        if not self._owns_rng:
332            self._rng = rng