eval_m1_monthly_mape_mlx.py

#!/usr/bin/env python3
"""Evaluate NextTerm-style MLX models on the M1 monthly forecasting dataset."""

from __future__ import annotations

import argparse
import gc
import json
import math
import re
import time
from dataclasses import dataclass
from decimal import Decimal, InvalidOperation, localcontext
from pathlib import Path
from statistics import mean, median

import mlx.core as mx
from mlx_lm import load
from mlx_lm.generate import BatchGenerator
from tqdm import tqdm


SCRIPT_DIR = Path(__file__).resolve().parent


def default_model_path() -> Path:
    if (SCRIPT_DIR / "model.safetensors").exists():
        return SCRIPT_DIR
    local_model = SCRIPT_DIR / "NextTerm-47M"
    if local_model.exists():
        return local_model
    return Path("N8Programs/NextTerm-47M")


DATA_PATH = SCRIPT_DIR / "m1_monthly_dataset.txt"
MODEL_PATH = default_model_path()
OUTPUT_PATH = Path("m1_eval_results/m1_monthly_nextterm47m_per_series.jsonl")
SUMMARY_PATH = Path("m1_eval_results/m1_monthly_nextterm47m_summary.json")


@dataclass
class SeriesRecord:
    row_index: int
    series_name: str
    start_timestamp: str
    raw_values: list[str]
    values: list[Decimal]
    context_values: list[Decimal]
    target_values: list[Decimal]
    scale: int
    scaled_context: list[int]


class SuppressTokenLogits:
    """Set selected token logits to a large negative value."""

    def __init__(self, token_ids: list[int]):
        self.token_ids = sorted({int(t) for t in token_ids if t is not None and int(t) >= 0})
        self._bias_by_width: dict[int, mx.array] = {}

    def __call__(self, tokens: mx.array, logits: mx.array) -> mx.array:
        width = int(logits.shape[-1])
        bias = self._bias_by_width.get(width)
        if bias is None:
            values = [0.0] * width
            for token_id in self.token_ids:
                if token_id < width:
                    values[token_id] = -1.0e9
            bias = mx.array(values, dtype=logits.dtype)
            self._bias_by_width[width] = bias
        return logits + bias


def parse_decimal(raw: str) -> Decimal:
    raw = raw.strip()
    try:
        value = Decimal(raw)
    except InvalidOperation as exc:
        raise ValueError(f"Could not parse decimal value {raw!r}") from exc
    if not value.is_finite():
        raise ValueError(f"Non-finite decimal value {raw!r}")
    return value


def context_scale(raw_context_values: list[str]) -> int:
    """Choose an integer scale using only visible decimal precision in context."""
    max_places = 0
    for raw in raw_context_values:
        value = parse_decimal(raw)
        exponent = value.as_tuple().exponent
        if exponent < 0:
            max_places = max(max_places, -exponent)
    return 10**max_places


def scale_decimal_to_int(value: Decimal, scale: int) -> int:
    scaled = value * Decimal(scale)
    with localcontext() as ctx:
        ctx.prec = max(50, len(scaled.as_tuple().digits) + 10)
        rounded = scaled.to_integral_value()
    if rounded != scaled:
        # This can only happen if the held-out side has more precision than the
        # context-derived scale. It is fine for scoring, but not for prompting.
        raise ValueError(f"Context scale {scale} does not make {value} integral")
    return int(rounded)


