mirror of https://github.com/razor-ai/soup.git
369 lines
13 KiB
Python
369 lines
13 KiB
Python
"""Classifier / reranker / cross_encoder trainer (v0.53.2 #132).
|
|
|
|
Wraps :class:`transformers.AutoModelForSequenceClassification` for the three
|
|
classifier-family tasks declared in v0.52.0 Part B:
|
|
|
|
* ``classifier`` — single-input sequence classification.
|
|
* ``reranker`` — single-input reranker (typically score head); shares the
|
|
classifier head topology.
|
|
* ``cross_encoder`` — paired-input scoring (passage / query, NLI, etc.).
|
|
|
|
Single-label uses ``CrossEntropyLoss``; multi-label uses ``BCEWithLogitsLoss``
|
|
(automatically when ``training.classifier_kind='multi_label'``).
|
|
|
|
Data format expected per row:
|
|
|
|
* ``text`` (or ``messages`` joined into a string) — single-input tasks.
|
|
* ``text_a`` + ``text_b`` — paired-input ``cross_encoder``.
|
|
* ``label`` — int (single-label), list[int] (multi-label), or string in
|
|
``training.label_names``.
|
|
|
|
Mirrors the BCO / pretrain wrapper pattern. Lazy imports for heavy deps;
|
|
trust_remote_code threaded through the v0.36.0 resolver.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, List, Union
|
|
|
|
from rich.console import Console
|
|
|
|
from soup_cli.config.schema import SoupConfig
|
|
|
|
console = Console()
|
|
|
|
# Cap on multi-label list entries — defense against malformed dataset rows
|
|
# (security review v0.53.2 H2). Matches v0.52.0 ``_MAX_LABELS=1024``.
|
|
_MAX_MULTI_LABEL_ENTRIES: int = 1024
|
|
|
|
|
|
def _row_to_text(row: dict) -> str:
|
|
"""Extract the single-input text from a row (``text`` or joined messages).
|
|
|
|
Raises:
|
|
TypeError: non-string ``content`` inside a messages list — silent skip
|
|
could poison training data with empty strings (security review
|
|
v0.53.2 M3).
|
|
ValueError: neither ``text`` nor ``messages`` present.
|
|
"""
|
|
if "text" in row and isinstance(row["text"], str):
|
|
return row["text"]
|
|
msgs = row.get("messages")
|
|
if isinstance(msgs, list):
|
|
parts: list[str] = []
|
|
for msg in msgs:
|
|
if not isinstance(msg, dict):
|
|
# Non-dict entries are silently skipped — caller's loader
|
|
# would normally have produced these; keep loud-fail at the
|
|
# row-level (missing text) rather than per-message.
|
|
continue
|
|
content = msg.get("content", "")
|
|
if not isinstance(content, str):
|
|
raise TypeError(
|
|
"Classifier row messages[i]['content'] must be str, got "
|
|
f"{type(content).__name__!r}"
|
|
)
|
|
parts.append(content)
|
|
return "\n".join(parts)
|
|
raise ValueError(
|
|
"Classifier row missing 'text' field and no joinable 'messages' list. "
|
|
f"Row keys: {sorted(row)!r}"
|
|
)
|
|
|
|
|
|
def _row_to_pair(row: dict) -> tuple[str, str]:
|
|
"""Extract (text_a, text_b) for paired ``cross_encoder`` rows.
|
|
|
|
Raises:
|
|
TypeError: either field present but not a string (security review M4 —
|
|
silent ``str()`` coercion of dicts/lists produced garbage training
|
|
text).
|
|
ValueError: neither pair of fields present.
|
|
"""
|
|
if "text_a" in row and "text_b" in row:
|
|
a, b = row["text_a"], row["text_b"]
|
|
if not isinstance(a, str) or not isinstance(b, str):
|
|
raise TypeError(
|
|
"cross_encoder rows require 'text_a' and 'text_b' to be str; "
|
|
f"got text_a={type(a).__name__}, text_b={type(b).__name__}"
|
|
)
|
|
return a, b
|
|
if "question" in row and "answer" in row:
|
|
q, ans = row["question"], row["answer"]
|
|
if not isinstance(q, str) or not isinstance(ans, str):
|
|
raise TypeError(
|
|
"cross_encoder rows require 'question' and 'answer' to be "
|
|
f"str; got question={type(q).__name__}, answer={type(ans).__name__}"
|
|
)
|
|
return q, ans
|
|
raise ValueError(
|
|
"cross_encoder row requires 'text_a' + 'text_b' (or 'question' + "
|
|
f"'answer'). Row keys: {sorted(row)!r}"
|
|
)
|
|
|
|
|
|
def _normalise_label(
|
|
raw: object,
|
|
label_names: List[str] | None,
|
|
num_labels: int,
|
|
multi_label: bool,
|
|
) -> Union[int, list[float]]:
|
|
"""Convert a raw label (int / str / list) to the trainer-expected form.
|
|
|
|
Single-label → int in [0, num_labels). Multi-label → list[float] of
|
|
length ``num_labels``.
|
|
|
|
Raises:
|
|
ValueError: invalid index, oversize multi-label list (defense-in-depth
|
|
against malformed datasets — security review H2).
|
|
TypeError: bool / unsupported scalar type.
|
|
"""
|
|
if multi_label:
|
|
if isinstance(raw, list):
|
|
if len(raw) > _MAX_MULTI_LABEL_ENTRIES:
|
|
raise ValueError(
|
|
f"multi-label list too long: {len(raw)} entries "
|
|
f"(max {_MAX_MULTI_LABEL_ENTRIES})"
|
|
)
|
|
vec = [0.0] * num_labels
|
|
for entry in raw:
|
|
idx = _label_index(entry, label_names, num_labels)
|
|
vec[idx] = 1.0
|
|
return vec
|
|
# Single label silently broadcast to one-hot multi-label.
|
|
idx = _label_index(raw, label_names, num_labels)
|
|
vec = [0.0] * num_labels
|
|
vec[idx] = 1.0
|
|
return vec
|
|
return _label_index(raw, label_names, num_labels)
|
|
|
|
|
|
def _label_index(
|
|
raw: object, label_names: List[str] | None, num_labels: int
|
|
) -> int:
|
|
if isinstance(raw, bool):
|
|
# Project policy (v0.30.0 Candidate / v0.39.0 ReLoRAPolicy / v0.41.0
|
|
# Part B): bool-as-int violations raise TypeError, not ValueError.
|
|
raise TypeError(f"label must not be bool, got {raw!r}")
|
|
if isinstance(raw, int):
|
|
if raw < 0 or raw >= num_labels:
|
|
raise ValueError(
|
|
f"label index {raw} out of range [0, {num_labels})"
|
|
)
|
|
return raw
|
|
if isinstance(raw, str):
|
|
if label_names is None:
|
|
raise ValueError(
|
|
f"label is str {raw!r} but training.label_names is unset"
|
|
)
|
|
try:
|
|
return label_names.index(raw)
|
|
except ValueError as exc:
|
|
raise ValueError(
|
|
f"label {raw!r} not in training.label_names={label_names!r}"
|
|
) from exc
|
|
raise TypeError(
|
|
f"label must be int / str / list, got {type(raw).__name__}"
|
|
)
|
|
|
|
|
|
class ClassifierTrainerWrapper:
|
|
"""High-level wrapper for classifier / reranker / cross_encoder training."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: SoupConfig,
|
|
device: str = "cuda",
|
|
report_to: str = "none",
|
|
deepspeed_config: str | None = None,
|
|
fsdp_config: dict | None = None,
|
|
trust_remote_code: bool = False,
|
|
) -> None:
|
|
self.config = config
|
|
self.device = device
|
|
self.report_to = report_to
|
|
self.deepspeed_config = deepspeed_config
|
|
self.fsdp_config = fsdp_config
|
|
self.trust_remote_code = trust_remote_code
|
|
|
|
from soup_cli.utils.trust_remote import (
|
|
model_requires_trust_remote_code,
|
|
resolve_trust_remote_code,
|
|
)
|
|
|
|
requires = model_requires_trust_remote_code(config.base) or False
|
|
self._trust_remote_code = resolve_trust_remote_code(
|
|
config.base,
|
|
requested=trust_remote_code,
|
|
console=console,
|
|
requires_remote_code=requires,
|
|
)
|
|
|
|
self.model: Any = None
|
|
self.tokenizer: Any = None
|
|
self.trainer: Any = None
|
|
self._output_dir: str | None = None
|
|
|
|
def setup(self, dataset: dict) -> None:
|
|
"""Load model + tokenizer, tokenise dataset, build HF Trainer."""
|
|
from datasets import Dataset
|
|
from transformers import (
|
|
AutoModelForSequenceClassification,
|
|
AutoTokenizer,
|
|
Trainer,
|
|
TrainingArguments,
|
|
)
|
|
|
|
cfg = self.config
|
|
tcfg = cfg.training
|
|
|
|
if tcfg.num_labels is None:
|
|
raise ValueError(
|
|
f"task={cfg.task!r} requires training.num_labels to be set"
|
|
)
|
|
num_labels = int(tcfg.num_labels)
|
|
multi_label = (tcfg.classifier_kind == "multi_label")
|
|
problem_type = (
|
|
"multi_label_classification" if multi_label else "single_label_classification"
|
|
)
|
|
|
|
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
cfg.base, trust_remote_code=self._trust_remote_code
|
|
)
|
|
if self.tokenizer.pad_token is None and self.tokenizer.eos_token is not None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
console.print(f"[dim]Loading classifier model: {cfg.base}[/]")
|
|
self.model = AutoModelForSequenceClassification.from_pretrained(
|
|
cfg.base,
|
|
num_labels=num_labels,
|
|
problem_type=problem_type,
|
|
trust_remote_code=self._trust_remote_code,
|
|
)
|
|
|
|
is_paired = (cfg.task == "cross_encoder")
|
|
label_names = (
|
|
list(tcfg.label_names) if tcfg.label_names is not None else None
|
|
)
|
|
|
|
def encode(row: dict) -> dict:
|
|
if is_paired:
|
|
a, b = _row_to_pair(row)
|
|
enc = self.tokenizer(
|
|
a, b,
|
|
truncation=True,
|
|
max_length=cfg.data.max_length,
|
|
)
|
|
else:
|
|
text = _row_to_text(row)
|
|
enc = self.tokenizer(
|
|
text,
|
|
truncation=True,
|
|
max_length=cfg.data.max_length,
|
|
)
|
|
label = _normalise_label(
|
|
row.get("label"), label_names, num_labels, multi_label
|
|
)
|
|
enc["labels"] = label
|
|
return enc
|
|
|
|
raw_train = Dataset.from_list(dataset["train"])
|
|
train_ds = raw_train.map(encode, remove_columns=raw_train.column_names)
|
|
eval_ds = None
|
|
if "val" in dataset and dataset["val"]:
|
|
raw_val = Dataset.from_list(dataset["val"])
|
|
eval_ds = raw_val.map(encode, remove_columns=raw_val.column_names)
|
|
|
|
output_dir = Path(cfg.output)
|
|
if cfg.experiment_name:
|
|
output_dir = output_dir / cfg.experiment_name
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
batch_size = tcfg.batch_size if tcfg.batch_size != "auto" else 8
|
|
total_steps = (
|
|
math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps)
|
|
* tcfg.epochs
|
|
)
|
|
warmup_steps = int(total_steps * tcfg.warmup_ratio)
|
|
|
|
args = TrainingArguments(
|
|
output_dir=str(output_dir),
|
|
num_train_epochs=tcfg.epochs,
|
|
per_device_train_batch_size=batch_size,
|
|
gradient_accumulation_steps=tcfg.gradient_accumulation_steps,
|
|
learning_rate=tcfg.lr,
|
|
warmup_steps=warmup_steps,
|
|
weight_decay=tcfg.weight_decay,
|
|
max_grad_norm=tcfg.max_grad_norm,
|
|
optim=tcfg.optimizer,
|
|
lr_scheduler_type=tcfg.scheduler,
|
|
logging_steps=tcfg.logging_steps,
|
|
save_steps=tcfg.save_steps,
|
|
save_total_limit=3,
|
|
bf16=self.device == "cuda",
|
|
report_to=self.report_to,
|
|
deepspeed=self.deepspeed_config,
|
|
**(self.fsdp_config or {}),
|
|
)
|
|
|
|
from transformers import DataCollatorWithPadding
|
|
|
|
self.trainer = Trainer(
|
|
model=self.model,
|
|
args=args,
|
|
train_dataset=train_ds,
|
|
eval_dataset=eval_ds,
|
|
tokenizer=self.tokenizer,
|
|
data_collator=DataCollatorWithPadding(tokenizer=self.tokenizer),
|
|
)
|
|
self._output_dir = str(output_dir)
|
|
|
|
def train(
|
|
self,
|
|
display: object | None = None,
|
|
tracker: object | None = None,
|
|
run_id: str = "",
|
|
resume_from_checkpoint: str | None = None,
|
|
) -> dict:
|
|
if self.trainer is None:
|
|
raise RuntimeError(
|
|
"ClassifierTrainerWrapper.train() called before setup(). "
|
|
"Call setup(dataset) first."
|
|
)
|
|
start = time.time()
|
|
if display is not None:
|
|
from soup_cli.monitoring.callback import SoupTrainerCallback
|
|
|
|
self.trainer.add_callback(
|
|
SoupTrainerCallback(
|
|
display, tracker=tracker, run_id=run_id,
|
|
loss_watchdog=self.config.training.loss_watchdog,
|
|
loss_watchdog_threshold=self.config.training.loss_watchdog_threshold,
|
|
loss_watchdog_patience=self.config.training.loss_watchdog_patience,
|
|
)
|
|
)
|
|
self.trainer.train(resume_from_checkpoint=resume_from_checkpoint)
|
|
duration = time.time() - start
|
|
|
|
self.trainer.save_model(self._output_dir)
|
|
self.tokenizer.save_pretrained(self._output_dir)
|
|
|
|
logs = self.trainer.state.log_history
|
|
train_losses = [entry["loss"] for entry in logs if "loss" in entry]
|
|
|
|
hours = int(duration // 3600)
|
|
minutes = int((duration % 3600) // 60)
|
|
duration_str = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m"
|
|
return {
|
|
"initial_loss": train_losses[0] if train_losses else 0,
|
|
"final_loss": train_losses[-1] if train_losses else 0,
|
|
"duration": duration_str,
|
|
"duration_secs": duration,
|
|
"output_dir": self._output_dir,
|
|
"total_steps": self.trainer.state.global_step,
|
|
}
|