#!/usr/bin/env python3
"""Generate from the fine-tuned model, and from the untouched one, for comparison.

    python3 generate.py "Write a field note about monitoring that failed silently."

    --base       generate from the base model instead, with no adapter
    --adapter D  use a specific checkpoint rather than ./out/adapter

Both runs use the same temperature and the same seed, which is the only way the comparison
means anything: sampling settings live outside the model and change its character more than
most people expect. Part ten is about judging the result without fooling yourself.
"""
import argparse
import os
from pathlib import Path

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

HERE = Path(__file__).resolve().parent
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-8B")
SYSTEM = "You are Doc Quatrain. You write short technical field notes."

parser = argparse.ArgumentParser()
parser.add_argument("prompt", nargs="?", default="Write a field note about monitoring that failed silently.")
parser.add_argument("--base", action="store_true", help="no adapter, for comparison")
parser.add_argument("--adapter", default=str(HERE / "out" / "adapter"))
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--seed", type=int, default=0)
args = parser.parse_args()

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
set_seed(args.seed)

quant = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, quantization_config=quant, dtype=torch.bfloat16, device_map={"": 0}
)
if not args.base:
    model = PeftModel.from_pretrained(model, args.adapter)
model.eval()

messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": args.prompt}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt"
).to(model.device)

with torch.no_grad():
    out = model.generate(
        inputs,
        max_new_tokens=600,
        do_sample=True,
        temperature=args.temperature,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id,
    )

print(f"# {'base model' if args.base else args.adapter}, temperature {args.temperature}, seed {args.seed}\n")
print(tokenizer.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))