def load_m1_monthly(path: Path, horizon: int | None = None) -> tuple[list[SeriesRecord], int]:
    records: list[SeriesRecord] = []
    parsed_horizon: int | None = horizon
    in_data = False
    with path.open("r", encoding="latin-1", newline=None) as f:
        for raw_line in f:
            line = raw_line.strip()
            if not line:
                continue
            if line.startswith("@horizon") and parsed_horizon is None:
                parts = line.split()
                if len(parts) >= 2:
                    parsed_horizon = int(parts[1])
                continue
            if line == "@data":
                in_data = True
                continue
            if not in_data or line.startswith("#") or line.startswith("@"):
                continue

            try:
                series_name, start_timestamp, values_blob = line.split(":", 2)
            except ValueError as exc:
                raise ValueError(f"Malformed M1 data line: {line[:120]!r}") from exc
            raw_values = [part.strip() for part in values_blob.split(",") if part.strip()]
            values = [parse_decimal(part) for part in raw_values]
            if parsed_horizon is None:
                raise ValueError("No horizon provided and no @horizon metadata found")
            if len(values) <= parsed_horizon:
                continue

            raw_context = raw_values[:-parsed_horizon]
            scale = context_scale(raw_context)
            context_values = values[:-parsed_horizon]
            target_values = values[-parsed_horizon:]
            scaled_context = [scale_decimal_to_int(v, scale) for v in context_values]
            records.append(
                SeriesRecord(
                    row_index=len(records),
                    series_name=series_name,
                    start_timestamp=start_timestamp,
                    raw_values=raw_values,
                    values=values,
                    context_values=context_values,
                    target_values=target_values,
                    scale=scale,
                    scaled_context=scaled_context,
                )
            )
    if parsed_horizon is None:
        raise ValueError("No horizon provided and no @horizon metadata found")
    return records, parsed_horizon


def parse_generated_terms(text: str, limit: int) -> tuple[list[int], bool]:
    terms: list[int] = []
    current: list[str] = []
    malformed = False

    def flush_current() -> None:
        nonlocal malformed
        if not current:
            return
        token = "".join(current)
        current.clear()
        if token == "-":
            malformed = True
            return
        try:
            terms.append(int(token))
        except ValueError:
            malformed = True

    for ch in text:
        if len(terms) >= limit:
            break
        if ch.isdigit() or (ch == "-" and not current):
            current.append(ch)
        elif ch == ",":
            flush_current()
        elif ch.isspace():
            continue
        else:
            malformed = True
            flush_current()
            break
    if len(terms) < limit:
        flush_current()
    return terms[:limit], malformed


def as_float(value: Decimal) -> float:
    return float(value)


def ape(actual: Decimal, prediction: Decimal) -> float | None:
    if actual == 0:
        return None
    return float(abs(actual - prediction) / abs(actual) * Decimal(100))


def mape_for_predictions(
    actuals: list[Decimal],
    predictions: list[Decimal | None],
    *,
    missing_policy: str,
    fallback: Decimal,
) -> tuple[float | None, list[float | None], int]:
    apes: list[float | None] = []
    missing = 0
    for actual, prediction in zip(actuals, predictions):
        pred = prediction
        if pred is None:
            missing += 1
            if missing_policy == "skip":
                apes.append(None)
                continue
            if missing_policy == "zero":
                pred = Decimal(0)
            elif missing_policy == "last_context":
                pred = fallback
            else:
                raise ValueError(f"Unknown missing policy: {missing_policy}")
        apes.append(ape(actual, pred))
    valid = [x for x in apes if x is not None and math.isfinite(x)]
    return (mean(valid) if valid else None), apes, missing


def seasonal_naive_predictions(context: list[Decimal], horizon: int, season: int = 12) -> list[Decimal]:
    if not context:
        return [Decimal(0)] * horizon
    if len(context) < season:
        return [context[-1]] * horizon
    last_season = context[-season:]
    return [last_season[i % season] for i in range(horizon)]


def last_value_predictions(context: list[Decimal], horizon: int) -> list[Decimal]:
    fallback = context[-1] if context else Decimal(0)
    return [fallback] * horizon


def load_completed(path: Path) -> dict[int, dict]:
    completed: dict[int, dict] = {}
    if not path.exists():
        return completed
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            record = json.loads(line)
            completed[int(record["row_index"])] = record
    return completed


