File size: 11,143 Bytes
f8768df
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8947dc2
f8768df
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8947dc2
f8768df
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c4aee46
f8768df
 
 
 
 
c4aee46
 
f8768df
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
"""Public API for deterministic financial-news relevance screening."""

from __future__ import annotations

import hashlib
import os
from collections.abc import Iterable, Sequence
from datetime import datetime
from functools import lru_cache
from pathlib import Path
from typing import Literal

from ._core.accelerated.indicator import AcceleratedRelevanceIndicator
from ._core.accelerated.compact import (
    ArrowCompactFacts,
    ColumnArticleBatch,
    CompactRelevanceIndicator,
)
from ._core.models import Article, ExposureFamily, RelevanceDecision

from .backends import available_backends, create_backend
from .customization import (
    DefaultRanker,
    DefaultTaxonomy,
    Ranker,
    Taxonomy,
)
from .models import (
    BackendInfo,
    Prediction,
    PythonScreeningBatch,
    RankedPrediction,
)

TimestampLike = datetime | str


def _workers(value: int | None) -> int:
    if value is None:
        return max(1, min(10, os.cpu_count() or 1))
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError("workers must be a positive integer or None")
    return value


def _article_id(text: str, timestamp: TimestampLike) -> str:
    payload = f"{timestamp}\0{text}".encode("utf-8")
    return f"article-{hashlib.sha256(payload).hexdigest()[:20]}"


def _optional_length(

    name: str,

    values: Sequence[object] | None,

    expected: int,

) -> None:
    if values is not None and len(values) != expected:
        raise ValueError(f"{name} length {len(values)} does not match {expected}")


class RelevanceEngine:
    """Exact rich prediction plus explicitly separate compact screening paths."""

    __slots__ = (
        "_backend",
        "_indicator",
        "_ranker",
        "_screening_indicator",
        "_taxonomy",
        "_workers",
    )

    def __init__(

        self,

        taxonomy: Taxonomy | None = None,

        *,

        ranker: Ranker | None = None,

        backend: Literal["cpu", "cuda"] = "cpu",

        workers: int | None = None,

    ) -> None:
        resolved_taxonomy = taxonomy or DefaultTaxonomy()
        if not isinstance(resolved_taxonomy, Taxonomy):
            raise TypeError("taxonomy must be a Taxonomy or None")
        resolved_ranker = ranker or DefaultRanker()
        if not isinstance(resolved_ranker, Ranker):
            raise TypeError("ranker must inherit Ranker")
        self._taxonomy = resolved_taxonomy
        self._ranker = resolved_ranker
        self._workers = _workers(workers)
        self._backend = create_backend(backend)
        self._indicator = AcceleratedRelevanceIndicator(
            config=resolved_taxonomy.config,
            catalogs=resolved_taxonomy.catalogs,
        )
        self._screening_indicator = CompactRelevanceIndicator(
            config=resolved_taxonomy.config,
            catalogs=resolved_taxonomy.catalogs,
        )

    @classmethod
    def from_yaml(

        cls,

        path: str | Path,

        *,

        ranker: Ranker | None = None,

        backend: Literal["cpu", "cuda"] = "cpu",

        workers: int | None = None,

    ) -> "RelevanceEngine":
        return cls(
            Taxonomy.from_yaml(path),
            ranker=ranker,
            backend=backend,
            workers=workers,
        )

    @property
    def taxonomy(self) -> Taxonomy:
        return self._taxonomy

    @property
    def ranker(self) -> Ranker:
        return self._ranker

    @property
    def backend(self) -> str:
        return self._backend.name

    @property
    def backend_info(self) -> BackendInfo:
        return self._backend.describe(self._screening_indicator)

    @property
    def workers(self) -> int:
        return self._workers

    def predict(

        self,

        text: str,

        *,

        timestamp: TimestampLike,

        article_id: str | None = None,

        source: str = "unknown",

    ) -> Prediction:
        """Build one rich decision with evidence and explanation provenance."""

        article = Article(
            article_id=(
                _article_id(text, timestamp) if article_id is None else article_id
            ),
            timestamp=timestamp,
            text=text,
            source=source,
        )
        return self._indicator.predict(article)

    def predict_batch(

        self,

        texts: Sequence[str],

        *,

        timestamps: Sequence[TimestampLike],

        article_ids: Sequence[str] | None = None,

        sources: Sequence[str] | None = None,

        workers: int | None = None,

    ) -> tuple[Prediction, ...]:
        """Build rich decisions; use `screen_batch()` for throughput work."""

        count = len(texts)
        if len(timestamps) != count:
            raise ValueError("texts and timestamps must have equal lengths")
        _optional_length("article_ids", article_ids, count)
        _optional_length("sources", sources, count)
        articles = tuple(
            Article(
                article_id=(
                    article_ids[index] if article_ids is not None else f"row-{index}"
                ),
                timestamp=timestamps[index],
                text=text,
                source=sources[index] if sources is not None else "unknown",
            )
            for index, text in enumerate(texts)
        )
        return self._indicator.predict_batch(
            articles,
            workers=_workers(workers) if workers is not None else self._workers,
        )

    def screen_batch(

        self,

        texts: Sequence[str],

        *,

        timestamps: Sequence[TimestampLike],

        article_ids: Sequence[str] | None = None,

        sources: Sequence[str] | None = None,

        workers: int | None = None,

    ) -> PythonScreeningBatch:
        """Run compact list screening with implicit zero-observation rows."""

        count = len(texts)
        if len(timestamps) != count:
            raise ValueError("texts and timestamps must have equal lengths")
        _optional_length("article_ids", article_ids, count)
        _optional_length("sources", sources, count)
        ids = (
            article_ids
            if article_ids is not None
            else tuple(f"row-{index}" for index in range(count))
        )
        batch = ColumnArticleBatch(
            article_ids=ids,
            timestamps=timestamps,
            texts=texts,
            sources=sources,
        )
        return self._backend.screen_batch(
            self._screening_indicator,
            batch,
            workers=_workers(workers) if workers is not None else self._workers,
        )

    def screen_arrow(

        self,

        data,

        *,

        text_column: str = "text",

        timestamp_column: str = "timestamp",

        workers: int | None = None,

        row_count: int | None = None,

    ) -> ArrowCompactFacts:
        """Run the sparse Arrow path without materializing per-row Python objects."""

        if row_count is not None and (
            isinstance(row_count, bool)
            or not isinstance(row_count, int)
            or row_count < 0
        ):
            raise ValueError("row_count must be a non-negative integer or None")
        discovered = getattr(data, "num_rows", None)
        if discovered is not None:
            discovered = int(discovered)
            if row_count is None:
                row_count = discovered
            elif row_count != discovered:
                raise ValueError(
                    f"row_count {row_count} does not match Arrow rows {discovered}"
                )

        return self._backend.screen_arrow(
            self._screening_indicator,
            self._taxonomy,
            data,
            text_column=text_column,
            timestamp_column=timestamp_column,
            workers=_workers(workers) if workers is not None else self._workers,
            row_count=row_count,
        )

    def explain(

        self,

        prediction: RelevanceDecision,

        *,

        format: Literal["text", "mermaid"] = "text",

        detail: Literal["summary", "details"] = "summary",

    ) -> str:
        """Format already-recorded evidence; this does not rerun prediction."""

        if not isinstance(prediction, RelevanceDecision):
            raise TypeError("prediction must be a rich Prediction")
        if detail not in {"summary", "details"}:
            raise ValueError("detail must be 'summary' or 'details'")
        return self._indicator.explain(
            prediction,
            output_format=format,
            detail=detail,
        )

    def rank(

        self,

        predictions: Iterable[Prediction],

        *,

        exposure_family: ExposureFamily | str,

        decision_bucket: datetime | str,

        limit: int | None = None,

    ) -> tuple[RankedPrediction, ...]:
        return self._ranker.rank(
            predictions,
            config=self._taxonomy.config,
            exposure_family=exposure_family,
            decision_bucket=decision_bucket,
            limit=limit,
        )


