Formatting an SFT dataset

pip install torch

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# A toy "chat" example. The prompt is unpredictable noise, drawn uniformly from
# 18 symbols. The answer is a rule the model CAN learn. This mirrors real SFT:
# the user's question is not yours to predict, the assistant's reply is.
V, PROMPT_LEN, ANS_LEN = 40, 12, 3
BOS, SEP = 38, 39


def make_batch(n, gen):
    p = torch.randint(0, 18, (n, PROMPT_LEN), generator=gen)
    ans = 20 + (p[:, :1] % 8).repeat(1, ANS_LEN) + torch.arange(ANS_LEN)
    return torch.cat([torch.full((n, 1), BOS), p, torch.full((n, 1), SEP), ans], 1)


gen = torch.Generator().manual_seed(0)
train, test = make_batch(512, gen), make_batch(128, gen)
print("one training row:", train[0].tolist())
print(f"  BOS, then {PROMPT_LEN} noise tokens, then SEP, then a {ANS_LEN}-token answer")


def labels_for(seq, mask_prompt):
    lab = seq[:, 1:].clone()                      # shift-by-one targets
    if mask_prompt:
        lab[:, :PROMPT_LEN + 1] = -100            # -100 = "do not train on this"
    return lab


print("\nlabels with prompt masking (-100 is ignored by cross_entropy):")
print(" ", labels_for(train[:1], True)[0].tolist())
print("labels without masking:")
print(" ", labels_for(train[:1], False)[0].tolist())

n_prompt = PROMPT_LEN + 1
n_ans = train.shape[1] - 1 - n_prompt
print(f"\nof {n_prompt + n_ans} supervised positions per row, {n_prompt} are prompt "
      f"and {n_ans} are answer")
print(f"the prompt tokens are uniform over 18 symbols, so their loss cannot fall "
      f"below ln(18) = {math.log(18):.3f}")


def run(mask_prompt, steps=150):
    torch.manual_seed(1)
    layer = nn.TransformerEncoderLayer(48, 4, 96, batch_first=True, dropout=0.0)
    m = nn.ModuleDict({"emb": nn.Embedding(V, 48), "pos": nn.Embedding(24, 48),
                       "blocks": nn.TransformerEncoder(layer, 2), "head": nn.Linear(48, V)})
    opt = torch.optim.AdamW(m.parameters(), lr=3e-3)

    def fwd(seq):
        h = m["emb"](seq) + m["pos"](torch.arange(seq.shape[1]))
        mask = nn.Transformer.generate_square_subsequent_mask(seq.shape[1])
        return m["head"](m["blocks"](h, mask=mask, is_causal=True))

    for _ in range(steps):
        logits = fwd(train)[:, :-1]
        loss = F.cross_entropy(logits.reshape(-1, V),
                               labels_for(train, mask_prompt).reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()

    with torch.no_grad():
        logits = fwd(test)[:, :-1]
        full = labels_for(test, False)
        per_tok = F.cross_entropy(logits.reshape(-1, V), full.reshape(-1),
                                  reduction="none").reshape(full.shape)
        prompt_loss = per_tok[:, :n_prompt].mean().item()
        ans_loss = per_tok[:, n_prompt:].mean().item()
        acc = (logits.argmax(-1)[:, n_prompt:] == full[:, n_prompt:]).float().mean().item()
    return prompt_loss, ans_loss, acc


print(f"\n{'training signal':<26} {'prompt loss':>12} {'answer loss':>12} {'answer acc':>11}")
for mask_prompt in (False, True):
    pl, al, acc = run(mask_prompt)
    name = "every token" if not mask_prompt else "answer tokens only"
    print(f"{name:<26} {pl:>12.4f} {al:>12.4f} {acc:>10.1%}")

# 1. language modelling - loss on everything
{"text": "The sky is blue."}

# 2. conversational language modelling
{"messages": [{"role": "user", "content": "What color is the sky?"},
              {"role": "assistant", "content": "It is blue."}]}

# 3. prompt-completion - loss on the completion only, BY DEFAULT
{"prompt": "The sky is", "completion": " blue."}

# 4. conversational prompt-completion
{"prompt": [{"role": "user", "content": "What color is the sky?"}],
 "completion": [{"role": "assistant", "content": "It is blue."}]}

import statistics

# stand-in for your tokenised dataset: (prompt_len, completion_len) per example
rows = [(31, 88), (410, 22), (77, 640), (18, 9), (1203, 40), (95, 210),
        (64, 55), (890, 700), (12, 5), (150, 120)]

MAX = 1024
total = [p + c for p, c in rows]
print(f"examples: {len(rows)}")
print(f"median total length: {statistics.median(total):.0f} tokens")
print(f"longest: {max(total)} tokens")
print(f"would be truncated at max_length={MAX}: "
      f"{sum(t > MAX for t in total)} of {len(rows)}")
print(f"completion is under 10 tokens in {sum(c < 10 for _, c in rows)} examples")
print(f"mean fraction of tokens that are completion: "
      f"{statistics.mean(c / (p + c) for p, c in rows):.2f}")