Benchmark_eval/ChatTS_test/_build/schema.py

"""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 "<ts></ts>" 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 <answer>\\boxed{NUMBER}</answer> "
        "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 "
        "<answer>\\boxed{LETTER}</answer> 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 <answer>...</answer>."
    ),
    SEQUENCE: (
        "You are a time-series analysis assistant. Forecast the requested values. "
        "Output strictly as <answer>v1, v2, v3, ...</answer>: comma-separated numbers "
        "and nothing else."
    ),
}

TS_PLACEHOLDER = "<ts></ts>"


# ----------------------------------------------------------------------------
# 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 <ts></ts> 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"<answer>\s*(.*?)\s*</answer>", re.DOTALL | re.IGNORECASE)
_THINK_RE = re.compile(r"<think>.*?</think>", 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 <answer>..</answer>, else the text minus <think>..</think>.

    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 <answer>..</answer>, 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 <answer> 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 <answer> tag only (no prose scraping).

    Order: \\boxed{} letter/option -> letter/option inside <answer>..</answer>.
    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 <answer> 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)]