stonks-oracle/services/intelligence_pipeline_v3/sentiment/finbert_adapter.py

"""FinBERT adapter for financial sentiment classification.

Provides a unified interface for FinBERT inference with:
- Production mode: loads ProsusAI/finbert and runs real inference
- Test mode: deterministic keyword-based mock probabilities

Model version is pinned and exposed for lineage tracking.
"""

from __future__ import annotations

import logging

logger = logging.getLogger(__name__)

# Pinned model configuration
FINBERT_MODEL_NAME = "ProsusAI/finbert"
FINBERT_MODEL_VERSION = "ProsusAI/finbert@v1.0"

# Keywords for deterministic test mode
_POSITIVE_KEYWORDS = frozenset({
    "growth", "profit", "beat", "raised", "upgrade", "strong",
    "surge", "gain", "bullish", "outperform", "exceeded", "record",
    "positive", "optimistic", "rally", "upside",
})
_NEGATIVE_KEYWORDS = frozenset({
    "loss", "decline", "miss", "cut", "downgrade", "weak",
    "plunge", "drop", "bearish", "underperform", "fell", "crash",
    "negative", "pessimistic", "risk", "downside", "slump",
})


class FinBERTAdapter:
    """Adapter for FinBERT financial sentiment classification.

    Parameters
    ----------
    test_mode
        When True, uses deterministic keyword-based classification
        instead of loading the actual FinBERT model. Useful for testing
        without GPU/large model dependencies.
    """

    def __init__(self, test_mode: bool = True) -> None:
        self._test_mode = test_mode
        self._model = None
        self._tokenizer = None
        self._model_name = FINBERT_MODEL_NAME
        self._model_version = FINBERT_MODEL_VERSION

        if not test_mode:
            self._load_model()

    @property
    def model_version(self) -> str:
        """Return the pinned model version string."""
        return self._model_version

    @property
    def model_name(self) -> str:
        """Return the model name."""
        return self._model_name

    def _load_model(self) -> None:
        """Load the FinBERT model and tokenizer for production inference."""
        try:
            from transformers import AutoModelForSequenceClassification, AutoTokenizer

            logger.info("Loading FinBERT model: %s", self._model_name)
            self._tokenizer = AutoTokenizer.from_pretrained(self._model_name)
            self._model = AutoModelForSequenceClassification.from_pretrained(self._model_name)
            self._model.eval()
            logger.info("FinBERT model loaded successfully")
        except ImportError:
            raise RuntimeError(
                "transformers and torch are required for production FinBERT inference. "
                "Install with: pip install transformers torch"
            )
        except Exception as e:
            raise RuntimeError(f"Failed to load FinBERT model: {e}") from e

    def classify(self, texts: list[str]) -> list[tuple[float, float, float]]:
        """Classify texts and return probability distributions.

        Parameters
        ----------
        texts
            List of text strings to classify.

        Returns
        -------
        list[tuple[float, float, float]]
            List of (positive_prob, negative_prob, neutral_prob) tuples.
            Probabilities sum to 1.0 for each text.
        """
        if not texts:
            return []

        if self._test_mode:
            return self._classify_test_mode(texts)

        return self._classify_production(texts)

    def _classify_test_mode(self, texts: list[str]) -> list[tuple[float, float, float]]:
        """Deterministic keyword-based classification for testing.

        Returns consistent probabilities based on keyword presence:
        - Positive keywords dominant -> (0.75, 0.10, 0.15)
        - Negative keywords dominant -> (0.10, 0.75, 0.15)
        - Both present (mixed signals) -> (0.40, 0.40, 0.20)
        - Neither present -> (0.15, 0.15, 0.70)
        """
        results: list[tuple[float, float, float]] = []

        for text in texts:
            lower_text = text.lower()
            words = set(lower_text.split())

            has_positive = bool(words & _POSITIVE_KEYWORDS)
            has_negative = bool(words & _NEGATIVE_KEYWORDS)

            if has_positive and has_negative:
                # Mixed signals
                results.append((0.40, 0.40, 0.20))
            elif has_positive:
                results.append((0.75, 0.10, 0.15))
            elif has_negative:
                results.append((0.10, 0.75, 0.15))
            else:
                # Neutral — no sentiment keywords
                results.append((0.15, 0.15, 0.70))

        return results

    def _classify_production(self, texts: list[str]) -> list[tuple[float, float, float]]:
        """Run FinBERT inference on texts using the loaded model."""
        import torch

        if self._model is None or self._tokenizer is None:
            raise RuntimeError("Model not loaded. Initialize with test_mode=False.")

        results: list[tuple[float, float, float]] = []

        # Process in batches to manage memory
        batch_size = 16
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            inputs = self._tokenizer(
                batch,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt",
            )

            with torch.no_grad():
                outputs = self._model(**inputs)
                # FinBERT output order: positive, negative, neutral
                probs = torch.softmax(outputs.logits, dim=-1)

            for prob in probs:
                pos = float(prob[0])
                neg = float(prob[1])
                neu = float(prob[2])
                results.append((pos, neg, neu))

        return results