training/base_trainer.py

"""Base trainer abstraction for all model training pipelines.

Enforces: split discipline, deterministic seed, artifact layout,
metrics schema, and model card generation. All 7 model trainers
inherit from BaseTrainer and implement the abstract methods.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import joblib
import json
import logging
import pandas as pd

from app.core.config import settings
from app.core.exceptions import DatasetError, TrainingError
from app.data.loader import DatasetLoader

logger = logging.getLogger(__name__)


@dataclass
class TrainingResult:
    """Result of a single model training run."""

    model_name: str
    model_version: str
    metrics: dict
    artifact_dir: Path
    trained_at: datetime
    split_counts: dict[str, int]


class BaseTrainer(ABC):
    """Abstract base for all model training pipelines.

    Enforces: split discipline, deterministic seed, artifact layout,
    metrics schema, and model card generation.
    """

    def __init__(self, dataset_dir: str | Path, artifact_base_dir: str | Path) -> None:
        self._loader = DatasetLoader(dataset_dir)
        self._artifact_base_dir = Path(artifact_base_dir)
        self._seed = settings.seed

    @property
    @abstractmethod
    def model_name(self) -> str:
        """Unique model identifier, e.g., 'lo_tagger', 'bloom_classifier'."""
        ...

    @property
    @abstractmethod
    def model_version(self) -> str:
        """Version string, e.g., 'lo_tagger_v2_baseline_001'."""
        ...

    @property
    @abstractmethod
    def table_name(self) -> str:
        """CSV table to load, e.g., 'training_lo_tagging'."""
        ...

    @property
    def artifact_dir(self) -> Path:
        """Resolved artifact directory for this model."""
        return self._artifact_base_dir / self.model_name

    def run(self) -> TrainingResult:
        """Execute the full training pipeline. Returns TrainingResult."""
        logger.info("Starting training for model '%s'", self.model_name)

        df = self._loader.load_table(self.table_name)
        train_df, val_df, test_df = self._split(df)
        self._validate_train_not_empty(train_df)

        logger.info(
            "Split counts — train: %d, validation: %d, test: %d",
            len(train_df), len(val_df), len(test_df),
        )

        model_artifacts = self.train(train_df, val_df)
        val_metrics = self.evaluate(model_artifacts, val_df, split_name="validation")
        test_metrics = self.evaluate(model_artifacts, test_df, split_name="test")

        metrics = self._build_metrics(val_metrics, test_metrics, train_df, val_df, test_df)
        self._check_baseline(metrics)

        self._save_artifacts(model_artifacts, metrics, train_df, val_df, test_df)

        trained_at = datetime.now(timezone.utc)
        split_counts = {
            "train": len(train_df),
            "validation": len(val_df),
            "test": len(test_df),
        }

        logger.info("Training complete for model '%s'", self.model_name)

        return TrainingResult(
            model_name=self.model_name,
            model_version=self.model_version,
            metrics=metrics,
            artifact_dir=self.artifact_dir,
            trained_at=trained_at,
            split_counts=split_counts,
        )

    def _split(self, df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """Filter by train_split column. Never shuffles or re-samples."""
        if "train_split" not in df.columns:
            raise DatasetError(
                f"Table '{self.table_name}' missing 'train_split' column"
            )
        train_df = df[df["train_split"] == "train"].copy()
        val_df = df[df["train_split"] == "validation"].copy()
        test_df = df[df["train_split"] == "test"].copy()
        return train_df, val_df, test_df

    def _validate_train_not_empty(self, train_df: pd.DataFrame) -> None:
        """Raise DatasetError if no training rows available."""
        if len(train_df) == 0:
            raise DatasetError(
                f"No rows with train_split='train' in table '{self.table_name}'. "
                f"Cannot proceed with empty training data."
            )

    @abstractmethod
    def train(self, train_df: pd.DataFrame, val_df: pd.DataFrame) -> dict:
        """Train the model. Return dict of artifacts (model, vectorizer, etc.)."""
        ...

    @abstractmethod
    def evaluate(self, artifacts: dict, df: pd.DataFrame, split_name: str) -> dict:
        """Evaluate model on a split. Return metrics dict."""
        ...

    @abstractmethod
    def _build_metrics(
        self,
        val_metrics: dict,
        test_metrics: dict,
        train_df: pd.DataFrame,
        val_df: pd.DataFrame,
        test_df: pd.DataFrame,
    ) -> dict:
        """Assemble the full metrics.json content."""
        ...

    @abstractmethod
    def _check_baseline(self, metrics: dict) -> None:
        """Verify metrics exceed random baseline. Raise TrainingError on failure."""
        ...

    @abstractmethod
    def _build_training_config(
        self,
        train_df: pd.DataFrame,
        val_df: pd.DataFrame,
        test_df: pd.DataFrame,
    ) -> dict:
        """Build training_config.json content with hyperparameters."""
        ...

    @abstractmethod
    def _build_model_card(self, metrics: dict) -> str:
        """Generate model_card.md content."""
        ...

    def _save_artifacts(
        self,
        model_artifacts: dict,
        metrics: dict,
        train_df: pd.DataFrame,
        val_df: pd.DataFrame,
        test_df: pd.DataFrame,
    ) -> None:
        """Save all artifacts to the model's artifact directory."""
        self.artifact_dir.mkdir(parents=True, exist_ok=True)

        # Save model objects
        for name, obj in model_artifacts.items():
            if name.endswith(".json"):
                (self.artifact_dir / name).write_text(
                    json.dumps(obj, indent=2, default=str), encoding="utf-8"
                )
            else:
                joblib.dump(obj, self.artifact_dir / f"{name}.joblib", compress=3)

        # Save metrics.json
        (self.artifact_dir / "metrics.json").write_text(
            json.dumps(metrics, indent=2, default=str), encoding="utf-8"
        )

        # Save training_config.json
        config = self._build_training_config(train_df, val_df, test_df)
        (self.artifact_dir / "training_config.json").write_text(
            json.dumps(config, indent=2, default=str), encoding="utf-8"
        )

        # Save model_card.md
        card = self._build_model_card(metrics)
        (self.artifact_dir / "model_card.md").write_text(card, encoding="utf-8")

        logger.info("Artifacts saved to %s", self.artifact_dir)