typedecide.training.config

TrainConfig and the fixed policy constants that are written into every manifest. Field-by-field guide: Configure training.

typedecide.training.config

The full description of a fine-tuning run, in one frozen object.

Everything that changes the weights lives here, so TrainConfig.to_manifest() plus the base model id and the dataset fingerprint is enough to re-run a job. Values that are really fixed policy rather than knobs (the LoRA target modules, the schedule) are module constants, and they are still written into the manifest -- a manifest that only records the fields someone happened to expose is a manifest you cannot reproduce from.

LORA_TARGET_MODULES module-attribute

LORA_TARGET_MODULES: tuple[str, ...] = (
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
)

LR_SCHEDULER_TYPE module-attribute

LR_SCHEDULER_TYPE = 'cosine'

WARMUP_RATIO module-attribute

WARMUP_RATIO = 0.03

LOGGING_STEPS module-attribute

LOGGING_STEPS = 25

MIN_MAX_LENGTH module-attribute

MIN_MAX_LENGTH = 16

TrainConfig dataclass

TrainConfig(
    base_model: str,
    output_dir: Path,
    epochs: float = 2.0,
    batch_size: int = 4,
    grad_accum: int = 4,
    learning_rate: float = 0.0001,
    lora_rank: int = 16,
    lora_alpha: int | None = None,
    lora_dropout: float = 0.05,
    max_length: int = 1024,
    seed: int = 0,
    randomise_option_order: bool = True,
)

One fine-tuning job.

output_dir receives the LoRA adapter, the tokenizer and manifest.json. lora_alpha defaults to twice the rank when left unset; read the resolved value from effective_lora_alpha, never from the field.

base_model instance-attribute

base_model: str

output_dir instance-attribute

output_dir: Path

epochs class-attribute instance-attribute

epochs: float = 2.0

batch_size class-attribute instance-attribute

batch_size: int = 4

grad_accum class-attribute instance-attribute

grad_accum: int = 4

learning_rate class-attribute instance-attribute

learning_rate: float = 0.0001

lora_rank class-attribute instance-attribute

lora_rank: int = 16

lora_alpha class-attribute instance-attribute

lora_alpha: int | None = None

lora_dropout class-attribute instance-attribute

lora_dropout: float = 0.05

max_length class-attribute instance-attribute

max_length: int = 1024

seed class-attribute instance-attribute

seed: int = 0

randomise_option_order class-attribute instance-attribute

randomise_option_order: bool = True

effective_lora_alpha property

effective_lora_alpha: int

The alpha actually handed to peft: twice the rank unless one was given.

effective_batch_size property

effective_batch_size: int

Decisions per optimiser step, which is the number that shapes the run.

from_file classmethod

from_file(path: str | Path) -> TrainConfig

Read a config from YAML or JSON.

The format is taken from the suffix; anything else is tried as JSON and then as YAML. A relative output_dir is resolved against the config file's own directory, so a config can be moved with its run without being rewritten.

Source code in src/typedecide/training/config.py

@classmethod
def from_file(cls, path: str | Path) -> TrainConfig:
    """Read a config from YAML or JSON.

    The format is taken from the suffix; anything else is tried as JSON and then as
    YAML. A relative `output_dir` is resolved against the config file's own
    directory, so a config can be moved with its run without being rewritten.
    """
    path = Path(path)
    try:
        text = path.read_text(encoding="utf-8-sig")
    except OSError as error:
        raise ConfigError(
            f"Could not read the training config at {path}: {error}"
        ) from error
    except UnicodeDecodeError as error:
        raise ConfigError(
            f"The training config at {path} is not UTF-8 ({error.reason} at byte "
            f"{error.start}). Save it as UTF-8 and retry."
        ) from error

    suffix = path.suffix.lower()
    if suffix == ".json":
        payload = _load_json(text, path)
    elif suffix in {".yaml", ".yml"}:
        payload = _load_yaml(text, path)
    else:
        try:
            payload = _load_json(text, path)
        except ConfigError:
            payload = _load_yaml(text, path)

    if not isinstance(payload, Mapping):
        raise ConfigError(
            f"{path} must hold a mapping of config keys, not {type(payload).__name__}."
        )
    return cls.from_mapping(payload, base_dir=path.parent)

