"""CPU reference for openjev (Qwen3_5ForSequenceClassification).

Writes ref/reference.pt with, per dtype (bf16, fp32): probs [N,3], logits [N,3], pooled hidden
[N,2560] (input of `score`), token ids and attention masks for a fixed NLI pair set.
"""
import json, os, sys, time
from pathlib import Path
import torch

SNAP = Path(os.environ.get("OPENJEV_DIR", "/home/ttuser/.cache/huggingface/hub/models--AlexWortega--openjev/snapshots/f8187e6e11d413d0771bcc7970b85f78e194264c/qwen3.5-4b-nli"))
OUT = Path(__file__).resolve().parent / "reference.pt"
TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"

PAIRS = [
    ("The bird is 0.05 below the centre of the gap.", "The bird is below the centre of the gap."),
    ("The bird is 0.05 below the centre of the gap.", "The bird is above the centre of the gap."),
    ("A man is playing a guitar.", "Someone is making music."),
    ("A man is playing a guitar.", "A woman is sleeping."),
    ("A man is playing a guitar.", "The man is famous."),
    ("Which gas do plants absorb during photosynthesis?", "The correct answer is: oxygen"),
    ("Which gas do plants absorb during photosynthesis?", "The correct answer is: carbon dioxide"),
    ("Which gas do plants absorb during photosynthesis?", "The correct answer is: nitrogen"),
    ("What is the capital of France?", "The correct answer is: Paris"),
    ("What is the capital of France?", "The correct answer is: Lyon"),
    ("Two dogs are running through a field.", "There are animals outdoors."),
    ("Two dogs are running through a field.", "The dogs are sleeping indoors."),
    ("The meeting was postponed because the CEO was ill.", "The meeting happened on schedule."),
    ("The meeting was postponed because the CEO was ill.", "The CEO was not well."),
    ("She bought three apples and two oranges at the market.", "She purchased fruit."),
    ("She bought three apples and two oranges at the market.", "She bought vegetables."),
    ("The Eiffel Tower is located in Paris, France, and was completed in 1889.", "The Eiffel Tower is in Germany."),
    ("The Eiffel Tower is located in Paris, France, and was completed in 1889.", "The tower was finished in the 19th century."),
    ("It is raining heavily outside.", "The weather is dry."),
    ("It is raining heavily outside.", "People might use umbrellas."),
]

def main():
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(str(SNAP))
    tok.padding_side = "right"
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    texts = [TEMPLATE.format(premise=p.strip(), hypothesis=h.strip()) for p, h in PAIRS]
    enc = tok(texts, truncation=True, max_length=4096, padding=True, return_tensors="pt")
    print("pad_token_id", tok.pad_token_id, "shape", enc["input_ids"].shape, flush=True)
    result = {"pairs": PAIRS, "texts": texts, "input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"],
              "pad_token_id": tok.pad_token_id, "lens": enc["attention_mask"].sum(1)}
    for dtype_name, dtype in [("bf16", torch.bfloat16), ("fp32", torch.float32)]:
        t0 = time.time()
        model = AutoModelForSequenceClassification.from_pretrained(str(SNAP), dtype=dtype)
        model.eval()
        tc = model.config.get_text_config()
        if tc.pad_token_id is None:
            tc.pad_token_id = tok.pad_token_id
        backbone = getattr(model, model.base_model_prefix)
        print(f"[{dtype_name}] loaded in {time.time()-t0:.1f}s; num_labels={model.config.num_labels}", flush=True)
        with torch.no_grad():
            t0 = time.time()
            out = backbone(**enc, output_hidden_states=True)
            h = out.last_hidden_state
            last = enc["attention_mask"].sum(1) - 1
            pooled = h[torch.arange(h.shape[0]), last]
            logits = model.score(pooled).float()
            probs = torch.softmax(logits, -1)
            print(f"[{dtype_name}] forward {time.time()-t0:.1f}s", flush=True)
        result[dtype_name] = {
            "probs": probs, "logits": logits, "pooled": pooled.float(),
            # per-layer hidden states at the pooled position (for layer-wise PCC debugging)
            "hidden_last": torch.stack([hs[torch.arange(h.shape[0]), last].float() for hs in out.hidden_states], 0),
            "score_weight": model.score.weight.detach().float(),
        }
        for i, (p, hy) in enumerate(PAIRS):
            pr = probs[i].tolist()
            print(f"[{dtype_name}] {i:2d} argmax={['CON','ENT','NEU'][int(probs[i].argmax())]} probs=[{pr[0]:.4f} {pr[1]:.4f} {pr[2]:.4f}] | {p[:50]} || {hy[:50]}")
        del model
    torch.save(result, OUT)
    print("saved", OUT)

if __name__ == "__main__":
    main()