hojichar.core.async_composition

  1from __future__ import annotations
  2
  3import asyncio
  4import logging
  5from concurrent.futures import Executor, ThreadPoolExecutor
  6from typing import Any, AsyncGenerator, AsyncIterable, Iterable, Sequence
  7
  8import numpy as np
  9
 10from hojichar.core import inspection
 11from hojichar.core.async_filter_interface import AsyncFilter
 12from hojichar.core.composition import Compose
 13from hojichar.core.filter_interface import Filter
 14from hojichar.core.models import Document, Statistics, get_doc_info
 15from hojichar.utils.async_handlers import handle_stream_as_async
 16
 17
 18class AsyncFilterAdapter(AsyncFilter):
 19    """
 20    Adapter class for executing hojichar.Filter asynchronously.
 21    """
 22
 23    def __init__(
 24        self,
 25        sync_filter: Filter,
 26        *args: Any,
 27        executor: Executor | None = None,
 28        use_batch: bool = True,
 29        **kwargs: Any,
 30    ):
 31        """
 32        Adapter class for executing hojichar.Filter asynchronously.
 33        Used to incorporate Filter into AsyncCompose.
 34
 35        To reduce the overhead of asynchronous context switching,
 36        use_batch is set to True by default to process in batches
 37        in apply_stream, regardless of the sync_filter's use_batch setting.
 38
 39        If performing CPU-bound and heavy processing, you can specify an executor
 40        to offload the processing to the executor. However, due to Python's GIL
 41        constraints, using ThreadPoolExecutor will not parallelize CPU-bound
 42        processing, and the entire process will be locked.
 43
 44        By using ProcessPoolExecutor as the executor, it may be possible to
 45        parallelize CPU-bound processing. However, for parallelizing CPU-bound
 46        processing, it is recommended to use the hojichar.Parallel class to
 47        parallelize synchronous Compose pipeline.
 48        """
 49        super().__init__(*args, use_batch=use_batch, **kwargs)
 50        self.sync_filter = sync_filter
 51        self._has_external_executor = executor is not None
 52        self._executor = executor or ThreadPoolExecutor()
 53        self.batch_size = sync_filter.batch_size
 54
 55    async def apply(self, document: Document) -> Document:
 56        loop = asyncio.get_running_loop()
 57        return await loop.run_in_executor(self._executor, self.sync_filter.apply, document)
 58
 59    async def apply_batch(self, batch: Sequence[Document]) -> list[Document]:
 60        loop = asyncio.get_running_loop()
 61        return await loop.run_in_executor(
 62            self._executor,
 63            lambda: self.sync_filter.apply_batch(batch),
 64        )
 65
 66    async def shutdown(self) -> None:
 67        self.sync_filter.shutdown()
 68        if not self._has_external_executor:
 69            self._executor.shutdown()
 70
 71
 72class AsyncCompose(AsyncFilter):
 73    def __init__(
 74        self,
 75        filters: list[AsyncFilter | Filter],
 76        random_state: int | np.random.Generator | None = None,
 77        executor: ThreadPoolExecutor | None = None,
 78        *args: Any,
 79        **kwargs: Any,
 80    ):
 81        super().__init__(random_state=random_state, *args, **kwargs)
 82        self.logger = logging.getLogger(f"{self.__module__}.{self.__class__.__name__}")
 83        self._statistics.name = "Total"
 84        self._has_external_executor = executor is not None
 85        self._executor = executor or ThreadPoolExecutor()
 86        self.set_filters(filters)
 87
 88    def set_filters(self, filters: list[AsyncFilter | Filter]) -> None:
 89        self.filters: list[AsyncFilter] = []
 90        filter_idx = 0
 91        for f in filters:
 92            if isinstance(f, (AsyncCompose, Compose)):
 93                for sub in f.filters:
 94                    name = f"{filter_idx}-{sub.__class__.__name__}"
 95                    if isinstance(sub, Filter):
 96                        name = f"{filter_idx}-{sub.__class__.__name__}"
 97                        sub = AsyncFilterAdapter(sub, executor=self._executor)
 98
 99                    sub._set_rng_if_not_initialized(self._rng)
