from ch31_llm_foundations import MiniGPT, make_char_dataset


def build_preference_batch(stoi, prompt_str, chosen_str, rejected_str):
    def encode(s):
        return torch.tensor([[stoi[c] for c in s]], dtype=torch.long)
    return encode(prompt_str), encode(chosen_str), encode(rejected_str)


if __name__ == "__main__":
    torch.manual_seed(0)
    vocab = sorted(set(
        "abcdefghijklmnopqrstuvwxyz .," + "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    ))
    stoi = {ch: i for i, ch in enumerate(vocab)}

    policy = MiniGPT(
        vocab_size=len(vocab), d_model=64, n_heads=2,
        n_layers=2, max_len=64,
    )
    ref_model = copy.deepcopy(policy)
    for p in ref_model.parameters():
        p.requires_grad = False

    # a toy preference: detailed, complete-sentence answers are
    # "chosen" over terse, unhelpful ones for the same prompt.
    x, y_w, y_l = build_preference_batch(
        stoi,
        prompt_str="The weather today is ",
        chosen_str="quite pleasant with clear skies.",
        rejected_str="fine.",
    )

    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
    for step in range(200):
        loss, acc = dpo_loss(policy, ref_model, x, y_w, y_l, beta=0.1)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step % 50 == 0:
            print(f"step {step}  loss {loss.item():.4f}  "
                  f"implicit-reward accuracy {acc:.2f}")
