multimodal_data_processing/rats_uni.py

# -*- coding: utf-8 -*-
"""
Convert the Time-RA RATs-Uni (univariate) reasoning dataset into the parquet
format consumed by AnomSeer's veRL training/eval pipeline.

Unlike AnomLLM (see ``anom.py``), the Time-RA univariate dataset does **not**
provide anomaly *intervals* — it provides, per univariate series:

  * ``Observation``  : the raw time series (list of floats, length 32/64/128)
  * ``FigurePath``   : a rendered plot (``figures/{train,test}/{idx}.jpg``)
  * ``ActionID``     : the anomaly class id (0-14, 0 == normal)
  * ``Action``       : the human-readable class name
  * ``Label``        : "Normal" / "Anomaly"
  * ``Thought``      : an expert chain-of-thought explanation

We therefore adapt AnomSeer to a **classification + reasoning** task (no
localization). Each output row matches the schema expected by
``verl.utils.dataset.rl_dataset.RLHFDataset`` and is scored by the
``timeseries_rats`` reward (``verl/utils/reward_score/rats.py``):

  data_source  : "timeseries_rats"
  prompt       : [{"role": "user", "content": "<image>\\n ... <class> ..."}]
  images       : [{"bytes": <jpg bytes>, "path": <abs path>}]
  ability      : "time_series_anomaly_detection"
  reward_model : {"style": "rule", "ground_truth": <canonical class name>}
  extra_info   : {category, source, split, anomaly_type, action_id, label,
                  series_length, image_path, expcot, numtext, index}

The ``numtext`` field is the expert CoT and is **required** by TimerPO's
semantic-alignment term (``compute_hidden_states_of_hint`` in fsdp_workers.py).

Usage
-----
    python multimodal_data_processing/rats_uni.py \
        --json_path /path/to/RATs40K/RATs-Uni-TSImage_Reason.json \
        --out_dir   ./data/rats_uni_processed

Produces ``train_full.parquet`` and ``test_full.parquet`` in ``--out_dir``.
"""

import os
import io
import json
import argparse
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
from PIL import Image


# Canonical 15-class taxonomy (id -> (name, description)), matching
# Time-RA/eval_utils.py::ANOMALY_DICT exactly.
ANOMALY_DICT: Dict[int, Tuple[str, str]] = {
    0:  ("Normal Sequence", "There are no abnormal situations in this time series."),
    1:  ("Point Anomaly", "A single data point significantly deviates from the local or global pattern of the sequence."),
    2:  ("Periodic Change Anomaly", "The original periodic pattern is disrupted, e.g. the period is broken or the amplitude becomes anomalous."),
    3:  ("Trend Change Anomaly", "A sudden change in the long-term trend of the time series."),
    4:  ("Change Point Anomaly", "Statistical properties (e.g. mean, variance) change abruptly at certain points."),
    5:  ("Distributional Change Anomaly", "The statistical distribution of the time series changes significantly."),
    6:  ("Amplitude Anomaly", "The amplitude of data points exceeds the normal upper and lower bounds."),
    7:  ("Pattern Change Anomaly", "The pattern of the time series suddenly changes from one form to another."),
    8:  ("Sparse Anomaly", "Isolated anomalous patterns occasionally appear in a long time series."),
    9:  ("Repeated Value Anomaly", "Continuous or intermittent repeated values disrupt the normal fluctuation pattern."),
    10: ("Sudden Flatline Anomaly", "The time series suddenly becomes a flat line with no normal fluctuations."),
    11: ("Drift Anomaly", "The data gradually drifts away from the normal level."),
    12: ("Sudden Spike Anomaly", "The data suddenly spikes or drops within a short time and then returns to normal."),
    13: ("Continuous Segment Anomaly", "A continuous segment of data points deviates from the normal pattern."),
    14: ("Nonlinear Pattern Anomaly", "Nonlinear changes appear in the sequence, breaking the original linear rule."),
}

