#!/usr/bin/env python3
"""Quick inference test with trained v6 model from Kaggle."""
import sys
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))

from architecture import MorphConfig, MorphModel

CKPT_DIR = Path(__file__).resolve().parent / "kaggle-output2" / "output" / "morph-model"
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"


def main():
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)

    print("Loading base model (fp16 for inference)...")
    base_model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
        torch_dtype=torch.float16,
        device_map="auto",
    )

    print("Building MorphModel v6...")
    config = MorphConfig(base_model=MODEL_NAME, max_seq_len=8192)
    model = MorphModel(config)
    model.base_model_raw = base_model
    model.apply_lora(target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ])
    model.tokenizer = tokenizer
    model.eval()

    print(f"Loading trained weights from {CKPT_DIR / 'checkpoint-393' / 'model.safetensors'} ...")
    from safetensors import safe_open
    state_dict = {}
    with safe_open(str(CKPT_DIR / "checkpoint-393" / "model.safetensors"), framework="pt") as f:
        for key in f.keys():
            state_dict[key] = f.get_tensor(key)

    print(f"Loaded {len(state_dict)} tensors from checkpoint")
    model.load_state_dict(state_dict)
    print("Weights loaded successfully")

    prompt = "What is 2+2? Think step by step."
    inputs = tokenizer(prompt, return_tensors="pt").to(model.base_model_raw.device)

    with torch.no_grad():
        outputs = model.base_model_raw.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=False,
        )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print("\n=== Inference Test ===")
    print(result)
    print("=== End ===")


if __name__ == "__main__":
    main()