from_mapping classmethod

from_mapping(
    payload: Mapping[str, Any],
    *,
    base_dir: Path | None = None
) -> TrainConfig

Build a config from an already-parsed mapping, naming every unusable key.

Source code in src/typedecide/training/config.py

@classmethod
def from_mapping(
    cls, payload: Mapping[str, Any], *, base_dir: Path | None = None
) -> TrainConfig:
    """Build a config from an already-parsed mapping, naming every unusable key."""
    known = {f.name for f in fields(cls)}
    # str() first: YAML happily produces int, bool or None keys, and sorting or
    # joining a mixed set raises a TypeError instead of the diagnosis below.
    unknown = sorted({str(key) for key in payload} - known)
    if unknown:
        raise ConfigError(
            f"Unknown training config key(s): {', '.join(unknown)}. "
            f"Permitted keys are: {', '.join(sorted(known))}."
        )
    missing = sorted({"base_model", "output_dir"} - set(payload))
    if missing:
        raise ConfigError(
            f"The training config is missing {', '.join(missing)}; "
            "both the base model and the output directory are required."
        )

    output_dir = Path(str(payload["output_dir"])).expanduser()
    if base_dir is not None and not output_dir.is_absolute():
        output_dir = base_dir / output_dir

    return cls(
        base_model=str(payload["base_model"]),
        output_dir=output_dir,
        epochs=_as_float(payload, "epochs", 2.0),
        batch_size=_as_int(payload, "batch_size", 4),
        grad_accum=_as_int(payload, "grad_accum", 4),
        learning_rate=_as_float(payload, "learning_rate", 1e-4),
        lora_rank=_as_int(payload, "lora_rank", 16),
        lora_alpha=None if payload.get("lora_alpha") is None
        else _as_int(payload, "lora_alpha", 0),
        lora_dropout=_as_float(payload, "lora_dropout", 0.05),
        max_length=_as_int(payload, "max_length", 1024),
        seed=_as_int(payload, "seed", 0),
        randomise_option_order=_as_bool(payload, "randomise_option_order", True),
    )

to_manifest

to_manifest() -> dict[str, Any]

A JSON-safe, fully resolved record of this config.

Resolved means no None standing for "work it out later": lora_alpha is the number peft receives, and the fixed policy constants are included so a manifest written by one release still describes the run if a later one changes them.

Source code in src/typedecide/training/config.py

def to_manifest(self) -> dict[str, Any]:
    """A JSON-safe, fully resolved record of this config.

    Resolved means no `None` standing for "work it out later": `lora_alpha` is the
    number peft receives, and the fixed policy constants are included so a manifest
    written by one release still describes the run if a later one changes them.
    """
    return {
        "base_model": self.base_model,
        "output_dir": str(self.output_dir),
        "epochs": float(self.epochs),
        "batch_size": int(self.batch_size),
        "grad_accum": int(self.grad_accum),
        "effective_batch_size": self.effective_batch_size,
        "learning_rate": float(self.learning_rate),
        "lora_rank": int(self.lora_rank),
        "lora_alpha": self.effective_lora_alpha,
        "lora_alpha_was_explicit": self.lora_alpha is not None,
        "lora_dropout": float(self.lora_dropout),
        "lora_target_modules": list(LORA_TARGET_MODULES),
        "max_length": int(self.max_length),
        "seed": int(self.seed),
        "randomise_option_order": bool(self.randomise_option_order),
        "lr_scheduler_type": LR_SCHEDULER_TYPE,
        "warmup_ratio": WARMUP_RATIO,
        "logging_steps": LOGGING_STEPS,
    }