NAME_BY_ID = {i: name for i, (name, _) in ANOMALY_DICT.items()}
ID_BY_NAME = {name.lower(): i for i, (name, _) in ANOMALY_DICT.items()}

# Splits in the JSON -> output split name used by AnomSeer ("eval" == val/test).
SPLIT_MAP = {"TSAD_train": "train", "TSAD_test": "eval"}


def _build_taxonomy_block() -> str:
    lines = []
    for i in range(len(ANOMALY_DICT)):
        name, desc = ANOMALY_DICT[i]
        lines.append(f"{i}: {name} — {desc}")
    return "\n".join(lines)


_TAXONOMY_BLOCK = _build_taxonomy_block()
_CLASS_NAMES = ", ".join(NAME_BY_ID[i] for i in range(len(NAME_BY_ID)))


def build_prompt(length: int, source: str) -> str:
    """Classification + reasoning prompt (no interval localization)."""
    return (
        "<image>\n"
        f"You are an expert in univariate time-series anomaly detection. The figure shows a "
        f"single-channel time series of length {length} from the \"{source}\" domain.\n\n"
        "Decide whether the series is normal or contains an anomaly. If it is anomalous, choose the "
        "single most appropriate anomaly type from the following 15 categories "
        "(format `id: name — description`):\n"
        f"{_TAXONOMY_BLOCK}\n\n"
        "Reason step by step inside <think>...</think> based on the visual shape of the series, "
        "then output exactly one line with your final answer:\n"
        "<class>one exact category name from the list above</class>\n"
        "If the series is normal, use <class>Normal Sequence</class>."
    )


def _canonical_class(action: Optional[str], action_id: Optional[int]) -> Tuple[str, int]:
    """Resolve the canonical (name, id) from the raw Action / ActionID fields.

    ActionID is treated as authoritative; Action is used only as a fallback.
    """
    if isinstance(action_id, (int, np.integer)) and int(action_id) in NAME_BY_ID:
        cid = int(action_id)
        return NAME_BY_ID[cid], cid
    if isinstance(action, str) and action.strip().lower() in ID_BY_NAME:
        cid = ID_BY_NAME[action.strip().lower()]
        return NAME_BY_ID[cid], cid
    # Unknown -> treat as normal (id 0) so downstream typing stays consistent.
    return NAME_BY_ID[0], 0


def _read_image_bytes(path: str) -> Optional[bytes]:
    """Return raw image bytes, re-encoding to a clean RGB JPEG if needed."""
    if not os.path.isfile(path):
        return None
    try:
        with open(path, "rb") as f:
            raw = f.read()
        # Validate; re-encode non-RGB to avoid downstream surprises.
        with Image.open(io.BytesIO(raw)) as im:
            if im.mode == "RGB":
                return raw
            buf = io.BytesIO()
            im.convert("RGB").save(buf, format="JPEG", quality=95)
            return buf.getvalue()
    except Exception as exc:  # noqa: BLE001
        print(f"[WARN] failed to read image {path}: {exc}")
        return None


# Some test entries point FigurePath at a ``.pdf`` that was never exported; the
# rendered raster lives next to it as a same-index ``.jpg``. Fall back across
# common raster extensions before giving up.
_IMG_EXT_FALLBACKS = (".jpg", ".png", ".jpeg")


def _resolve_figure_path(data_root: str, figure_path: str) -> str:
    """FigurePath is stored relative to the dataset root (e.g. figures/train/0.jpg).

    Returns the first existing file, trying the stored path first and then the
    same stem with a raster extension (handles ``.pdf`` FigurePaths)."""
    base = figure_path if os.path.isabs(figure_path) else os.path.join(data_root, figure_path)
    if os.path.isfile(base):
        return base
    stem, _ = os.path.splitext(base)
    for ext in _IMG_EXT_FALLBACKS:
        cand = stem + ext
        if os.path.isfile(cand):
            return cand
    return base  # non-existent; caller reports it as a skip


