"""Lesson 11: instruction-tuning data - chat templates and masking the loss so the model learns only the answers."""
import torch
pairs = [("What does VLOOKUP do?", "It finds a value in the first column of a table and returns a value from another column."),
         ("When does the Indian financial year start?", "On 1 April.")]

def to_chat(q, a):   # a simple template; each real model has its own (see its model card)
    return "<|user|>\n" + q + "\n<|assistant|>\n", a + "<|end|>\n"

for q, a in pairs:
    prompt, answer = to_chat(q, a)
    mask = [0] * len(prompt) + [1] * len(answer)     # 1 = learn from this position
    print(repr(prompt + answer))
    print("characters learned from:", sum(mask), "of", len(mask))
ids = torch.arange(10); m = torch.tensor([0] * 6 + [1] * 4)
print("labels:", torch.where(m == 1, ids, torch.tensor(-100)).tolist())   # -100 is ignored by cross_entropy