mirror of https://github.com/razor-ai/soup.git
144 lines
4.2 KiB
Python
144 lines
4.2 KiB
Python
"""HuggingFace Trainer callback that feeds metrics to our display and tracker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from transformers import (
|
|
TrainerCallback,
|
|
TrainerControl,
|
|
TrainerState,
|
|
TrainingArguments,
|
|
)
|
|
|
|
from soup_cli.monitoring.display import TrainingDisplay
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SoupTrainerCallback(TrainerCallback):
|
|
"""Bridges HF Trainer events to Soup's Rich live display and experiment tracker."""
|
|
|
|
def __init__(
|
|
self,
|
|
display: TrainingDisplay,
|
|
tracker: Optional[object] = None,
|
|
run_id: str = "",
|
|
eval_config: Optional[object] = None,
|
|
output_dir: str = "",
|
|
):
|
|
self.display = display
|
|
self.tracker = tracker
|
|
self.run_id = run_id
|
|
self.eval_config = eval_config
|
|
self.output_dir = output_dir
|
|
|
|
def on_train_begin(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, **kwargs,
|
|
):
|
|
self.display.start(total_steps=state.max_steps)
|
|
|
|
def on_log(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, logs=None, **kwargs,
|
|
):
|
|
if logs is None:
|
|
return
|
|
|
|
# Try to get GPU memory
|
|
gpu_mem = ""
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
used = torch.cuda.memory_allocated() / (1024**3)
|
|
total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
|
|
gpu_mem = f"{used:.1f}/{total:.1f} GB"
|
|
except Exception:
|
|
pass
|
|
|
|
step = state.global_step
|
|
epoch = state.epoch or 0
|
|
loss = logs.get("loss", 0.0)
|
|
lr = logs.get("learning_rate", 0.0)
|
|
grad_norm = logs.get("grad_norm", 0.0)
|
|
speed = logs.get("train_steps_per_second", 0.0)
|
|
|
|
self.display.update(
|
|
step=step,
|
|
epoch=epoch,
|
|
loss=loss,
|
|
lr=lr,
|
|
grad_norm=grad_norm,
|
|
speed=speed,
|
|
gpu_mem=gpu_mem,
|
|
)
|
|
|
|
# Log to experiment tracker
|
|
if self.tracker and self.run_id:
|
|
self.tracker.log_metrics(
|
|
run_id=self.run_id,
|
|
step=step,
|
|
epoch=epoch,
|
|
loss=loss,
|
|
lr=lr,
|
|
grad_norm=grad_norm,
|
|
speed=speed,
|
|
gpu_mem=gpu_mem,
|
|
)
|
|
|
|
def on_train_end(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, **kwargs,
|
|
):
|
|
self.display.stop()
|
|
self._run_auto_eval()
|
|
|
|
def _run_auto_eval(self) -> None:
|
|
"""Run automatic evaluation if configured."""
|
|
if self.eval_config is None:
|
|
return
|
|
if not getattr(self.eval_config, "auto_eval", False):
|
|
return
|
|
if not self.output_dir:
|
|
return
|
|
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
console.print("\n[bold blue]Running auto-eval...[/]")
|
|
|
|
benchmarks = getattr(self.eval_config, "benchmarks", None) or []
|
|
custom_tasks = getattr(self.eval_config, "custom_tasks", None)
|
|
|
|
# Run standard benchmarks
|
|
if benchmarks:
|
|
try:
|
|
from soup_cli.commands.eval import benchmark
|
|
benchmark(
|
|
model=self.output_dir,
|
|
benchmarks=",".join(benchmarks),
|
|
num_fewshot=None,
|
|
batch_size=8,
|
|
run_id=self.run_id,
|
|
device=None,
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("Auto-eval benchmark failed")
|
|
console.print(f"[yellow]Auto-eval benchmark failed: {exc}[/]")
|
|
|
|
# Run custom eval
|
|
if custom_tasks:
|
|
try:
|
|
from soup_cli.commands.eval import custom
|
|
custom(
|
|
tasks=custom_tasks,
|
|
model=self.output_dir,
|
|
run_id=self.run_id,
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("Auto-eval custom failed")
|
|
console.print(f"[yellow]Auto-eval custom failed: {exc}[/]")
|