Le chat et le fine-tuning SFT.

1"""Étape 5 — le chat template et le fine-tuning supervisé (leçons 12 et 14).2
3Le chat template transforme une liste de messages en une seule suite de tokens :4
5    <|im_start|>system\\n…<|im_end|>\\n<|im_start|>user\\n…<|im_end|>\\n<|im_start|>assistant\\n6
7et le modèle écrit la réponse, puis <|im_end|>. Le fine-tuning part du modèle8pré-entraîné de l'étape 3, sur de courtes conversations à propos du roman. La loss ne9compte que la réponse de l'assistant et son <|im_end|> : le modèle lit la question, mais10n'est corrigé que sur ce qu'il doit apprendre à écrire, et sur le moment où s'arrêter.11
12    uv run python -m packages.minilab.chat13"""14
15import json16import time17from dataclasses import dataclass18from pathlib import Path19
20import torch21
22from packages.minilab.generate import generate23from packages.minilab.model import MiniLabModel, load, save24from packages.minilab.tokenizer import CHECKPOINTS, IM_END, IM_START, Tokenizer, get_tokenizer25
26SFT_DATA = Path(__file__).parent / "data" / "sft_fr.jsonl"27CHAT_CHECKPOINT = CHECKPOINTS / "chat.pt"28SYSTEM_PROMPT = "Tu réponds en une phrase."29IGNORE = -100  # les cibles ignorées par la loss (masquage de la loss)30
31Message = dict[str, str]32
33
34def render(messages: list[Message], tokenizer: Tokenizer) -> tuple[list[int], list[bool]]:35    """Les identifiants de la conversation, et pour chacun, s'il compte dans la loss."""36    start, end = tokenizer.special[IM_START], tokenizer.special[IM_END]37    ids: list[int] = []38    trained: list[bool] = []39    for message in messages:40        header = [start, *tokenizer.encode(f"{message['role']}\n")]41        content = [*tokenizer.encode(message["content"]), end]42        is_answer = message["role"] == "assistant"43        ids += header + content + tokenizer.encode("\n")44        trained += [False] * len(header) + [is_answer] * len(content) + [False]45    return ids, trained46
47
48def prompt_ids(messages: list[Message], tokenizer: Tokenizer) -> list[int]:49    """La conversation, terminée par un tour assistant vide : là où le modèle écrit."""50    ids, _ = render(messages, tokenizer)51    return ids + [tokenizer.special[IM_START], *tokenizer.encode("assistant\n")]

54def with_system(messages: list[Message]) -> list[Message]:55    if messages and messages[0]["role"] == "system":56        return messages57    return [{"role": "system", "content": SYSTEM_PROMPT}, *messages]58
59
60def fit_to_window(61    messages: list[Message], tokenizer: Tokenizer, context_length: int, max_new_tokens: int62) -> tuple[list[Message], list[int]]:63    """Retire les tours les plus anciens, une question et sa réponse ensemble, jusqu'à ce64    que prompt + max_new_tokens tienne dans la fenêtre de contexte. Le system prompt et65    le dernier message sont toujours gardés (leçon 14)."""66    kept = with_system(messages)67    while True:68        ids = prompt_ids(kept, tokenizer)69        if len(ids) + max_new_tokens <= context_length:70            return kept, ids71        if len(kept) <= 2:72            raise ValueError("Même seul, le dernier message ne tient pas dans la fenêtre.")73        drop = 2 if kept[1]["role"] == "user" and kept[2]["role"] == "assistant" else 174        kept = [kept[0], *kept[1 + drop :]]  # le plus ancien tour après le system prompt part75
76
77def load_dialogues(path: Path = SFT_DATA) -> list[list[Message]]:78    lines = path.read_text(encoding="utf-8").splitlines()79    return [with_system(json.loads(line)["messages"]) for line in lines if line.strip()]80
81
82def make_examples(83    dialogues: list[list[Message]], tokenizer: Tokenizer, context_length: int84) -> list[tuple[list[int], list[int]]]:85    """Des paires (entrée, cible), la cible décalée d'un token et masquée hors des réponses."""86    examples = []87    for dialogue in dialogues:88        ids, trained = render(dialogue, tokenizer)89        if len(ids) > context_length + 1:90            raise ValueError("Un dialogue d'entraînement dépasse la fenêtre de contexte.")91        targets = [92            token if keep else IGNORE for token, keep in zip(ids[1:], trained[1:], strict=True)93        ]94        examples.append((ids[:-1], targets))95    return examples

