src/feature_extraction.py

"""TF-IDF featurizer for Traditional ML (RF / SVM / NB).

Single class ``TfidfFeaturizer`` wrapping ``sklearn.feature_extraction.text.TfidfVectorizer``.
Defaults: `ngram_range=(1,2)`, `min_df=2`, `max_df=0.95`, `sublinear_tf=True`.
Persists via joblib.

Other featurizers live elsewhere:
- FastText → ``src/fasttext_featurizer.py`` (Hybrid DL)
- Transformer tokenizers → owned by each model class in ``src/models/``
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Sequence

import joblib
from sklearn.feature_extraction.text import TfidfVectorizer

logger = logging.getLogger(__name__)


class TfidfFeaturizer:
    """TF-IDF featurizer for Traditional ML (RF, NB, SVM)."""

    def __init__(
        self,
        max_features: int | None = None,
        ngram_range: tuple[int, int] = (1, 2),
        min_df: int = 2,
        max_df: float = 0.95,
        sublinear_tf: bool = True,
    ) -> None:
        self.vectorizer = TfidfVectorizer(
            max_features=max_features,
            ngram_range=ngram_range,
            min_df=min_df,
            max_df=max_df,
            sublinear_tf=sublinear_tf,
        )

    def fit(self, texts: Sequence[str]) -> "TfidfFeaturizer":
        """Fit the vocabulary + IDF weights on ``texts``."""
        self.vectorizer.fit(list(texts))
        logger.info("TfidfFeaturizer fitted: vocab_size=%d", self.vocab_size)
        return self

    def transform(self, texts: Sequence[str]):
        """Transform ``texts`` into a sparse TF-IDF matrix ``(n, vocab_size)``."""
        return self.vectorizer.transform(list(texts))

    def fit_transform(self, texts: Sequence[str]):
        """Fit then transform in one step."""
        out = self.vectorizer.fit_transform(list(texts))
        logger.info("TfidfFeaturizer fitted: vocab_size=%d", self.vocab_size)
        return out

    @property
    def vocab_size(self) -> int:
        """Number of features learned during fit."""
        vocab = getattr(self.vectorizer, "vocabulary_", None)
        return len(vocab) if vocab is not None else 0

    def save(self, path: Path) -> None:
        """Persist the fitted vectorizer to ``path`` via joblib."""
        joblib.dump(self.vectorizer, str(path))

    @classmethod
    def load(cls, path: Path) -> "TfidfFeaturizer":
        """Load a previously saved vectorizer and return a fresh wrapper."""
        instance = cls.__new__(cls)
        instance.vectorizer = joblib.load(str(path))
        return instance