#!/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()