evaluate.py

"""Original unit-conversion robustness case study. Dataset/code: MIT license.

No networking occurs during evaluation. The separately downloaded model uses
its original Apache-2.0 license. Keep every completion, including failures.
"""
import argparse
import hashlib
import json
import platform
import re
import time
from fractions import Fraction
from pathlib import Path

REVISION = "448e61eb392c00f2c403185c5b56d5e0665bfaab"
SUFFIX = " Show your calculation briefly. End with exactly FINAL: <number> <unit>."


def corpus():
    rows = []
    cases = [
        ("length", "m", "A cable is 2.4 m long. You cut off 75 cm. How many meters remain?", Fraction(24, 10)-Fraction(75, 100)),
        ("volume", "mL", "A jug contains 1.25 L. You pour out 350 mL. How many milliliters remain?", Fraction(125, 100)*1000-350),
        ("area", "cm2", "A rectangular tile is 0.8 m long and 25 cm wide. What is its area in square centimeters?", Fraction(8, 10)*100*25),
        ("speed", "m", "A cart moves at 18 km/h for 40 seconds. How many meters does it travel?", Fraction(18*1000, 3600)*40),
        ("mass", "g", "A box has 1.6 kg of beads. Four equal bags use 275 g each. How many grams stay in the box?", Fraction(16, 10)*1000-4*275),
        ("flow", "L", "A tap delivers 450 mL per minute for 8 minutes. How many liters does it deliver?", Fraction(450*8, 1000)),
        ("length_swap", "m", "A cable is 3.2 m long. You cut off 45 cm. How many meters remain?", Fraction(32, 10)-Fraction(45, 100)),
        ("speed_swap", "m", "A cart moves at 27 km/h for 24 seconds. How many meters does it travel?", Fraction(27*1000, 3600)*24),
    ]
    for name, unit, question, answer in cases:
        for variant, prefix in [("base", ""), ("distractor", "There are 17 green stickers on the storage shelf; the stickers do not affect the quantities in this problem. ")]:
            rows.append(dict(id=f"{name}_{variant}", pair=name, variant=variant,
                             unit=unit, expected=str(answer), prompt=prefix+question+SUFFIX))
    return rows