@lru_cache(maxsize=1)
def _default_engine() -> RelevanceEngine:
    return RelevanceEngine()


def predict(

    text: str,

    *,

    timestamp: TimestampLike,

    article_id: str | None = None,

    source: str = "unknown",

) -> Prediction:
    return _default_engine().predict(
        text,
        timestamp=timestamp,
        article_id=article_id,
        source=source,
    )


def predict_batch(

    texts: Sequence[str],

    *,

    timestamps: Sequence[TimestampLike],

    article_ids: Sequence[str] | None = None,

    sources: Sequence[str] | None = None,

    workers: int | None = None,

) -> tuple[Prediction, ...]:
    return _default_engine().predict_batch(
        texts,
        timestamps=timestamps,
        article_ids=article_ids,
        sources=sources,
        workers=workers,
    )


def screen_batch(

    texts: Sequence[str],

    *,

    timestamps: Sequence[TimestampLike],

    article_ids: Sequence[str] | None = None,

    sources: Sequence[str] | None = None,

    workers: int | None = None,

) -> PythonScreeningBatch:
    return _default_engine().screen_batch(
        texts,
        timestamps=timestamps,
        article_ids=article_ids,
        sources=sources,
        workers=workers,
    )


def screen_arrow(

    data,

    *,

    text_column: str = "text",

    timestamp_column: str = "timestamp",

    workers: int | None = None,

    row_count: int | None = None,

) -> ArrowCompactFacts:
    return _default_engine().screen_arrow(
        data,
        text_column=text_column,
        timestamp_column=timestamp_column,
        workers=workers,
        row_count=row_count,
    )


__all__ = [
    "RelevanceEngine",
    "available_backends",
    "predict",
    "predict_batch",
    "screen_arrow",
    "screen_batch",
]