training/SFT/__init__.py

"""
SFT (Supervised Fine-Tuning) module for LKAlert dual-head training.

Dual-head: HazardHead (binary ego-threat) + TTAHead (time-to-accident regression).
Manifest-based dataset, initialized from pretrain_v2 stage_b checkpoint.
"""

from .dataset import (
    SFTDataset,
    sft_collate_fn,
    TTASample,
    VideoInfo,
)

from .trainer import (
    SFTModel,
    SFTTrainer,
    HazardHead,
    TTAHead,
    BeliefAggregator,
    compute_sft_loss,
    compute_calibration_error,
)

__all__ = [
    # Dataset
    "SFTDataset",
    "sft_collate_fn",
    "TTASample",
    "VideoInfo",
    # Trainer
    "SFTModel",
    "SFTTrainer",
    "HazardHead",
    "TTAHead",
    "BeliefAggregator",
    "compute_sft_loss",
    "compute_calibration_error",
]