98@dataclass(frozen=True)99class SFTSettings:100    steps: int = 300101    batch_size: int = 16102    learning_rate: float = 1e-3103    seed: int = 0104
105
106def fine_tune(107    model: MiniLabModel, examples: list[tuple[list[int], list[int]]], settings: SFTSettings108) -> list[float]:109    torch.manual_seed(settings.seed)110    generator = torch.Generator().manual_seed(settings.seed)111    optimizer = torch.optim.AdamW(model.parameters(), lr=settings.learning_rate, weight_decay=0.1)112    losses = []113    model.train()114    for _ in range(settings.steps):115        picked = torch.randint(0, len(examples), (settings.batch_size,), generator=generator)116        batch = [examples[i] for i in picked.tolist()]117        width = max(len(inputs) for inputs, _ in batch)118        # On complète à la même longueur ; le remplissage n'est jamais entraîné.119        inputs = torch.tensor([x + [0] * (width - len(x)) for x, _ in batch])120        targets = torch.tensor([y + [IGNORE] * (width - len(y)) for _, y in batch])121        _, loss = model(inputs, targets)122        optimizer.zero_grad()123        loss.backward()124        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)125        optimizer.step()126        losses.append(loss.item())127    return losses128
129
130def answer(131    model: MiniLabModel, tokenizer: Tokenizer, messages: list[Message], max_new_tokens: int = 60132) -> tuple[str, bool]:133    """La réponse en glouton, arrêtée par <|im_end|> ; dit aussi si le modèle s'est134    arrêté de lui-même."""135    _, ids = fit_to_window(messages, tokenizer, model.config.context_length, max_new_tokens)136    end = tokenizer.special[IM_END]137    tokens = [step.token_id for step in generate(model, ids, max_new_tokens, stop_id=end)]138    stopped = bool(tokens) and tokens[-1] == end139    return tokenizer.decode(tokens[:-1] if stopped else tokens), stopped

142def main() -> None:143    from packages.minilab.train import BASE_CHECKPOINT144
145    if not BASE_CHECKPOINT.exists():146        raise SystemExit("Checkpoint absent : lancez d'abord l'étape 3 (packages.minilab.train).")147    tokenizer = get_tokenizer()148    model = load(BASE_CHECKPOINT)149    dialogues = load_dialogues()150    examples = make_examples(dialogues, tokenizer, model.config.context_length)151    sample, trained = render(dialogues[0], tokenizer)152    print("Le chat template d'un exemple, et ce que la loss corrige ([…]) :")153    shown = "".join(154        f"[{tokenizer.decode([t])}]" if keep else tokenizer.decode([t])155        for t, keep in zip(sample, trained, strict=True)156    )157    print("  " + shown.replace("\n", "\n  "))158
159    question = "Qui est Passepartout ?"160    print(f"\nModèle de base, avant le SFT : « {question} »")161    print(162        f"  → « {answer(model, tokenizer, [{'role': 'user', 'content': question}], 40)[0]} »".replace(163            "\n", " ↵ "164        )165    )166
167    settings = SFTSettings()168    print(169        f"\nFine-tuning supervisé : {len(examples)} dialogues, {settings.steps} pas, lots de {settings.batch_size}"170    )171    start = time.perf_counter()172    losses = fine_tune(model, examples, settings)173    print(174        f"  loss sur les réponses : {losses[0]:.2f} → {sum(losses[-20:]) / 20:.2f} en {time.perf_counter() - start:.0f} s"175    )176    model.eval()177    save(model, CHAT_CHECKPOINT)178    print(f"  checkpoint sauvegardé : {CHAT_CHECKPOINT.relative_to(CHECKPOINTS.parents[1])}")179
180    print("\nL'assistant, en glouton, arrêté par <|im_end|> :")181    for question in [182        "Qui est Passepartout ?",183        "Quel est le pari de Phileas Fogg ?",184        "Peux-tu me présenter Fix ?",  # une formulation absente des données185        "Comment Phileas Fogg traverse-t-il l'Inde ?",186        "Quelle est la capitale de l'Espagne ?",  # hors du roman187    ]:188        text, stopped = answer(model, tokenizer, [{"role": "user", "content": question}])189        print(190            f"  {question}\n    → {text}  {'[<|im_end|>]' if stopped else '[coupé : limite de tokens]'}"191        )192
193    history = [194        {"role": "user", "content": "Qui est Phileas Fogg ?"},195        {"role": "assistant", "content": "Un gentleman anglais très ponctuel."},196    ] * 4 + [{"role": "user", "content": "Qui est Aouda ?"}]197    kept, ids = fit_to_window(history, tokenizer, model.config.context_length, 60)198    print(199        f"\nFenêtre de contexte : {len(history)} messages + system → {len(kept)} gardés ({len(ids)} tokens de prompt + 60 réservés ≤ {model.config.context_length})"200    )201
202
203if __name__ == "__main__":204    main()