"""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()