def score(text, expected, unit, require_closed_thinking=False):
    # Only the final response after thinking is scored, never a reasoning number.
    if require_closed_thinking and "</think>" not in text:
        return {"parsed": False, "correct": False, "reason": "unfinished_thinking"}
    answer = text.split("</think>")[-1]
    matches = re.findall(r"(?m)^FINAL:[ \t]*([-+]?\d+(?:\.\d+)?(?:/\d+)?)[ \t]+([^\n\r]+?)[ \t]*$", answer)
    if len(matches) != 1:
        return {"parsed": False, "correct": False, "reason": "missing_or_ambiguous_final"}
    value, found_unit = matches[0]
    aliases = {"m": "m", "meter": "m", "meters": "m", "metre": "m", "metres": "m",
               "ml": "mL", "milliliter": "mL", "milliliters": "mL", "millilitres": "mL",
               "l": "L", "liter": "L", "liters": "L", "litre": "L", "litres": "L",
               "g": "g", "gram": "g", "grams": "g", "cm2": "cm2", "cm^2": "cm2",
               "cm²": "cm2", "square centimeters": "cm2", "square centimetres": "cm2"}
    normalized_unit = aliases.get(found_unit.strip().rstrip(".").lower())
    try:
        correct = Fraction(value) == Fraction(expected) and normalized_unit == unit
    except (ValueError, ZeroDivisionError):
        return {"parsed": False, "correct": False, "reason": "invalid_number"}
    return {"parsed": True, "correct": correct, "value": value, "unit": found_unit}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--prepare", action="store_true")
    parser.add_argument("--max-tokens", type=int, default=512)
    parser.add_argument("--batch-size", type=int, default=1)
    args = parser.parse_args()
    root = Path(__file__).resolve().parent
    rows = corpus()
    (root/"dataset.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    if args.prepare:
        assert score("<think>wrong 99</think>FINAL: 1.65 m", "33/20", "m")["correct"]
        assert not score("FINAL: 1.65 cm2", "33/20", "m")["correct"]
        assert not score("answer 1.65", "33/20", "m")["parsed"]
        assert not score("FINAL: 1 m\nFINAL: 2 m", "1", "m")["parsed"]
        assert not score("FINAL: 1/0 m", "1", "m")["parsed"]
        assert score("FINAL: 1.65 meters", "33/20", "m")["correct"]
        assert score("FINAL: 2000 cm²", "2000", "cm2")["correct"]
        assert not score("I plan to write:\nFINAL: 200 m", "200", "m", True)["correct"]
        assert score("Calculating.</think>FINAL: 200 m", "200", "m", True)["correct"]
        print(f"Prepared {len(rows)} cases; scoring edge checks pass", flush=True)
        return
    import torch
    import transformers
    from transformers import AutoTokenizer, AutoModelForCausalLM
    torch.set_num_threads(4)
    torch.manual_seed(20260910)
    model_path = root/"model"
    tok = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
    tok.padding_side = "left"
    tok.chat_template = (model_path/"chat_template.jinja").read_text(encoding="utf-8")
    start = time.perf_counter()
    model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True,
             local_files_only=True, dtype=torch.float32, attn_implementation="eager").eval()
    print(f"Model loaded in {time.perf_counter()-start:.2f}s", flush=True)
    chosen = rows[:2] if args.smoke else rows
    out = root/("smoke-results.jsonl" if args.smoke else "results.jsonl")
    metadata = dict(model="XHToken/Spark-X2.5-1.7B", revision=REVISION, python=platform.python_version(),
                    torch=torch.__version__, transformers=transformers.__version__, device="CPU",
                    hardware="Intel Core i7-8750H, 32 GB RAM", precision="float32", threads=4,
                    seed=20260910, do_sample=False, temperature=None, top_p=None, top_k=None,
                    max_new_tokens=args.max_tokens, enable_thinking=True, samples_per_problem=1,
                    dataset_sha256=hashlib.sha256((root/"dataset.json").read_bytes()).hexdigest(),
                    cases=len(chosen), smoke_only=args.smoke, batch_size=args.batch_size,
                    scorer_version="3: semantic unit aliases and explicit thinking-boundary check")
    out.with_suffix(".metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
    with out.open("x", encoding="utf-8") as f:
        for offset in range(0, len(chosen), args.batch_size):
            batch = chosen[offset:offset+args.batch_size]
            rendered_batch = [tok.apply_chat_template([dict(role="user", content=row["prompt"])], tokenize=False, add_generation_prompt=True, enable_thinking=True) for row in batch]
            inputs = tok(rendered_batch, return_tensors="pt", add_special_tokens=False, padding=True)
            start = time.perf_counter()
            with torch.inference_mode():
                result = model.generate(**inputs, max_new_tokens=args.max_tokens, do_sample=False, use_cache=True, pad_token_id=tok.pad_token_id, logits_to_keep=1)
            elapsed = time.perf_counter()-start
            for i, row in enumerate(batch):
                tokens = result[i, inputs.input_ids.shape[1]:].tolist()
                ended = tok.eos_token_id in tokens
                if ended:
                    tokens = tokens[:tokens.index(tok.eos_token_id)+1]
                raw = tok.decode(tokens, skip_special_tokens=False)
                record = dict(**row, rendered_prompt=rendered_batch[i], raw_output=raw, generated_tokens=len(tokens),
                              batch_seconds=elapsed, batch_offset=offset, seconds=elapsed/len(batch),
                              hit_token_limit=not ended and len(tokens)>=args.max_tokens,
                              score=score(raw.replace(tok.eos_token or "<unused>", ""), row["expected"], row["unit"], True))
                f.write(json.dumps(record)+"\n")
                f.flush()
                print(json.dumps({k:record[k] for k in ["id","generated_tokens","seconds","hit_token_limit","score"]}), flush=True)


if __name__ == "__main__":
    main()