scripts/dump_oumi_schema.py

"""Generate data/oumi_training_schema.json by introspecting oumi's config classes.

Emits a JSON Schema for a *curated subset* of oumi's TrainingConfig — the fields
the config-copilot task covers. Field names, types, defaults, and enum values are
read from the real classes so drift from oumi HEAD fails loudly here rather than
silently in the app.

Requires an environment with oumi installed (heavy: torch et al.):

    uv run --with oumi python scripts/dump_oumi_schema.py
"""

import dataclasses
import enum
import json
import types
import typing
from pathlib import Path

from oumi.core.configs import TrainingConfig
from oumi.core.configs.params.data_params import DatasetParams, DatasetSplitParams
from oumi.core.configs.params.model_params import ModelParams
from oumi.core.configs.params.peft_params import PeftParams
from oumi.core.configs.params.training_params import (
    MixedPrecisionDtype,
    TrainerType,
    TrainingParams,
)

OUT_PATH = Path(__file__).resolve().parent.parent / "data" / "oumi_training_schema.json"

# section -> (params class, {field_name: description})
CURATED: dict[str, tuple[type, dict[str, str]]] = {
    "model": (ModelParams, {
        "model_name": "HF Hub id or local path of the base model (required)",
        "model_max_length": "Max sequence length; null lets the model default apply",
        "torch_dtype_str": 'Model dtype, e.g. "auto", "bfloat16", "float16", "float32"',
        "trust_remote_code": "Allow models with custom code from the Hub",
        "chat_template": "Name of the chat template to apply; null uses the tokenizer default",
    }),
    "dataset": (DatasetParams, {
        "dataset_name": "Registered dataset name, HF Hub id, or format name for local files (required)",
        "dataset_path": "Path to a local dataset file (e.g. a .jsonl), if not loading from the Hub",
        "split": 'Dataset split to load, e.g. "train"',
        "sample_count": "Cap the number of examples drawn from this dataset",
    }),
    "dataset_split": (DatasetSplitParams, {
        "pack": "Pack multiple short examples into each sequence",
    }),
    "training": (TrainingParams, {
        "trainer_type": "Which trainer implementation to use",
        "use_peft": "Train with parameter-efficient fine-tuning (LoRA/QLoRA); pair with the peft section",
        "output_dir": "Directory where checkpoints and the final model are written",
        "num_train_epochs": "Number of passes over the training data (ignored if max_steps > 0)",
        "max_steps": "Hard cap on optimizer steps; -1 disables",
        "learning_rate": "Peak learning rate",
        "per_device_train_batch_size": "Micro-batch size per device",
        "gradient_accumulation_steps": "Steps to accumulate before each optimizer update",
        "lr_scheduler_type": 'LR schedule, e.g. "linear", "cosine", "constant"',
        "warmup_steps": "LR warmup steps (alternative to warmup_ratio)",
        "warmup_ratio": "LR warmup as a fraction of total steps",
        "optimizer": 'Optimizer name, e.g. "adamw_torch", "adamw_torch_fused", "sgd"',
        "weight_decay": "Weight decay coefficient",
        "mixed_precision_dtype": "Mixed-precision mode",
        "enable_gradient_checkpointing": "Trade compute for memory during backprop",
        "eval_strategy": '"no", "steps", or "epoch"',
        "eval_steps": "Evaluate every N steps (when eval_strategy=steps)",
        "save_steps": "Checkpoint every N steps",
        "save_final_model": "Save the model at the end of training",
        "logging_steps": "Log metrics every N steps",
        "seed": "Random seed",
        "run_name": "Human-readable name for the run",
    }),
    "peft": (PeftParams, {
        "lora_r": "LoRA rank",
        "lora_alpha": "LoRA scaling alpha",
        "lora_dropout": "Dropout on LoRA layers",
        "lora_target_modules": 'Module names to adapt, e.g. ["q_proj", "v_proj"]; null lets oumi pick',
        "q_lora": "Quantize the base model (QLoRA)",
        "q_lora_bits": "Quantization bits for QLoRA (typically 4)",
    }),
}

