"""Unified request/prediction schema + answer extraction for external eval. A *request* (what an adapter emits, consumed by multivar/eval/run_ts_align.py) is a dict with at least: request_id : str unique id : any sample id within the benchmark task_type : str one of TASK_TYPES (also mirrored into meta["task_type"]) variant : str "normal" (kept for run_ts_align compatibility) raw_ts : list[list[float]] [n_channels, seq_len] -> MOMENT prompt : str user prompt, contains "" placeholder system_prompt : str output-format instruction chosen by task_type gt : number | str | list typed by task_type meta : dict {benchmark, task, task_type, metric, options?, split, lang, ...} A *prediction* (what run_ts_align.py writes) carries every request field through plus "prediction" (raw model text, or None on error). Scoring reads predictions. This module is dependency-free (stdlib only) so it imports anywhere. """ from __future__ import annotations import json import re from pathlib import Path from typing import Any, Iterable # ---------------------------------------------------------------------------- # task types (routing keys) # ---------------------------------------------------------------------------- NUMERIC = "numeric" # single scalar -> relative-accuracy / abs-err CHOICE = "choice" # MCQ / T-F / label -> exact-match Acc (+ F1) TEXT = "text" # free-form -> Rouge-L / BLEU (or LLM-judge) SEQUENCE = "sequence" # numeric series -> MSE / MAE TASK_TYPES = (NUMERIC, CHOICE, TEXT, SEQUENCE) # default output-format system prompts per task_type (English; language deferred). SYSTEM_PROMPTS = { NUMERIC: ( "You are a time-series analysis assistant. Use only the given time series " "to answer. Output the final answer strictly as \\boxed{NUMBER} " "with a single number and nothing else." ), CHOICE: ( "You are a time-series analysis assistant. Use only the given time series " "to answer. Choose exactly one option. Output strictly as " "\\boxed{LETTER} with a single option letter and nothing else." ), TEXT: ( "You are a time-series analysis assistant. Use only the given time series " "to answer. Put your final answer inside ...." ), SEQUENCE: ( "You are a time-series analysis assistant. Forecast the requested values. " "Output strictly as v1, v2, v3, ...: comma-separated numbers " "and nothing else." ), } TS_PLACEHOLDER = "" # ---------------------------------------------------------------------------- # jsonl io # ---------------------------------------------------------------------------- def read_jsonl(path: str | Path) -> list[dict]: rows: list[dict] = [] with open(path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if line: rows.append(json.loads(line)) return rows def write_jsonl(path: str | Path, rows: Iterable[dict]) -> int: path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) n = 0 with open(path, "w", encoding="utf-8") as f: for row in rows: f.write(json.dumps(row, ensure_ascii=False) + "\n") n += 1 return n # ---------------------------------------------------------------------------- # request builder (adapters call this to stay schema-correct) # ---------------------------------------------------------------------------- def build_request( *, request_id: str, sample_id: Any, task_type: str, raw_ts: list[list[float]], question: str, gt: Any, benchmark: str, task: str, metric: str, options: list[str] | None = None, split: str = "test", lang: str = "en", extra_meta: dict | None = None, system_prompt: str | None = None, reference_block: str | None = None, prompt_override: str | None = None, ) -> dict: if task_type not in TASK_TYPES: raise ValueError(f"unknown task_type {task_type!r}; expected one of {TASK_TYPES}") # Compose the user prompt. If prompt_override is given (benchmark already has a # complete question text, e.g. TimeOmni's `problem`), use it verbatim and only # ensure a slot exists for MOMENT injection. if prompt_override is not None: prompt = prompt_override.strip() if TS_PLACEHOLDER not in prompt: prompt = f"{prompt}\nTime series: {TS_PLACEHOLDER}" else: parts = [] if reference_block: parts.append(reference_block.strip()) parts.append(question.strip()) if options: parts.append("Options:\n" + "\n".join(options)) parts.append(f"Time series: {TS_PLACEHOLDER}") prompt = "\n".join(parts) meta = { "benchmark": benchmark, "task": task, "task_type": task_type, "metric": metric, "split": split, "lang": lang, "n_channels": len(raw_ts), "seq_len": len(raw_ts[0]) if raw_ts else 0, } if options is not None: meta["options"] = options if extra_meta: meta.update(extra_meta) return { "request_id": request_id, "id": sample_id, "task_type": task_type, "variant": "normal", "raw_ts": raw_ts, "prompt": prompt, "system_prompt": system_prompt or SYSTEM_PROMPTS[task_type], "gt": gt, "meta": meta, } # ---------------------------------------------------------------------------- # answer extraction (from raw model text) # ---------------------------------------------------------------------------- _BOXED_RE = re.compile(r"\\boxed\s*\{([^{}]*)\}") _ANSWER_RE = re.compile(r"\s*(.*?)\s*", re.DOTALL | re.IGNORECASE) _THINK_RE = re.compile(r".*?", re.DOTALL | re.IGNORECASE) _NUM_RE = re.compile(r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?") _LETTER_RE = re.compile(r"\b([A-Ea-e])\b") # MCQ up to 5 options (EngineMT uses a-e) def answer_block(text: str) -> str: """Content of the last .., else the text minus ... Loose: used for TEXT tasks where the whole reply is the answer. """ if not text: return "" matches = _ANSWER_RE.findall(text) if matches: return matches[-1].strip() return _THINK_RE.sub("", text).strip() def answer_block_strict(text: str) -> str | None: """Content of the last .., or None if no answer tag present. Strict: used for choice/numeric/sequence so that SR (format validity) only counts answers emitted in a structured location, never scraped from free prose. Matches TimeOmni's own extract_answer semantics. """ if not text: return None matches = _ANSWER_RE.findall(text) return matches[-1].strip() if matches else None def extract_boxed(text: str) -> str | None: """Content of the last \\boxed{...}, or None.""" if not text: return None m = _BOXED_RE.findall(text) return m[-1].strip() if m else None def _to_float(token: str) -> float | None: token = token.strip().replace(",", "") is_pct = token.endswith("%") if is_pct: token = token[:-1] m = _NUM_RE.search(token) if not m: return None try: val = float(m.group(0)) except ValueError: return None return val def parse_number(text: str) -> float | None: """Strict: \\boxed{} content, else the block. No prose fallback. Returning None means "format-invalid" and counts against SR — we do NOT scrape a number from free reasoning text, which would inflate SR with accidental matches. """ boxed = extract_boxed(text) if boxed is not None: v = _to_float(boxed) if v is not None: return v blk = answer_block_strict(text) if blk is not None: return _to_float(blk) return None def parse_choice(text: str, options: list[str] | None = None) -> str | None: """Strict choice extraction: \\boxed{} or tag only (no prose scraping). Order: \\boxed{} letter/option -> letter/option inside ... Returns normalized letter (upper) or matched option text, else None. """ boxed = extract_boxed(text) if boxed: m = _LETTER_RE.search(boxed) if m: return m.group(1).upper() if options and boxed in options: return boxed blk = answer_block_strict(text) if blk: m = _LETTER_RE.search(blk) if m: return m.group(1).upper() if options: low = blk.lower() for opt in options: if opt and opt.lower() in low: return opt return None def parse_sequence(text: str) -> list[float]: """Strict: floats from \\boxed{} or the block only (no whole-text).""" src = extract_boxed(text) if not src: src = answer_block_strict(text) if not src: return [] return [float(m) for m in _NUM_RE.findall(src)]