src/models/random_forest.py

"""Random Forest classifier - single-stage binary cyberbullying detection."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Optional

import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier

logger = logging.getLogger(__name__)


class RandomForestModel:
    """Random Forest wrapper with sklearn-style fit/predict/predict_proba/save/load."""

    # Target hyperparams to sweep during tuning:
    #   n_estimators ∈ [100, 200, 500]
    def __init__(
        self,
        n_estimators: int = 200,
        max_depth: Optional[int] = None,
        min_samples_split: int = 2,
        class_weight: str = "balanced",
        random_state: int = 42,
        n_jobs: int = -1,
    ) -> None:
        self.model = RandomForestClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            class_weight=class_weight,
            random_state=random_state,
            n_jobs=n_jobs,
        )

    def fit(self, X_train, y_train, X_val=None, y_val=None) -> None:
        """Fit on TF-IDF features. ``X_val``/``y_val`` accepted for API parity but unused."""
        self.model.fit(X_train, y_train)
        logger.info("RandomForest fitted on %d samples", X_train.shape[0])

    def predict(self, X) -> np.ndarray:
        """Predict 0/1 labels."""
        return self.model.predict(X)

    def predict_proba(self, X) -> np.ndarray:
        """Predict class probabilities; shape ``(n, 2)``."""
        return self.model.predict_proba(X)

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

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