"""Attach LaunchHelm to a Hugging Face Trainer.

Requires ``pip install "launchhelm[hf]"`` plus the LAUNCHHELM_* environment.
The callback opens the run at train start, logs at each ``logging_steps``
boundary, applies dashboard commands at ``on_step_end`` (the collective
safe point), uploads checkpoint saves, and finishes at train end.

    python hf_trainer.py
"""

import torch
from torch import nn
from torch.utils.data import Dataset
from transformers import Trainer, TrainingArguments

from launchhelm.torch.hf import LaunchHelmCallback


class Tiny(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(4, 2)

    def forward(self, x=None, y=None, labels=None, **kwargs):
        logits = self.linear(x if x is not None else torch.ones(4))
        loss = (logits - (y if y is not None else torch.zeros(2))).pow(2).mean()
        return {"loss": loss, "logits": logits}


class Rows(Dataset):
    def __len__(self) -> int:
        return 16

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        return {"x": torch.ones(4), "y": torch.tensor([1.0, 0.0])}


args = TrainingArguments(
    output_dir="/tmp/lh-hf-example",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=1,
    save_strategy="no",
    report_to=[],  # LaunchHelm is the only logger
    disable_tqdm=True,
    run_name="hf-example",
)
trainer = Trainer(
    model=Tiny(),
    args=args,
    train_dataset=Rows(),
    callbacks=[LaunchHelmCallback("quickstart")],
)
trainer.train()