simulate_inputs.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Resolve dataset and algo paths for simulate."""
from __future__ import annotations
 
import csv
import json
import os
from pathlib import Path
from typing import Any
 
_HINT_KEYS = {
    "rank", "alpha", "lr", "learning_rate", "epochs", "num_train_epochs",
    "dropout", "target_modules", "targetModules", "warmup_steps", "warmupSteps",
    "grad_clip", "gradClip", "weight_decay", "weightDecay", "grad_accum_steps",
    "gradAccumSteps", "optimizer", "scheduler", "max_seq_len", "maxSeqLen",
    "use_qlora", "useQlora", "use_rlhf", "rlhf_beta",
}
_KEY_ALIASES = {
    "learning_rate": "lr",
    "num_train_epochs": "epochs",
    "target_modules": "targetModules",
    "warmup_steps": "warmupSteps",
    "grad_clip": "gradClip",
    "weight_decay": "weightDecay",
    "grad_accum_steps": "gradAccumSteps",
    "max_seq_len": "maxSeqLen",
    "use_qlora": "useQlora",
}
 
_INSTR_FIELDS = ("instruction", "prompt", "input", "query", "question", "text")
_RESP_FIELDS = ("response", "output", "answer", "completion", "target", "label")
 
 
def _expand(path: str) -> Path:
    return Path(os.path.expanduser(path)).resolve()
 
 
def _normalize_row(row: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(row, dict):
        raise ValueError(f"Each dataset row must be an object, got {type(row).__name__}")
    if row.get("instruction") and row.get("response"):
        return row
    instr = next((row[k] for k in _INSTR_FIELDS if row.get(k)), None)
    resp = next((row[k] for k in _RESP_FIELDS if row.get(k)), None)
    if instr is not None and resp is not None:
        out = dict(row)
        out.setdefault("instruction", str(instr))
        out.setdefault("response", str(resp))
        return out
    return row
 
 
def load_dataset_file(path: str) -> list[dict]:
    """Load rows from the user's dataset file (.json, .jsonl, .csv)."""
    p = _expand(path)
    if not p.is_file():
        raise FileNotFoundError(f"Dataset not found: {p}")
 
    suffix = p.suffix.lower()
    if suffix == ".jsonl":
        rows: list[dict] = []
        with p.open(encoding="utf-8") as f:
            for i, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    rows.append(_normalize_row(json.loads(line)))
                except json.JSONDecodeError as e:
                    raise ValueError(f"Invalid JSON on line {i} of {p}: {e}") from e
        if not rows:
            raise ValueError(f"Dataset file is empty: {p}")
        return rows
 
    if suffix == ".json":
        with p.open(encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, list):
            rows = [_normalize_row(r) for r in data]
        elif isinstance(data, dict):
            for key in ("rows", "data", "dataset", "examples", "train"):
                if isinstance(data.get(key), list):
                    rows = [_normalize_row(r) for r in data[key]]
                    break
            else:
                raise ValueError(f"Unrecognized JSON dataset layout: {p}")
        else:
            raise ValueError(f"JSON dataset must be an array or object: {p}")
        if not rows:
            raise ValueError(f"Dataset file is empty: {p}")
        return rows
 
    if suffix == ".csv":
        with p.open(encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            rows = [_normalize_row(dict(r)) for r in reader]
        if not rows:
            raise ValueError(f"CSV dataset is empty: {p}")
        return rows
 
    raise ValueError(
        f"Unsupported dataset format {suffix!r} at {p} — use your pipeline's .json, .jsonl, or .csv export."
    )
 
 
def _parse_structured_file(p: Path) -> Any | None:
    suffix = p.suffix.lower()
    if suffix in (".py", ".sh", ".bash"):
        return None
    text = p.read_text(encoding="utf-8")
    if suffix in (".yaml", ".yml"):
        try:
            import yaml  # type: ignore
        except ImportError:
            return None
        return yaml.safe_load(text)
    if suffix == ".json":
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return None
    return None
 
 
def _collect_hints(obj: Any, found: dict[str, Any], depth: int = 0) -> None:
    if depth > 4 or obj is None:
        return
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k in _HINT_KEYS and k not in found:
                key = _KEY_ALIASES.get(k, k)
                found[key] = v
            elif isinstance(v, (dict, list)):
                _collect_hints(v, found, depth + 1)
    elif isinstance(obj, list):
        for item in obj[:20]:
            _collect_hints(item, found, depth + 1)
 
 
def _try_training_hints(path: Path) -> dict[str, Any]:
    """Best-effort scrape of rank/lr/etc. from the user's existing config — never required."""
    data = _parse_structured_file(path)
    if data is None:
        return {}
    found: dict[str, Any] = {}
    _collect_hints(data, found)
    return found
 
 
def prepare_simulation_args(args: dict) -> dict:
    """
    Resolve --dataset and --algo paths from the user's existing pipeline.
    No Aquin-specific config format — paths only. Optional CLI flags override hints.
    """
    out = dict(args)
 
    algo_path = out.pop("algo_path", None) or out.get("algo")
    if isinstance(algo_path, str):
        p = _expand(algo_path)
        if not p.is_file():
            raise FileNotFoundError(f"Algo path not found: {p}")
        out["algo_path"] = str(p)
        out.pop("algo", None)
        for k, v in _try_training_hints(p).items():
            out.setdefault(k, v)
 
    dataset_ref = out.get("dataset")
    if isinstance(dataset_ref, str):
        p = _expand(dataset_ref)
        if not p.is_file():
            raise FileNotFoundError(f"Dataset not found: {p}")
        out["rows"] = load_dataset_file(str(p))
        out["dataset_path"] = str(p)
        out.pop("dataset", None)
    elif isinstance(dataset_ref, list):
        out["rows"] = dataset_ref
        out.pop("dataset", None)
 
    return out