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