| """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(), |
| |
| "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() |
|
|
