cbrkit.reuse.build

  1import itertools
  2from collections.abc import Sequence
  3from dataclasses import dataclass
  4from inspect import signature as inspect_signature
  5from multiprocessing.pool import Pool
  6from typing import cast, override
  7
  8from ..helpers import (
  9    batchify_adaptation,
 10    batchify_sim,
 11    chunkify,
 12    get_logger,
 13    mp_count,
 14    mp_map,
 15    mp_starmap,
 16    produce_factory,
 17    use_mp,
 18)
 19from ..typing import (
 20    AnyAdaptationFunc,
 21    AnySimFunc,
 22    Casebase,
 23    Float,
 24    MapAdaptationFunc,
 25    MaybeFactory,
 26    ReduceAdaptationFunc,
 27    ReuserFunc,
 28    SimMap,
 29    SimpleAdaptationFunc,
 30)
 31
 32logger = get_logger(__name__)
 33
 34__all__ = ["build"]
 35
 36
 37@dataclass(slots=True, frozen=True)
 38class build[K, V, S: Float](ReuserFunc[K, V, S]):
 39    """Builds a casebase by adapting cases using an adaptation function and a similarity function.
 40
 41    Args:
 42        adaptation_func: The adaptation function that will be applied to the cases.
 43        similarity_func: The similarity function that will be used to compare the adapted cases to the query.
 44        multiprocessing: Multiprocessing configuration for adaptation.
 45        chunksize: Number of batches to process at a time using the adaptation function.
 46            If 0, it will be set to the number of batches divided by the number of processes.
 47
 48    Returns:
 49        The adapted casebases and the similarities between the adapted cases and the query.
 50    """
 51
 52    adaptation_func: MaybeFactory[AnyAdaptationFunc[K, V]]
 53    similarity_func: MaybeFactory[AnySimFunc[V, S]]
 54    multiprocessing: Pool | int | bool = False
 55    chunksize: int = 0
 56
 57    @override
 58    def __call__(
 59        self,
 60        batches: Sequence[tuple[Casebase[K, V], V]],
 61    ) -> Sequence[tuple[Casebase[K, V], SimMap[K, S]]]:
 62        adaptation_func = cast(
 63            AnyAdaptationFunc[K, V], produce_factory(self.adaptation_func)
 64        )
 65        adapted_casebases = self._adapt(batches, adaptation_func)
 66        adapted_batches = [
 67            (adapted_casebase, query)
 68            for adapted_casebase, (_, query) in zip(
 69                adapted_casebases, batches, strict=True
 70            )
 71        ]
 72
 73        # Score adapted cases against queries
 74        sim_func = batchify_sim(produce_factory(self.similarity_func))
 75
 76        flat_batches: list[tuple[V, V]] = []
 77        flat_index: list[tuple[int, K]] = []
 78
 79        for idx, (casebase, query) in enumerate(adapted_batches):
 80            for key, case in casebase.items():
 81                flat_index.append((idx, key))
 82                flat_batches.append((case, query))
 83
 84        scores = sim_func(flat_batches)
 85
 86        sim_maps: list[dict[K, S]] = [{} for _ in adapted_batches]
 87        for (idx, key), score in zip(flat_index, scores, strict=True):
 88            sim_maps[idx][key] = score
 89
 90        return [
 91            (adapted_casebase, sim_map)
 92            for adapted_casebase, sim_map in zip(
 93                adapted_casebases, sim_maps, strict=True
 94            )
 95        ]
 96
 97    def _adapt(
 98        self,
 99        batches: Sequence[tuple[Casebase[K, V], V]],
100        adaptation_func: AnyAdaptationFunc[K, V],
101    ) -> Sequence[Casebase[K, V]]:
102        adaptation_func_signature = inspect_signature(adaptation_func)
103
104        if "casebase" in adaptation_func_signature.parameters:
105            adapt_func = cast(
106                MapAdaptationFunc[K, V] | ReduceAdaptationFunc[K, V],
107                adaptation_func,
108            )
109            adaptation_results = mp_starmap(
110                adapt_func,
111                batches,
112                self.multiprocessing,
113                logger,
114            )
115
116            if all(isinstance(item, tuple) for item in adaptation_results):
117                adaptation_results = cast(Sequence[tuple[K, V]], adaptation_results)
118                return [
119                    {adapted_key: adapted_case}
120                    for adapted_key, adapted_case in adaptation_results
121                ]
122
123            return cast(Sequence[Casebase[K, V]], adaptation_results)
124
125        adapt_func = batchify_adaptation(cast(SimpleAdaptationFunc[V], adaptation_func))
126        batches_index: list[tuple[int, K]] = []
127        flat_batches: list[tuple[V, V]] = []
128
129        for idx, (casebase, query) in enumerate(batches):
130            for key, case in casebase.items():
131                batches_index.append((idx, key))
132                flat_batches.append((case, query))
133
134        adapted_cases: Sequence[V]
135
136        if use_mp(self.multiprocessing) or self.chunksize > 0:
137            chunksize = (
138                self.chunksize
139                if self.chunksize > 0
140                else len(flat_batches) // mp_count(self.multiprocessing)
141            )
142            batch_chunks = list(chunkify(flat_batches, chunksize))
143            adapted_chunks = mp_map(
144                adapt_func, batch_chunks, self.multiprocessing, logger
145            )
146            adapted_cases = list(itertools.chain.from_iterable(adapted_chunks))
147        else:
148            adapted_cases = list(adapt_func(flat_batches))
149
150        adapted_casebases: list[dict[K, V]] = [{} for _ in range(len(batches))]
151
152        for (idx, key), adapted_case in zip(batches_index, adapted_cases, strict=True):
153            adapted_casebases[idx][key] = adapted_case
154
155        return adapted_casebases