100                    sub.name = name
101                    sub._statistics.name = name
102                    self.filters.append(sub)
103                    filter_idx += 1
104            else:
105                name = f"{filter_idx}-{f.__class__.__name__}"
106                if isinstance(f, Filter):
107                    name = f"{filter_idx}-{f.__class__.__name__}"
108                    f = AsyncFilterAdapter(f, executor=self._executor)
109                f._set_rng_if_not_initialized(self._rng)
110                f.name = name
111                f._statistics.name = name
112                self.filters.append(f)
113                filter_idx += 1
114
115    async def apply(self, document: Document) -> Document:
116        stat = get_doc_info(document)
117        for filter_idx, filt in enumerate(self.filters):
118            document = await filt._apply(document)
119        new_stat = get_doc_info(document)
120        async with self._stats_lock:
121            self._statistics.update_by_diff(stat, new_stat)
122        return document
123
124    async def apply_batch(self, batch: Sequence[Document]) -> list[Document]:
125        stats = [get_doc_info(doc) for doc in batch]
126        for i, filt in enumerate(self.filters):
127            batch = await filt._apply_batch(batch)
128        batch = await self._finalize_batch(batch, stats)
129        return list(batch)
130
131    async def apply_stream(
132        self,
133        stream: AsyncIterable[Document] | Iterable[Document],
134    ) -> AsyncGenerator[Document, None]:
135        async_stream = handle_stream_as_async(stream, chunk_size=1000, executor=self._executor)
136        async_stream = self._count_input_stats(async_stream)
137
138        for i, filt in enumerate(self.filters):
139            async_stream = filt.apply_stream(async_stream)
140
141        async for doc in async_stream:
142            in_stat = doc.extras["__init_stats"]
143            out_stat = get_doc_info(doc)
144            async with self._stats_lock:
145                self._statistics.update_by_diff(in_stat, out_stat)
146            del doc.extras["__init_stats"]
147            yield doc
148
149    async def _count_input_stats(
150        self, async_stream: AsyncIterable[Document]
151    ) -> AsyncGenerator[Document, None]:
152        async for doc in async_stream:
153            doc.extras["__init_stats"] = get_doc_info(doc)
154            yield doc
155
156    def get_total_statistics(self) -> list[Statistics]:
157        """
158        Get the statistics of the Compose object and sub filters.
159
160        The statistics of the Compose class are stored in an object with the name "Total",
161        and sub-filters's are stored with names in the format {filter_index}-{filter class name}.
162        """
163        stats = []
164        stats.append(self.get_statistics())
165        for i, filt in enumerate(self.filters):
166            stats.append(filt.get_statistics())
167        return stats
168
169    def get_total_statistics_map(self) -> list[dict[str, Any]]:
170        """
171        Get the statistics of the Compose object and sub filters as a list of dictionaries.
172        """
173        stats = self.get_total_statistics()
174        return [stat.to_dict() for stat in stats]
175
176    @property
177    def statistics(self) -> dict:
178        """
179        Deprecated
180
181        Get the statistics of the Compose object and sub filters.
182
183        This property is retained for compatibility with previous versions.
184        Please use `get_total_statistics` or `get_total_statistics_map` instead.
185        """
186        return inspection.statistics_obj_adapter(  # type: ignore
187            self.get_total_statistics()
188        ).get_human_readable_values()
189
190    @property
191    def statistics_obj(self) -> inspection.StatsContainer:
192        """
193        Deprecated
194
195        Get the statistics of the AsyncCompose object and sub filters.
196        This method returns a StatsContainer object which contains the statistics
197        of the AsyncCompose object and sub filters.
198
199        This property is retained for compatibility with previous versions.
200        Please use `get_total_statistics` or `get_total_statistics_map` instead.
201        """
202        return inspection.statistics_obj_adapter(self.get_total_statistics())  # type: ignore
203
204    async def shutdown(self) -> None:
205        for filt in self.filters:
206            await filt.shutdown()
207        if not self._has_external_executor:
208            self._executor.shutdown()
209
210    async def __aenter__(self) -> "AsyncCompose":
211        return self
212
213    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
214        await self.shutdown()
215        if exc_type is not None:
216            raise exc_value