#!/usr/bin/env python3 """Evaluate NextTerm-style MLX models on the canonical 111-series M1 set.""" from __future__ import annotations import argparse import gc import json import math import time from dataclasses import dataclass from decimal import Decimal 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 from eval_m1_monthly_mape_mlx import ( SuppressTokenLogits, context_scale, mape_for_predictions, parse_decimal, parse_generated_terms, scale_decimal_to_int, ) 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-440M" if local_model.exists(): return local_model return Path("N8Programs/NextTerm-440M") DATA_PATH = SCRIPT_DIR / "m1_competition_111.jsonl" MODEL_PATH = default_model_path() OUTPUT_PATH = Path("m1_eval_results/m1_competition111_nextterm440m_greedy_per_series.jsonl") SUMMARY_PATH = Path("m1_eval_results/m1_competition111_nextterm440m_greedy_summary.json") @dataclass class SeriesRecord: row_index: int series_name: str original_id: str period: str frequency: int type: str description: str horizon: int raw_context: list[str] raw_target: list[str] context_values: list[Decimal] target_values: list[Decimal] naive2_forecast: list[Decimal] naive2_mape: float scale: int scaled_context: list[int] def load_m1_competition(path: Path) -> list[SeriesRecord]: records: list[SeriesRecord] = [] with path.open("r", encoding="utf-8") as f: for line in f: if not line.strip(): continue row = json.loads(line) raw_context = [str(x) for x in row["context"]] raw_target = [str(x) for x in row["target"]] context_values = [parse_decimal(x) for x in raw_context] target_values = [parse_decimal(x) for x in raw_target] scale = context_scale(raw_context) scaled_context = [scale_decimal_to_int(v, scale) for v in context_values] records.append( SeriesRecord( row_index=int(row["row_index"]), series_name=str(row["series_name"]), original_id=str(row["original_id"]), period=str(row["period"]), frequency=int(row["frequency"]), type=str(row["type"]), description=str(row["description"]), horizon=int(row["horizon"]), raw_context=raw_context, raw_target=raw_target, context_values=context_values, target_values=target_values, naive2_forecast=[parse_decimal(str(x)) for x in row["naive2_forecast"]], naive2_mape=float(row["naive2_mape"]), scale=scale, scaled_context=scaled_context, ) ) return records def decimal_or_none(value: Decimal | None) -> str | None: return str(value) if value is not None else None 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]) -> dict: def collect(key: str, rows: list[dict] = records) -> list[float]: vals: list[float] = [] for r in rows: value = r.get(key) if value is not None and math.isfinite(float(value)): vals.append(float(value)) return vals model_series = collect("mape") naive2_series = collect("naive2_mape") all_model_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_naive2_apes = [ float(x) for r in records for x in (r.get("naive2_apes") or []) if x is not None and math.isfinite(float(x)) ] periods = sorted({str(r["period"]) for r in records}) by_period = {} for period in periods: rows = [r for r in records if r["period"] == period] period_model = collect("mape", rows) period_naive2 = collect("naive2_mape", rows) period_model_apes = [ float(x) for r in rows for x in (r.get("apes") or []) if x is not None and math.isfinite(float(x)) ] period_naive2_apes = [ float(x) for r in rows for x in (r.get("naive2_apes") or []) if x is not None and math.isfinite(float(x)) ] by_period[period] = { "series_count": len(rows), "model_macro_mape": mean(period_model) if period_model else None, "model_point_mape": mean(period_model_apes) if period_model_apes else None, "naive2_macro_mape": mean(period_naive2) if period_naive2 else None, "naive2_point_mape": mean(period_naive2_apes) if period_naive2_apes else None, "parsed_full_series": sum( 1 for r in rows if int(r.get("parsed_terms", 0)) >= int(r.get("horizon", 0)) ), } horizon_counts: dict[int, int] = {} for r in records: h = int(r["horizon"]) horizon_counts[h] = horizon_counts.get(h, 0) + 1 return { "series_count": len(records), "horizon_counts": dict(sorted(horizon_counts.items())), "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_model_apes) if all_model_apes else None, "naive2_macro_mape": mean(naive2_series) if naive2_series else None, "naive2_point_mape": mean(all_naive2_apes) if all_naive2_apes else None, "parsed_full_series": sum( 1 for r in records if int(r.get("parsed_terms", 0)) >= int(r.get("horizon", 0)) ), "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)), "by_period": by_period, } 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("--batch-size", type=int, default=64) parser.add_argument("--max-new-tokens", type=int, default=1024) parser.add_argument("--max-context-tokens", type=int, default=2048) parser.add_argument("--max-series", type=int, default=0) parser.add_argument( "--allow-eos-before-horizon", action="store_true", help="Do not suppress EOS/pad while waiting for all forecast terms.", ) parser.add_argument( "--missing-policy", choices=["zero", "last_context", "skip"], default="zero", ) parser.add_argument("--overwrite", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() started = time.perf_counter() series = load_m1_competition(args.data_path) if args.max_series > 0: series = series[: args.max_series] print(f"Loaded {len(series)} canonical M1 competition series from {args.data_path}") print("Horizons:", {h: sum(1 for s in series if s.horizon == h) for h in sorted({s.horizon for s in series})}) 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 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 = [] suppress_token_ids: list[int] = [] if getattr(tokenizer, "eos_token_id", None) is not None: suppress_token_ids.append(tokenizer.eos_token_id) if args.allow_eos_before_horizon: stop_tokens.append([tokenizer.eos_token_id]) if getattr(tokenizer, "pad_token_id", None) is not None: suppress_token_ids.append(tokenizer.pad_token_id) if args.allow_eos_before_horizon: stop_tokens.append([tokenizer.pad_token_id]) suppress_processor = None if args.allow_eos_before_horizon else SuppressTokenLogits(suppress_token_ids) 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: if uid in finished: return 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, s.horizon) predictions: list[Decimal | None] = [ Decimal(term) / Decimal(s.scale) for term in scaled_terms ] predictions.extend([None] * (s.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, ) naive2_mape, naive2_apes, _ = mape_for_predictions( s.target_values, s.naive2_forecast, missing_policy="skip", fallback=fallback, ) record = { "row_index": s.row_index, "series_name": s.series_name, "original_id": s.original_id, "period": s.period, "frequency": s.frequency, "type": s.type, "description": s.description, "horizon": s.horizon, "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": [decimal_or_none(x) for x in predictions], "parsed_terms": len(scaled_terms), "missing_terms": missing, "malformed": malformed, "mape": mape, "apes": apes, "naive2_forecast": [str(x) for x in s.naive2_forecast], "naive2_mape": naive2_mape, "naive2_apes": naive2_apes, "generated_text": generated_text, "finish_reason": finish_reason, } out.write(json.dumps(record) + "\n") out.flush() progress.update(1) while todo_indices and len(finished) < len(todo_indices): responses = gen.next() if isinstance(responses, tuple) and len(responses) == 2: prompt_responses, generation_responses = responses elif isinstance(responses, list): prompt_responses, generation_responses = [], responses else: raise RuntimeError( "Unexpected mlx_lm BatchGenerator.next() API. Update mlx_lm version." ) 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 idx = uid_to_idx[uid] horizon = series[idx].horizon 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() eval_index_set = set(eval_indices) all_records = [ r for idx, r in sorted(load_completed(args.output).items()) if idx in eval_index_set ] aggregate = aggregate_summary(all_records) elapsed = time.perf_counter() - started summary = { "data_path": str(args.data_path), "model": str(args.model), "output": str(args.output), "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": not args.allow_eos_before_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", "naive2_macro_mape", "naive2_point_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()