Spaces:
Running on Zero
Running on Zero
| """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, | |
| ) | |
| 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, | |
| ) | |
| def taxonomy(self) -> Taxonomy: | |
| return self._taxonomy | |
| def ranker(self) -> Ranker: | |
| return self._ranker | |
| def backend(self) -> str: | |
| return self._backend.name | |
| def backend_info(self) -> BackendInfo: | |
| return self._backend.describe(self._screening_indicator) | |
| 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, | |
| ) | |
| 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", | |
| ] | |