JSON_TYPES = {str: "string", int: "integer", float: "number", bool: "boolean"}


def field_schema(f: dataclasses.Field, description: str) -> dict:
    """Map a dataclass field's annotation + default to a JSON Schema fragment."""
    ann, nullable = f.type, False
    if isinstance(ann, str):  # from __future__ annotations
        ann = eval(ann, vars(typing) | {"torch": None}, {})  # noqa: S307 - trusted source
    origin = typing.get_origin(ann)
    if origin in (typing.Union, types.UnionType):
        args = [a for a in typing.get_args(ann) if a is not type(None)]
        nullable = len(args) < len(typing.get_args(ann))
        ann, origin = args[0], typing.get_origin(args[0])

    if isinstance(ann, type) and issubclass(ann, enum.Enum):
        # oumi YAML accepts enum *names* (see configs/recipes) and value-style
        # strings for str-enums; offer names plus str-enum values.
        allowed = sorted({m.name for m in ann} | {
            m.value for m in ann if isinstance(m.value, str)
        })
        schema: dict = {"enum": allowed + ([None] if nullable else [])}
    elif origin is list:
        item = typing.get_args(ann)[0] if typing.get_args(ann) else str
        schema = {"type": ["array", "null"] if nullable else "array",
                  "items": {"type": JSON_TYPES.get(item, "string")}}
    else:
        jtype = JSON_TYPES.get(ann, "string")
        schema = {"type": [jtype, "null"] if nullable else jtype}

    if f.default is not dataclasses.MISSING and f.default is not None:
        default = f.default.name if isinstance(f.default, enum.Enum) else f.default
        if not (isinstance(default, str) and default == "???"):  # omegaconf MISSING
            schema["default"] = default
    schema["description"] = description
    return schema


def section_schema(cls: type, wanted: dict[str, str]) -> dict:
    by_name = {f.name: f for f in dataclasses.fields(cls)}
    missing = set(wanted) - set(by_name)
    if missing:
        raise SystemExit(f"fields gone from {cls.__name__}: {missing} — update CURATED")
    return {
        "type": "object",
        "additionalProperties": False,
        "properties": {n: field_schema(by_name[n], desc) for n, desc in wanted.items()},
    }


def main() -> None:
    sections = {name: section_schema(cls, wanted) for name, (cls, wanted) in CURATED.items()}

    dataset = sections.pop("dataset")
    dataset["required"] = ["dataset_name"]
    split = sections.pop("dataset_split")
    split["properties"]["datasets"] = {
        "type": "array", "minItems": 1, "items": dataset,
        "description": "Datasets to mix for this split",
    }
    split["required"] = ["datasets"]

    sections["model"]["required"] = ["model_name"]
    schema = {
        "$schema": "https://json-schema.org/draft/2020-12/schema",
        "title": "Oumi TrainingConfig (curated subset)",
        "description": (
            "Subset of oumi.core.configs.TrainingConfig supported by the config "
            "copilot. Generated by scripts/dump_oumi_schema.py from oumi "
            f"{__import__('importlib.metadata').metadata.version('oumi')}."
        ),
        "type": "object",
        "additionalProperties": False,
        "required": ["model", "data", "training"],
        "properties": {
            "model": sections["model"],
            "data": {
                "type": "object",
                "additionalProperties": False,
                "required": ["train"],
                "properties": {
                    "train": split,
                    "validation": {**split, "required": []},
                },
            },
            "training": sections["training"],
            "peft": sections["peft"],
        },
    }

    # Sanity: defaults pulled from live classes must round-trip.
    assert "TRL_SFT" in schema["properties"]["training"]["properties"]["trainer_type"]["enum"]
    assert {m.value for m in MixedPrecisionDtype} <= set(
        schema["properties"]["training"]["properties"]["mixed_precision_dtype"]["enum"]
    )
    assert TrainerType and TrainingConfig  # imported for drift-checking

    OUT_PATH.write_text(json.dumps(schema, indent=2) + "\n")
    print(f"wrote {OUT_PATH}")


if __name__ == "__main__":
    main()