"""Canonical, versioned configuration schema for pstrain.
This module is the only place semantic configuration fields and their defaults
are declared. Runtime dataclasses are projections of :class:`Profile`.
"""
from __future__ import annotations
from typing import Annotated, Any, Literal
from pydantic import (
BaseModel,
ConfigDict,
Field,
PlainValidator,
TypeAdapter,
ValidationError,
WithJsonSchema,
field_validator,
model_validator,
)
from pstrain.lib.retry_ladder import validate_retry_ladder
CURRENT_CONFIG_VERSION: Literal[1] = 1
_RETRY_FACTOR_NUMBER: TypeAdapter[float] = TypeAdapter(Annotated[float, Field(gt=0)])
_RETRY_FACTOR_LIST: TypeAdapter[list[float]] = TypeAdapter(list[float])
def _one_error(error: ValidationError) -> ValueError:
detail = error.errors()[0]
location = "".join(f"[{part}]" for part in detail["loc"])
return ValueError(f"{location}: {detail['msg']}" if location else detail["msg"])
def _validate_retry_beam_factor(value: Any) -> float | list[float]:
"""Validate a number as one retry factor, or a list as a ladder, with one error.
A plain union would try both shapes and report both failures, so a bad
list would also be told it is not a number. This picks the shape from the
input and reports only that shape's problem.
"""
if isinstance(value, list | tuple):
try:
factors = _RETRY_FACTOR_LIST.validate_python(list(value))
except ValidationError as error:
raise _one_error(error) from None
return list(validate_retry_ladder(factors))
try:
return _RETRY_FACTOR_NUMBER.validate_python(value)
except ValidationError as error:
raise _one_error(error) from None
# One number keeps its long-standing meaning, a single retry at that factor; an
# ascending list of factors, each above 1 and each relative to the nominal beam,
# is a ladder whose rungs run in order until one succeeds. JSON Schema cannot say
# "ascending", so validation adds that to what the schema states.
RetryBeamFactorSetting = Annotated[
float | list[float],
PlainValidator(_validate_retry_beam_factor),
WithJsonSchema(
{
"anyOf": [
{"type": "number", "exclusiveMinimum": 0},
{
"type": "array",
"items": {"type": "number", "exclusiveMinimum": 1},
"minItems": 1,
},
]
}
),
]
class StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
[docs]
class FeatureConfig(StrictModel):
"""Acoustic front-end parameters."""
samprate: Annotated[int, Field(gt=0, description="Audio sample rate in Hz")] = 16000
ncep: Annotated[int, Field(gt=0, description="Number of cepstral coefficients")] = 13
nfilt: Annotated[int, Field(gt=0, description="Number of mel filters")] = 25
nfft: Annotated[int, Field(gt=0, description="FFT size")] = 512
lowerf: Annotated[float, Field(ge=0, description="Lower filter-bank frequency in Hz")] = 130.0
upperf: Annotated[float, Field(gt=0, description="Upper filter-bank frequency in Hz")] = 6800.0
alpha: Annotated[float, Field(description="Pre-emphasis coefficient")] = 0.97
dither: Annotated[bool, Field(description="Add half-bit dither to input audio")] = True
seed: Annotated[int, Field(description="Seed for deterministic input dithering")] = -1
remove_dc: Annotated[bool, Field(description="Remove DC offset from each frame")] = True
remove_noise: Annotated[bool, Field(description="Remove noise with spectral subtraction")] = (
True
)
frate: Annotated[int, Field(gt=0, description="Feature frame rate in Hz")] = 100
wlen: Annotated[float, Field(gt=0, description="Analysis window length in seconds")] = 0.025625
feat_type: Annotated[str, Field(description="Sphinx feature stream type")] = "1s_c_d_dd"
lifter: Annotated[int, Field(ge=0, description="Cepstral lifter window")] = 22
transform: Annotated[str, Field(description="Filter-bank transform")] = "dct"
agc: Annotated[str, Field(description="Automatic gain-control mode")] = "none"
# Keep the profile default at batch: training and decoder evaluation both
# consume this field, so it preserves the established matched front end.
# Vendored cmn.c computes one mean from the complete utterance (skipping
# negative-c0 frames) and subtracts it from every frame; "current" is an
# exact parser alias for the same native mode. A batch/current output
# comparison is therefore tautological and supplies no evidence for this
# default or the native mode's behavior. There is no evidence here for
# changing this compatibility default.
cmn: Annotated[str, Field(description="Cepstral mean-normalization mode")] = "batch"
cmninit: Annotated[str, Field(description="Initial cepstral mean vector for live CMN")] = (
"40,3,-1"
)
varnorm: Annotated[str, Field(description="Cepstral variance-normalization mode")] = "no"
[docs]
@model_validator(mode="after")
def validate_band(self) -> FeatureConfig:
if self.upperf <= self.lowerf:
raise ValueError("upperf must be greater than lowerf")
return self
[docs]
class TrainingScheduleConfig(StrictModel):
"""Convergence controller for one Baum-Welch stage family."""
max_iterations: Annotated[int, Field(ge=1, description="Maximum training passes")] = 10
min_iterations: Annotated[int, Field(ge=1, description="Minimum training passes")] = 1
convergence_ratio: Annotated[
float,
Field(
gt=0,
description=(
"Converge after min_iterations when the finite per-frame log-likelihood "
"increase is between zero and this many nats, inclusive. Negative or "
"nonfinite changes do not indicate convergence. Despite the name -- kept because SphinxTrain's "
"$CFG_CONVERGENCE_RATIO is the same signed per-frame delta -- this is an "
"absolute difference, not a ratio. At the default, corpora of Arctic's size "
"run all ten passes in every schedule, which is the more accurate outcome as "
"measured; treat max_iterations as the operative control. The sphinxtrain "
"profile carries SphinxTrain's own 0.1"
),
),
] = 0.001
[docs]
@model_validator(mode="after")
def validate_iterations(self) -> TrainingScheduleConfig:
if self.min_iterations > self.max_iterations:
raise ValueError("min_iterations must not exceed max_iterations")
return self
[docs]
class TrainingConfig(StrictModel):
"""Acoustic-model training parameters."""
n_state: Annotated[int, Field(ge=1, description="Emitting states per HMM")] = 3
skip_state: Annotated[
bool,
Field(
description=(
"Enable SphinxTrain's $CFG_SKIPSTATE topology, adding an arc from each "
"eligible emitting state to the state two positions ahead so a phone can be "
"realized with fewer frames than states. SphinxTrain writes raw 3/1/1 "
"weights and normalizes them on read; pstrain writes the behaviorally "
"equivalent normalized values"
)
),
] = False
n_senones: Annotated[int, Field(ge=1, description="Target tied-state count")] = 200
a_beam: Annotated[float, Field(gt=0, description="Forward alignment beam")] = 1e-90
b_beam: Annotated[float, Field(gt=0, description="Backward alignment beam")] = 1e-10
ci: TrainingScheduleConfig = Field(default_factory=TrainingScheduleConfig)
tied: TrainingScheduleConfig = Field(default_factory=TrainingScheduleConfig)
untied: TrainingScheduleConfig = Field(default_factory=TrainingScheduleConfig)
split_variance_floor_fraction: Annotated[
float,
Field(
ge=0,
le=1,
allow_inf_nan=False,
strict=True,
description=(
"Experimental variance lower bound for split training stages, as a fraction "
"of each matching coordinate in the fixed CI-1g or CD-1g variance reference. "
"Zero disables regularization. The reference never advances with later splits; "
"zero reference coordinates contribute no positive floor. Select a nonzero "
"fraction explicitly after evaluation; no universal nonzero value is assumed"
),
),
] = 0.0
max_skip_fraction: Annotated[
float, Field(ge=0, le=1, description="Maximum skipped-update fraction")
] = 0.05
retry_beam_factor: Annotated[
RetryBeamFactorSetting,
Field(
description=(
"Factor that widens the forward beam for one retry after an utterance fails "
"to reach its final state, or an ascending list of factors, each greater than "
"1 and relative to the nominal beam, tried in order until one succeeds; a retry "
"is counted only when that attempt runs"
),
),
] = 1e10
failed_alignment: Annotated[
Literal["recover", "abort", "omit"],
Field(
description=(
"Action when an utterance fails to reach its final state: ``recover`` runs the "
"wider-beam retries in ``retry_beam_factor`` and, if they all fail, reports the "
"utterance and continues without it; ``abort`` fails the run on the first failure; and ``omit`` reports "
"and excludes it without retrying. Skips are counted either way, and "
"``max_skip_fraction`` still fails the run when they stop being incidental"
)
),
] = "recover"
bw_checkpoint_iterations: Annotated[
bool,
Field(
description=(
"Retain the compact model files from every completed Baum-Welch pass under "
"``iterations/NN``; costs roughly one additional model copy per pass and does "
"not retain the much larger ``.bw-accum`` shard accumulators or change which "
"checkpoint is loaded by training. The deprecated "
"``PSTRAIN_BW_CHECKPOINTS=1`` environment variable can also enable retention, "
"but cannot disable a true profile setting"
)
),
] = False
arctic_a0302_zero_codebook_band: Annotated[
tuple[int, int] | None,
Field(
description=(
"Accepted inclusive exact-zero codebook occupancy band for the singular "
"Arctic a0302 terminal-alignment exception"
)
),
] = None
accept_arctic_a0587_known_skip: Annotated[
bool,
Field(
description=(
"Deprecated for live profiles: retained solely for the Arctic pin's retired "
"off-profile provenance; live benchmark cells run exception-free"
)
),
] = False
tree_state_weights: Annotated[
tuple[float, ...], Field(min_length=1, description="Decision-tree state weights")
] = (1.0, 0.05, 0.0)
tree_rotate_state_weights: Annotated[
bool,
Field(
description="Apply target-relative tree state weights; disable only for isolation measurements"
),
] = True
tree_directional_questions: Annotated[
bool,
Field(
description="Honor _L/_R tree-question suffixes; disable only for isolation measurements"
),
] = True
tree_ssplitmax: Annotated[int, Field(ge=0, description="Maximum state splits")] = 7
tree_ssplitthr: Annotated[float, Field(ge=0, description="State split threshold")] = 0.0
tree_csplitmax: Annotated[int, Field(ge=0, description="Maximum phone-context splits")] = 2000
tree_csplitthr: Annotated[float, Field(ge=0, description="Phone-context split threshold")] = 0.0
tree_mwfloor: Annotated[float, Field(gt=0, description="Tree mixture-weight floor")] = 1e-8
tree_intermediate_dumps: Annotated[
bool,
Field(description="Dump intermediate decision trees to worker diagnostics"),
] = False
question_npermute: Annotated[int, Field(ge=1, description="Question permutations")] = 12
question_quests_per_state: Annotated[
int, Field(ge=1, description="Questions generated per state")
] = 20
question_niter: Annotated[int, Field(ge=1, description="Question generation iterations")] = 1
multipron_training: Annotated[
bool,
Field(
description=(
"Sum posteriors over pronunciation variants; when disabled without an explicit "
"inventory policy, untied inventory resolves to upstream-compatible ``linear``"
)
),
] = True
optional_final_silence: Annotated[
bool,
Field(
description=(
"Permit final transcript silence to consume zero frames; stock SphinxTrain "
"requires that silence to consume at least one frame"
)
),
] = True
untied_inventory: Annotated[
Literal["all-triphone", "transcript-reachable", "linear"],
Field(
description=(
"Untied-model phone inventory policy: ``transcript-reachable`` includes contexts "
"reachable through every pronunciation when multipron training is enabled; "
"upstream-compatible ``linear`` includes contexts observed through each "
"transcript word's first pronunciation; ``all-triphone`` includes the complete "
"phoneset cross-product"
)
),
] = "transcript-reachable"
exclusion_schedule: Annotated[
dict[str, dict[Annotated[int, Field(strict=True)] | str, list[str]]],
Field(description="Experimental stage/pass utterance exclusions"),
] = Field(default_factory=dict)
[docs]
@model_validator(mode="before")
@classmethod
def select_inventory_default(cls, data: Any) -> Any:
if isinstance(data, dict) and data.get("multipron_training") is False:
data = dict(data)
data.setdefault("untied_inventory", "linear")
return data
[docs]
@model_validator(mode="after")
def validate_tree_state_weights(self) -> TrainingConfig:
if len(self.tree_state_weights) != self.n_state:
raise ValueError(
"tree_state_weights must contain exactly one weight per emitting state "
f"(expected {self.n_state}, got {len(self.tree_state_weights)})"
)
return self
[docs]
@field_validator("exclusion_schedule")
@classmethod
def validate_exclusions(
cls, value: dict[str, dict[int | str, list[str]]]
) -> dict[str, dict[int | str, list[str]]]:
stages = {
"ci-1g",
"ci-2g",
"ci-4g",
"ci-8g",
"cd-untied",
"cd-1g",
"cd-2g",
"cd-4g",
"cd-8g",
"cd-16g",
"cd-32g",
}
normalized: dict[str, dict[int | str, list[str]]] = {}
for stage, passes in value.items():
if stage not in stages:
raise ValueError(f"unknown BW stage {stage!r}")
normalized[stage] = {}
for selector, utterances in passes.items():
if selector != "*" and (isinstance(selector, bool) or int(selector) < 1):
raise ValueError("pass selectors must be positive integers or '*'")
if not all(utterances):
raise ValueError("utterance IDs must be non-empty")
key = selector if selector == "*" else int(selector)
if key in normalized[stage]:
raise ValueError(f"duplicate pass selector {selector!r} for {stage!r}")
normalized[stage][key] = utterances
return normalized
[docs]
@model_validator(mode="after")
def validate_training(self) -> TrainingConfig:
if (
self.arctic_a0302_zero_codebook_band is not None
and self.arctic_a0302_zero_codebook_band[0] > self.arctic_a0302_zero_codebook_band[1]
):
raise ValueError("arctic_a0302_zero_codebook_band lower bound exceeds upper bound")
if not self.multipron_training and self.untied_inventory == "transcript-reachable":
raise ValueError(
"training.untied_inventory 'transcript-reachable' requires "
"training.multipron_training: true; linear mode's equivalent is the "
"'linear' policy"
)
return self
[docs]
class SplitConfig(StrictModel):
"""Train/test split parameters."""
train_ratio: Annotated[float | None, Field(gt=0, lt=1, description="Training fraction")] = None
test_count: Annotated[
int | None,
Field(ge=0, description="Fixed test utterance count; zero disables an additional holdout"),
] = None
seed: Annotated[int, Field(description="Deterministic split seed")] = 42
[docs]
@model_validator(mode="after")
def validate_choice(self) -> SplitConfig:
if self.train_ratio is not None and self.test_count is not None:
raise ValueError("train_ratio and test_count are mutually exclusive")
return self
[docs]
class RunnerConfig(StrictModel):
"""Local pipeline execution policy."""
jobs: Annotated[int | None, Field(ge=1, description="Parallel workers; null means auto")] = None
nice: Annotated[int, Field(ge=0, description="Worker niceness increment")] = 5
class ShardingConfig(StrictModel):
"""Baum-Welch shard construction policy."""
partition_position: Annotated[
Literal["remainder-first", "remainder-last"],
Field(
description=(
"Position of uneven Baum-Welch partition capacity: ``remainder-first`` "
"distributes one extra utterance to each leading shard (the pstrain policy); "
"``remainder-last`` gives the entire remainder to the final shard (the "
"upstream SphinxTrain policy)"
)
),
] = "remainder-first"
class AlignmentConfig(StrictModel):
"""Forced-alignment transcript policy."""
beam: Annotated[float, Field(gt=0, description="Viterbi pruning beam")] = 1e-64
retry_beam_factor: Annotated[
RetryBeamFactorSetting,
Field(
description=(
"Factor that widens the beam for one retry after an utterance fails to reach "
"its final state, where values at or below 1 disable the retry; or an ascending "
"list of factors, each greater than 1 and relative to the nominal beam, tried "
"in order until one succeeds. The default retries once at 1e-200 on the "
"default beam. Retry-recovered alignments must pass the acceptance check "
"(``retry_acceptance_target``)"
),
),
] = 1e136
retry_acceptance_target: Annotated[
float | None,
Field(
gt=0,
lt=0.5,
description=(
"Fraction of this run's normal first-pass alignments that the retry acceptance "
"check would reject, used to calibrate its threshold at each retry beam; a "
"retry-recovered alignment scoring below that threshold is treated as not "
"recovered. null accepts retries unchecked. First-pass alignments are never "
"checked"
),
),
] = 0.05
failed_alignment: Annotated[
Literal["recover", "abort", "omit"],
Field(
description=(
"Forced-alignment failure policy: ``recover`` retries final-state failures at "
"each ``retry_beam_factor`` in turn; ``abort`` and ``omit`` do not retry"
)
),
] = "recover"
verbatim_tokens: Annotated[
bool,
Field(
description=(
"Honor explicit pronunciation tokens such as WORD(2) exactly during forced "
"alignment, matching PocketSphinx token handling. When false, suffixes collapse "
"to the base word and the vendored aligner considers its alternatives. This does "
"not alter training"
)
),
] = False
[docs]
class Profile(StrictModel):
"""One complete named model-training profile."""
description: Annotated[str, Field(description="Human-readable profile purpose")] = ""
features: FeatureConfig = Field(default_factory=FeatureConfig)
training: TrainingConfig = Field(default_factory=TrainingConfig)
split: SplitConfig = Field(default_factory=SplitConfig)
runner: RunnerConfig = Field(default_factory=RunnerConfig)
sharding: ShardingConfig = Field(default_factory=ShardingConfig)
alignment: AlignmentConfig = Field(default_factory=AlignmentConfig)
class ProfileDefinition(StrictModel):
"""On-disk profile; inheritance is allowed only through ``extends``."""
extends: str | None = Field(None, description="Profile to deep-merge before validation")
description: str | None = None
features: dict[str, Any] | None = None
training: dict[str, Any] | None = None
split: dict[str, Any] | None = None
runner: dict[str, Any] | None = None
sharding: dict[str, Any] | None = None
alignment: dict[str, Any] | None = None
[docs]
class ProfilesDocument(StrictModel):
"""Canonical ``etc/configs.yaml`` document."""
config_version: Literal[1] = CURRENT_CONFIG_VERSION
profiles: dict[str, ProfileDefinition]
[docs]
class OverlayDocument(StrictModel):
"""Canonical user, project, or experiment field overlay."""
config_version: Literal[1] = CURRENT_CONFIG_VERSION
profile: str | None = None
features: dict[str, Any] | None = None
training: dict[str, Any] | None = None
split: dict[str, Any] | None = None
runner: dict[str, Any] | None = None
sharding: dict[str, Any] | None = None
alignment: dict[str, Any] | None = None
SEMANTIC_BLOCKS = ("features", "training", "split", "runner", "sharding", "alignment")
def default_profile() -> Profile:
"""Return the schema-default profile."""
return Profile()