File size: 750 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """
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",
]
|