def aggregate_summary(records: list[dict], horizon: int) -> dict:
    def collect(key: str) -> list[float]:
        return [
            float(r[key])
            for r in records
            if r.get(key) is not None and math.isfinite(float(r[key]))
        ]

    model_series = collect("mape")
    seasonal_series = collect("seasonal_naive_mape")
    last_series = collect("last_value_mape")
    per_horizon: list[dict] = []
    for h in range(horizon):
        vals = []
        seasonal_vals = []
        last_vals = []
        for r in records:
            for source, target in [
                ("apes", vals),
                ("seasonal_naive_apes", seasonal_vals),
                ("last_value_apes", last_vals),
            ]:
                xs = r.get(source) or []
                if h < len(xs) and xs[h] is not None and math.isfinite(float(xs[h])):
                    target.append(float(xs[h]))
        per_horizon.append(
            {
                "horizon": h + 1,
                "mape": mean(vals) if vals else None,
                "seasonal_naive_mape": mean(seasonal_vals) if seasonal_vals else None,
                "last_value_mape": mean(last_vals) if last_vals else None,
                "n": len(vals),
            }
        )

    all_apes = [
        float(x)
        for r in records
        for x in (r.get("apes") or [])
        if x is not None and math.isfinite(float(x))
    ]
    all_seasonal_apes = [
        float(x)
        for r in records
        for x in (r.get("seasonal_naive_apes") or [])
        if x is not None and math.isfinite(float(x))
    ]
    all_last_apes = [
        float(x)
        for r in records
        for x in (r.get("last_value_apes") or [])
        if x is not None and math.isfinite(float(x))
    ]
    return {
        "series_count": len(records),
        "model_macro_mape": mean(model_series) if model_series else None,
        "model_median_series_mape": median(model_series) if model_series else None,
        "model_point_mape": mean(all_apes) if all_apes else None,
        "seasonal_naive_macro_mape": mean(seasonal_series) if seasonal_series else None,
        "seasonal_naive_point_mape": mean(all_seasonal_apes) if all_seasonal_apes else None,
        "last_value_macro_mape": mean(last_series) if last_series else None,
        "last_value_point_mape": mean(all_last_apes) if all_last_apes else None,
        "parsed_full_series": sum(1 for r in records if r.get("parsed_terms", 0) >= horizon),
        "total_missing_terms": sum(int(r.get("missing_terms", 0)) for r in records),
        "malformed_series": sum(1 for r in records if r.get("malformed", False)),
        "per_horizon": per_horizon,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-path", type=Path, default=DATA_PATH)
    parser.add_argument("--model", type=Path, default=MODEL_PATH)
    parser.add_argument("--output", type=Path, default=OUTPUT_PATH)
    parser.add_argument("--summary-output", type=Path, default=SUMMARY_PATH)
    parser.add_argument("--horizon", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--max-new-tokens", type=int, default=384)
    parser.add_argument("--max-context-tokens", type=int, default=2048)
    parser.add_argument("--max-series", type=int, default=0)
    parser.add_argument(
        "--suppress-eos-until-horizon",
        action="store_true",
        help="Suppress EOS/pad logits and stop only after the requested separator count or length cap.",
    )
    parser.add_argument(
        "--rerun-incomplete",
        action="store_true",
        help="When resuming, rerun rows whose previous record parsed fewer than horizon terms.",
    )
    parser.add_argument(
        "--missing-policy",
        choices=["zero", "last_context", "skip"],
        default="zero",
        help="How to score missing/unparseable forecast terms.",
    )
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    started = time.perf_counter()
    series, horizon = load_m1_monthly(args.data_path, args.horizon)
    if args.max_series > 0:
        series = series[: args.max_series]
    print(f"Loaded {len(series)} M1 monthly series from {args.data_path}; horizon={horizon}")

    model, tokenizer = load(str(args.model))
    print(f"Loaded model: {args.model}")

    prompts_text = [",".join(str(x) for x in s.scaled_context) + "," for s in series]
    prompts = [tokenizer.encode(text) for text in prompts_text]
    sep_token = tokenizer.encode("1,")[-1]
    eval_indices = [
        i
        for i, prompt in enumerate(prompts)
        if args.max_context_tokens <= 0 or len(prompt) < args.max_context_tokens
    ]
    skipped_long = len(prompts) - len(eval_indices)
    eval_indices = sorted(eval_indices, key=lambda i: len(prompts[i]))
    print(
        f"Evaluating {len(eval_indices)} series; skipped_long={skipped_long}; "
        f"batch_size={args.batch_size}; max_new_tokens={args.max_new_tokens}; sep_token={sep_token}"
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.summary_output.parent.mkdir(parents=True, exist_ok=True)
    if args.overwrite and args.output.exists():
        args.output.unlink()

    completed = load_completed(args.output)
    if args.rerun_incomplete:
        incomplete = [
            idx
            for idx, record in completed.items()
            if int(record.get("parsed_terms", 0)) < horizon
        ]
        for idx in incomplete:
            completed.pop(idx, None)
        if incomplete:
            print(
                "Rerunning incomplete rows: "
                + ", ".join(str(idx) for idx in sorted(incomplete))
            )
    if completed:
        print(f"Resuming from {args.output}: {len(completed)} series already done")
    todo_indices = [i for i in eval_indices if i not in completed]

    stop_tokens = [] if args.suppress_eos_until_horizon else [[tokenizer.eos_token_id]]
    suppress_token_ids: list[int] = []
    if getattr(tokenizer, "eos_token_id", None) is not None:
        suppress_token_ids.append(tokenizer.eos_token_id)
    if getattr(tokenizer, "pad_token_id", None) is not None:
        if not args.suppress_eos_until_horizon:
            stop_tokens.append([tokenizer.pad_token_id])
        suppress_token_ids.append(tokenizer.pad_token_id)
    suppress_processor = (
        SuppressTokenLogits(suppress_token_ids) if args.suppress_eos_until_horizon else None
    )
    gen = BatchGenerator(
        model,
        stop_tokens=stop_tokens or None,
        completion_batch_size=args.batch_size,
        prefill_batch_size=args.batch_size,
    )
    uid_to_idx: dict[int, int] = {}
    generated_tokens: dict[int, list[int]] = {}
    separator_counts: dict[int, int] = {}
    if todo_indices:
        logits_processors = (
            [[suppress_processor] for _ in todo_indices]
            if suppress_processor is not None
            else None
        )
        uids = gen.insert(
            [prompts[i] for i in todo_indices],
            [args.max_new_tokens] * len(todo_indices),
            logits_processors=logits_processors,
        )
        uid_to_idx = {uid: idx for uid, idx in zip(uids, todo_indices)}
        generated_tokens = {uid: [] for uid in uids}
        separator_counts = {uid: 0 for uid in uids}
    finished: set[int] = set()

    try:
        with args.output.open("a", encoding="utf-8") as out:
            with tqdm(total=len(todo_indices), desc="Generating") as progress:
                def finalize_uid(uid: int, finish_reason: str) -> None:
                    finished.add(uid)
                    idx = uid_to_idx[uid]
                    s = series[idx]
                    generated_text = tokenizer.decode(generated_tokens[uid])
                    scaled_terms, malformed = parse_generated_terms(generated_text, horizon)
                    predictions: list[Decimal | None] = [
                        Decimal(term) / Decimal(s.scale) for term in scaled_terms
                    ]
                    predictions.extend([None] * (horizon - len(predictions)))
                    fallback = s.context_values[-1] if s.context_values else Decimal(0)
                    mape, apes, missing = mape_for_predictions(
                        s.target_values,
                        predictions,
                        missing_policy=args.missing_policy,
                        fallback=fallback,
                    )
                    seasonal_preds = seasonal_naive_predictions(s.context_values, horizon)
                    seasonal_mape, seasonal_apes, _ = mape_for_predictions(
                        s.target_values,
                        seasonal_preds,
                        missing_policy="skip",
                        fallback=fallback,
                    )
                    last_preds = last_value_predictions(s.context_values, horizon)
                    last_mape, last_apes, _ = mape_for_predictions(
                        s.target_values,
                        last_preds,
                        missing_policy="skip",
                        fallback=fallback,
                    )
                    record = {
                        "row_index": idx,
                        "series_name": s.series_name,
                        "start_timestamp": s.start_timestamp,
                        "scale": s.scale,
                        "context_length": len(s.context_values),
                        "prompt_tokens": len(prompts[idx]),
                        "target": [str(x) for x in s.target_values],
                        "scaled_prediction_terms": scaled_terms,
                        "prediction": [str(x) if x is not None else None for x in predictions],
                        "parsed_terms": len(scaled_terms),
                        "missing_terms": missing,
                        "malformed": malformed,
                        "mape": mape,
                        "apes": apes,
                        "seasonal_naive_mape": seasonal_mape,
                        "seasonal_naive_apes": seasonal_apes,
                        "last_value_mape": last_mape,
                        "last_value_apes": last_apes,
                        "generated_text": generated_text,
                        "finish_reason": finish_reason,
                    }
                    out.write(json.dumps(record) + "\n")
                    out.flush()
                    progress.update(1)

                while todo_indices:
                    responses = gen.next()
                    if not isinstance(responses, tuple) or len(responses) != 2:
                        raise RuntimeError(
                            "Unexpected mlx_lm BatchGenerator.next() API. "
                            "Update your mlx_lm version."
                        )
                    prompt_responses, generation_responses = responses
                    if not prompt_responses and not generation_responses:
                        break
                    remove_uids: list[int] = []
                    for response in generation_responses:
                        uid = response.uid
                        if uid in finished:
                            continue
                        if response.finish_reason != "stop":
                            token = int(response.token)
                            generated_tokens[uid].append(token)
                            if token == sep_token:
                                separator_counts[uid] += 1
                                if separator_counts[uid] >= horizon:
                                    finalize_uid(uid, "separator_count")
                                    remove_uids.append(uid)
                                    continue
                        if response.finish_reason is not None:
                            finalize_uid(uid, str(response.finish_reason))
                    if remove_uids:
                        gen.remove(remove_uids)
                if len(finished) != len(todo_indices):
                    raise RuntimeError(f"Finished {len(finished)}/{len(todo_indices)} series")
    finally:
        gen.close()
        mx.clear_cache()
        gc.collect()

    all_records = [
        r
        for idx, r in sorted(load_completed(args.output).items())
        if idx in set(eval_indices)
    ]
    aggregate = aggregate_summary(all_records, horizon)
    elapsed = time.perf_counter() - started
    summary = {
        "data_path": str(args.data_path),
        "model": str(args.model),
        "output": str(args.output),
        "horizon": horizon,
        "series_loaded": len(series),
        "series_evaluated": len(all_records),
        "skipped_long": skipped_long,
        "max_context_tokens": args.max_context_tokens,
        "max_new_tokens": args.max_new_tokens,
        "batch_size": args.batch_size,
        "missing_policy": args.missing_policy,
        "suppress_eos_until_horizon": args.suppress_eos_until_horizon,
        "seconds": elapsed,
        **aggregate,
    }
    args.summary_output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({k: summary[k] for k in [
        "series_evaluated",
        "model_macro_mape",
        "model_point_mape",
        "seasonal_naive_macro_mape",
        "last_value_macro_mape",
        "parsed_full_series",
        "total_missing_terms",
        "malformed_series",
        "seconds",
    ]}, indent=2))
    print(f"Wrote {args.output}")
    print(f"Wrote {args.summary_output}")


if __name__ == "__main__":
    main()