def process_split(records: dict, split_name: str, data_root: str,
                  max_samples: Optional[int] = None) -> List[dict]:
    rows: List[dict] = []
    n_skipped_none = 0
    n_skipped_img = 0

    # Iterate in numeric key order for reproducibility.
    keys = sorted(records.keys(), key=lambda k: int(k) if str(k).isdigit() else k)
    for out_idx, key in enumerate(keys):
        if max_samples is not None and len(rows) >= max_samples:
            break

        entry = records[key]
        if entry is None:  # the JSON contains a handful of null test entries
            n_skipped_none += 1
            continue

        obs = entry.get("Observation") or []
        length = len(obs)
        source = entry.get("Source", "unknown")
        figure_path = entry.get("FigurePath", "")
        thought = entry.get("Thought", "") or ""
        label = entry.get("Label", "Anomaly")

        img_path = _resolve_figure_path(data_root, figure_path)
        img_bytes = _read_image_bytes(img_path)
        if img_bytes is None:
            n_skipped_img += 1
            continue

        class_name, class_id = _canonical_class(entry.get("Action"), entry.get("ActionID"))
        prompt_text = build_prompt(length, source)

        rows.append({
            "data_source": "timeseries_rats",
            "prompt": [{"role": "user", "content": prompt_text}],
            "images": [{"bytes": img_bytes, "path": img_path}],
            "ability": "time_series_anomaly_detection",
            "reward_model": {"style": "rule", "ground_truth": class_name},
            "extra_info": {
                "index": int(key) if str(key).isdigit() else out_idx,
                "category": class_name,          # per-class metric grouping
                "source": source,                # domain grouping
                "split": split_name,
                "anomaly_type": class_name,       # read by the reward as GT class
                "action_id": int(class_id),
                "label": label,
                "series_length": int(length),
                "image_path": img_path,
                "expcot": thought,                # expert CoT
                "numtext": thought,               # required by TimerPO hint term
            },
        })

    print(f"  [{split_name}] kept={len(rows)} "
          f"skipped(null)={n_skipped_none} skipped(missing-image)={n_skipped_img}")
    return rows


def write_parquet(rows: List[dict], path: str) -> None:
    if not rows:
        print(f"  [skip] no rows for {path}")
        return
    df = pd.DataFrame(rows)
    df.to_parquet(path, index=False, engine="pyarrow")
    print(f"  ✅ wrote {len(df)} rows -> {path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Convert Time-RA RATs-Uni JSON into AnomSeer parquet (classification + reasoning).")
    parser.add_argument("--json_path", type=str,
                        default="/mnt/share01/sqk/datasets/RATs40K/RATs-Uni-TSImage_Reason.json",
                        help="Path to RATs-Uni-TSImage_Reason.json")
    parser.add_argument("--data_root", type=str, default=None,
                        help="Dataset root that FigurePath is relative to "
                             "(defaults to the directory of --json_path).")
    parser.add_argument("--out_dir", type=str, default="./data/rats_uni_processed",
                        help="Output directory for the parquet files.")
    parser.add_argument("--max_samples", type=int, default=None,
                        help="Optional cap on samples per split (for quick tests).")
    args = parser.parse_args()

    data_root = args.data_root or os.path.dirname(os.path.abspath(args.json_path))
    os.makedirs(args.out_dir, exist_ok=True)

    print(f"Reading {args.json_path}")
    print(f"Figure root: {data_root}")
    with open(args.json_path, "r") as f:
        data = json.load(f)

    for json_split, out_split in SPLIT_MAP.items():
        if json_split not in data:
            print(f"[WARN] split '{json_split}' not found in JSON; skipping.")
            continue
        print(f"\nProcessing {json_split} -> {out_split}")
        rows = process_split(data[json_split], out_split, data_root, args.max_samples)
        fname = "train_full.parquet" if out_split == "train" else "test_full.parquet"
        write_parquet(rows, os.path.join(args.out_dir, fname))

    print("\nDone.")


if __name__ == "__main__":
    main()