"""What else did training change? Ask each model things that are not orders.

    python other_questions.py

Run it after "python train.py lora" and "python train.py full". Each question
goes to three models: the one you downloaded, the one with the LoRA add-ons,
and the fully trained one. Each model gets it twice: with no instruction at
all, and with the noodle-shop instruction that the training examples carry.
"""
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
QUESTIONS = [
    "What is the capital of Thailand?",
    "Write one short sentence about elephants.",
    "What time do you close today?",
]


def load(kind):
    if kind == "full":
        return AutoModelForCausalLM.from_pretrained(PAGE_FILES / "tuned_full")
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
    if kind == "lora":
        model = PeftModel.from_pretrained(model, PAGE_FILES / "tuned_lora")
    return model


def reply_to(model, tokenizer, chat):
    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=40, do_sample=False)
    reply = tokenizer.decode(output[0][question_ids.shape[1]:], skip_special_tokens=True)
    return " ".join(reply.split())


def main():
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    replies = {}
    for kind in ["before", "lora", "full"]:
        model = load(kind)
        for question in QUESTIONS:
            plain = [{"role": "user", "content": question}]
            shop = [{"role": "system", "content": INSTRUCTION}] + plain
            replies[kind, question, "no instruction"] = reply_to(model, tokenizer, plain)
            replies[kind, question, "shop instruction"] = reply_to(model, tokenizer, shop)
    for question in QUESTIONS:
        print(question)
        for instruction in ["no instruction", "shop instruction"]:
            for kind in ["before", "lora", "full"]:
                print(f"    {instruction:17} {kind:7} {replies[kind, question, instruction]}")
        print()


if __name__ == "__main__":
    main()