from __future__ import annotations import json import sys from collections.abc import Iterable from pathlib import Path from typing import Any EVAL_DIR = Path(__file__).resolve().parent MULTIVAR_DIR = EVAL_DIR.parent REPO_ROOT = MULTIVAR_DIR.parent SRC_DIR = MULTIVAR_DIR / "src" INFER_DIR = MULTIVAR_DIR / "infer" DEFAULT_DATA_DIR = MULTIVAR_DIR / "train" / "data_multivar" / "multivar" def ensure_src_path() -> None: src = str(SRC_DIR) if src not in sys.path: sys.path.insert(0, src) def ensure_infer_path(path: str | Path | None = None) -> None: infer = str(Path(path) if path is not None else INFER_DIR) if infer not in sys.path: sys.path.insert(0, infer) def read_jsonl(path: str | Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with Path(path).open("r", encoding="utf-8") as handle: for line_number, raw_line in enumerate(handle, start=1): line = raw_line.strip() if not line: continue try: row = json.loads(line) except json.JSONDecodeError as exc: raise ValueError(f"Invalid JSON on line {line_number} of {path}.") from exc if not isinstance(row, dict): raise ValueError(f"Expected JSON object on line {line_number} of {path}.") rows.append(row) return rows def write_jsonl(path: str | Path, rows: Iterable[dict[str, Any]]) -> None: output_path = Path(path) output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False, default=json_default) + "\n") def append_jsonl(path: str | Path, row: dict[str, Any]) -> None: output_path = Path(path) output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(row, ensure_ascii=False, default=json_default) + "\n") def write_json(path: str | Path, payload: dict[str, Any]) -> None: output_path = Path(path) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text( json.dumps(payload, ensure_ascii=False, indent=2, default=json_default), encoding="utf-8", ) def json_default(obj: Any) -> Any: if hasattr(obj, "item"): return obj.item() if hasattr(obj, "tolist"): return obj.tolist() if isinstance(obj, Path): return str(obj) raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable.") def parse_csv_arg(raw: str | None) -> list[str] | None: if raw is None: return None values = [item.strip() for item in raw.split(",") if item.strip()] return values or None