"""Is this model quicker on a CPU as 16-bit numbers or as 32-bit numbers?

    python number_size.py

The model's file holds 16-bit numbers, and transformers loads them as they are.
Every file on the page asks for 32-bit numbers instead (dtype=torch.float32).
This times one reply and three training steps both ways, to show why. It uses
the page's own train.py, one folder up, for the reply and for the steps.
"""
import sys
import time
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

PAGE_FILES = Path(__file__).resolve().parent.parent


def measure(train, dtype):
    """Seconds for one reply and for each of three training steps."""
    torch.manual_seed(0)
    tokenizer = AutoTokenizer.from_pretrained(train.MODEL)
    if dtype is None:
        model = AutoModelForCausalLM.from_pretrained(train.MODEL)   # as the file is
    else:
        model = AutoModelForCausalLM.from_pretrained(train.MODEL, dtype=dtype)
    name = str(next(model.parameters()).dtype).replace("torch.", "")
    started = time.perf_counter()
    train.reply_to(model, tokenizer, train.WATCHED)
    reply = time.perf_counter() - started
    print(f"    {name}: one reply took {reply:.1f} s", flush=True)

    model = train.prepare(model, "lora")
    trainable = [tensor for tensor in model.parameters() if tensor.requires_grad]
    optimizer = torch.optim.AdamW(trainable, lr=train.LEARNING_RATE["lora"])
    steps = []
    for example in train.read_examples(PAGE_FILES / "train.jsonl")[:3]:
        started = time.perf_counter()
        question_tokens = train.question_as_tokens(tokenizer, example["message"])
        answer_tokens = train.answer_as_tokens(tokenizer, example["order"])
        example_tokens = torch.tensor([question_tokens + answer_tokens])
        graded_tokens = torch.tensor([[-100] * len(question_tokens) + answer_tokens])
        loss = model(input_ids=example_tokens, labels=graded_tokens).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        steps.append(time.perf_counter() - started)
        print(f"    {name}: a training step took {steps[-1]:.2f} s", flush=True)
    return name, reply, sorted(steps)[1]


def main():
    sys.path.insert(0, str(PAGE_FILES))
    import train                    # the page's train.py

    print(f"torch {torch.__version__}, {torch.get_num_threads()} threads")
    print("As transformers loads the model when you ask for nothing:")
    plain = measure(train, None)
    print("With dtype=torch.float32, as on the page:")
    asked = measure(train, torch.float32)
    print()
    for name, reply, step in [plain, asked]:
        print(f"{name:9}  one reply {reply:5.1f} s   one training step {step:6.2f} s"
              f" (the middle one of three)")
    print(f"A training step in {plain[0]} takes {plain[2] / asked[2]:.0f} times"
          f" as long as in {asked[0]}.")


if __name__ == "__main__":
    main()