"""Is the model quicker once the LoRA add-ons are merged into it?

    python merged_speed.py

Run it after "python train.py lora". It answers the first ten handwritten test
messages twice: with the add-ons kept beside the model, as score.py loads them,
and with the add-ons merged into the model's own numbers, as use.py does. It
also says whether the two gave the same replies.
"""
import json
import time
from pathlib import Path

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"
INSTRUCTION = "You take orders for a noodle shop. Reply with the order as JSON."
PAGE_FILES = Path(__file__).resolve().parent.parent


def reply_to(model, tokenizer, message):
    chat = [{"role": "system", "content": INSTRUCTION}, {"role": "user", "content": message}]
    question = tokenizer.apply_chat_template(
        chat, add_generation_prompt=True, return_tensors="pt")
    question_ids = question["input_ids"]
    output = model.generate(question_ids, max_new_tokens=50, do_sample=False)
    return tokenizer.decode(output[0][question_ids.shape[1]:], skip_special_tokens=True)


def answer_all(model, tokenizer, messages):
    """Every reply, and the seconds one reply took on average."""
    reply_to(model, tokenizer, messages[0])         # the first reply is always slow
    started = time.perf_counter()
    replies = [reply_to(model, tokenizer, message) for message in messages]
    return replies, (time.perf_counter() - started) / len(messages)


def main():
    with open(PAGE_FILES / "test_handwritten.jsonl", encoding="utf-8") as file:
        messages = [json.loads(line)["message"] for line in file][:10]
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
    model = PeftModel.from_pretrained(model, PAGE_FILES / "tuned_lora")
    beside, beside_seconds = answer_all(model, tokenizer, messages)
    model = model.merge_and_unload()
    merged, merged_seconds = answer_all(model, tokenizer, messages)
    same = sum(one == other for one, other in zip(beside, merged))
    print(f"add-ons beside the model: {beside_seconds:.2f} s a reply")
    print(f"add-ons merged into it:   {merged_seconds:.2f} s a reply")
    print(f"the same reply both ways: {same} of {len(messages)}")


if __name__ == "__main__":
    main()