Training Pipeline Plan

Implementation Status Grid

Step

Pipeline Task

Step Function

Lib Function

CFFI Binding

Status

Features

step_extract_features

✅

features.extract_features()

✅ pstrain_fe_create

✅ Done

Flat

step_flat

✅

flat.init_flat_model()

✅ pstrain_flat_*, pstrain_init_gau

✅ Done

CI Train

step_ci_1g

✅

bw.BWTrainer

✅ pstrain_bw_*

✅ Done

CI Split

step_ci_2g/4g/8g

✅

split.split_gaussians()

✅ pstrain_inc_comp

✅ Done

CD Untied

step_cd_untied

✅

mdef.create_untied_mdef()

✅ pstrain_mdef_gen_untied

✅ Done

Make Quests

step_questions

✅

dtree.make_quests()

✅ pstrain_make_quests

✅ Done

Trees

step_trees

✅

dtree.build_tree()

✅ pstrain_build_tree

✅ Done

Prune

step_prune_trees

✅

dtree.prune_tree()

✅ pstrain_prune_tree

✅ Done

Tiestate

step_cd_1g

✅

dtree.tie_states()

✅ pstrain_tie_states

✅ Done

Init CD

step_cd_1g

✅

dtree.init_mixw()

✅ pstrain_init_mixw

✅ Done

CD Train

step_cd_2g/4g/8g

✅

bw.BWTrainer

✅ pstrain_bw_*

✅ Done

CD Split

step_cd_2g/4g/.../32g

✅

split.split_gaussians()

✅ pstrain_inc_comp

✅ Done

Package

package_ci_8g/cd_8g/cd_32g

✅

package.package_model()

N/A

✅ Done

Legend

  • ✅ Done - Fully implemented and tested

  • ⏳ Wire up - CFFI binding exists, need step function to orchestrate

  • ❌ Missing - Not yet implemented

Key Insight: CFFI Bindings Are Complete

All core C functions are already wrapped via CFFI:

pstrain/lib/
├── bw.py           ✅ BWTrainer class: pstrain_bw_init, pstrain_bw_process_utt,
│                      pstrain_bw_normalize, pstrain_bw_save, pstrain_bw_get_stats
├── split.py        ✅ split_gaussians(): pstrain_inc_comp
│                   ✅ kmeans(), kmeans_init(): pstrain_kmeans, pstrain_kmeans_init
├── dtree.py        ✅ build_tree(): pstrain_build_tree
│                   ✅ tie_states(): pstrain_tie_states
│                   ✅ make_quests(): pstrain_make_quests
│                   ✅ prune_tree(): pstrain_prune_tree
├── mdef.py         ✅ create_ci_mdef(): pstrain_mdef_gen_ci
│                   ✅ create_untied_mdef(): pstrain_mdef_gen_untied
│                   ✅ create_alltriphones_mdef(): pstrain_mdef_gen_alltriphones
├── flat.py         ✅ init_flat_model(): pstrain_flat_tmat, pstrain_flat_mixw,
│                      pstrain_init_gau, pstrain_norm_gau
├── features.py     ✅ FeatureExtractor: pstrain_fe_create
├── agg_seg.py      ✅ agg_seg(): pstrain_agg_seg
├── commands.py     ✅ core command paths: param_cnt, map_adapt, kdtree, delint
├── param_cnt.py    ✅ contained, declared capability (not pipeline-wired)
├── map_adapt.py    ✅ contained, declared capability (not pipeline-wired)
├── kdtree.py       ✅ contained, declared capability (not pipeline-wired)
├── delint.py       ✅ contained, declared capability (not pipeline-wired)
└── _cffi/io.py     ✅ read_gau, write_gau, read_mixw, write_mixw,
                       read_tmat, write_tmat, read_dnom, write_dnom

What’s Missing: Step Orchestration

The CFFI bindings provide low-level operations. We need higher-level “step” functions that orchestrate these into complete training steps:

Example: CI Training Step

The CFFI binding exists (BWTrainer), but we need a step function:

# pstrain/lib/steps/ci_hmm.py - needs implementation

def run_ci_training(
    model_dir: Path,
    output_dir: Path,
    features_dir: Path,
    train_fileids: Path,
    transcription: Path,
    dictionary: Path,
    n_iter: int = 10,
    convergence_ratio: float = 0.001,
) -> TrainingResult:
    """Run CI HMM training using BWTrainer.

    This orchestrates the existing CFFI bindings:
    1. Load model from model_dir
    2. For each iteration:
       a. Create BWTrainer with current model
       b. For each utterance in train_fileids:
          - Load features from features_dir
          - Look up transcript from transcription
          - Call trainer.process_utterance()
       c. Call trainer.normalize()
       d. Call trainer.save() to output_dir
       e. Check convergence
    3. Return final stats
    """
    from pstrain.lib.bw import BWTrainer, BWConfig

    trainer = BWTrainer(
        mdef_path=model_dir / "mdef",
        means_path=model_dir / "means",
        vars_path=model_dir / "variances",
        mixw_path=model_dir / "mixture_weights",
        tmat_path=model_dir / "transition_matrices",
    )

    # Set dictionary for text-based processing
    trainer.set_dict(dictionary, filler_dict)

    # Process each utterance
    for fileid in read_fileids(train_fileids):
        features = load_features(features_dir / f"{fileid}.mfc")
        transcript = get_transcript(transcription, fileid)
        trainer.process_utterance_text(features, transcript)

    # Normalize and save
    trainer.normalize()
    trainer.save(output_dir / "means", ...)

    return trainer.get_stats()

Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                              USER INTERFACE                                  │
│  pstrain build cd-8g --config wideband                                          │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                           PIPELINE RUNNER                                    │
│  pstrain.lib.pipeline.runner.Pipeline: resolves DAG from file inputs/outputs    │
│  pstrain.lib.pipeline.tasks.TARGETS:   target definitions (ci-1g, cd-8g, ...)   │
│  etc/configs.yaml:                 named configs (wideband, telephone, ...) │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                           STEP FUNCTIONS                                     │
│  pstrain.lib.steps.ci_hmm.run_ci_training()                                     │
│  pstrain.lib.steps.cd_hmm.run_cd_training()                                     │
│  pstrain.lib.steps.trees.build_trees()                                          │
│  Orchestrate multiple CFFI calls into complete training steps               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                           LIBRARY FUNCTIONS (CFFI)                           │
│  pstrain.lib.bw.BWTrainer    pstrain.lib.split.split_gaussians()                    │
│  pstrain.lib.dtree.*         pstrain.lib.mdef.*    pstrain.lib.flat.*                   │
│  All already wrapped via CFFI - no shell-outs needed                        │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                           C LIBRARY (libpstrainc)                                │
│  pstrain_bw_*, pstrain_inc_comp, pstrain_build_tree, pstrain_tie_states, pstrain_make_quests   │
│  pstrain_prune_tree, pstrain_mdef_gen_*, pstrain_flat_*, pstrain_init_gau, pstrain_norm_gau    │
└─────────────────────────────────────────────────────────────────────────────┘

What’s Complete ✅

CFFI Bindings (All Done)

  • [x] Feature extraction (pstrain_fe_create)

  • [x] Baum-Welch training (pstrain_bw_*)

  • [x] Gaussian splitting (pstrain_inc_comp)

  • [x] Flat model init (pstrain_flat_*, pstrain_init_gau, pstrain_norm_gau)

  • [x] Mdef generation (pstrain_mdef_gen_*)

  • [x] Decision trees (pstrain_build_tree, pstrain_make_quests, pstrain_prune_tree)

  • [x] State tying (pstrain_tie_states)

  • [x] Model I/O (read_gau, write_gau, read_mixw, etc.)

  • [x] K-means (pstrain_kmeans, pstrain_kmeans_init)

  • [x] Parameter counting (contained capability and core param_cnt command)

  • [x] Segment aggregation (pstrain_agg_seg)

  • [x] MAP adaptation (contained capability and core map_adapt command)

  • [x] KD-tree building (contained capability and core kdtree command)

  • [x] Deleted interpolation (contained capability and core delint command)

Orchestration Layer

  • [x] Pipeline runner with all tasks registered (pstrain.lib.pipeline)

  • [x] Dependency resolution (runner topo-sorts from inputs/outputs)

  • [x] Named configs (etc/configs.yaml: wideband, telephone, etc.)

  • [x] Target definitions (targets.yaml: ci-1g through cd-32g)

  • [x] Shared output paths (same config → same outputs)

  • [x] CLI (pstrain build cd-8g --config wideband)

  • [x] Dry-run support (pstrain --dry-run build cd-8g)

Working Steps

  • [x] step_extract_features - Feature extraction via CFFI

  • [x] step_flat - Flat model initialization via CFFI

What Needs Implementation ⏳

Complete Training Pipeline (execution order)

User Setup (manual):
  0. pstrain split                    # Create train/test fileids (optional)

Feature Extraction:
  1. step_extract_features        # ✅ DONE - MFCC extraction

Flat Model:
  2. step_flat                    # ✅ DONE - Initial uniform model

CI Training:
  3. step_ci_1g                   # ⏳ BW training on flat model
  4. step_ci_2g                   # ⏳ Split 1→2 Gaussians, BW
  5. step_ci_4g                   # ⏳ Split 2→4 Gaussians, BW
  6. step_ci_8g                   # ⏳ Split 4→8 Gaussians, BW

CD Untied:
  7. step_cd_untied               # ⏳ Generate triphone mdef, BW

Decision Trees:
  8. step_make_quests             # ⏳ Generate phonetic questions
  9. step_build_trees             # ⏳ Build decision trees
 10. step_prune_trees             # ⏳ Prune to target senones

State Tying:
 11. step_tie_states              # ⏳ Create tied mdef

CD Training:
 12. step_cd_1g                   # ⏳ BW on tied model
 13. step_cd_2g                   # ⏳ Split 1→2 Gaussians, BW
 14. step_cd_4g                   # ⏳ Split 2→4 Gaussians, BW
 15. step_cd_8g                   # ⏳ Split 4→8 Gaussians, BW
 16. step_cd_16g                  # ⏳ Split 8→16 Gaussians, BW
 17. step_cd_32g                  # ⏳ Split 16→32 Gaussians, BW

Step Functions to Implement (total order)

Many pipeline tasks share the same step function:

#

Step Function

File

CFFI Wrapper

Used By

1

run_bw_training()

steps/train.py

BWTrainer

ci-1g→8g, cd-untied, cd-1g→32g

2

run_gaussian_split()

steps/split.py

split_gaussians()

ci-2g→8g, cd-2g→32g

3

run_cd_untied_setup()

steps/cd_untied.py

create_untied_mdef()

cd-untied

4

run_make_quests()

steps/trees.py

make_quests()

step 8

5

run_build_trees()

steps/trees.py

build_tree()

step 9

6

run_prune_trees()

steps/trees.py

prune_tree()

step 10

7

run_tie_states()

steps/tiestate.py

tie_states()

step 11

Implementation order: 1 → 2 → 3 → 4 → 5 → 6 → 7

After implementing #1 and #2: pstrain build ci-8g works After implementing #3-#7: pstrain build cd-8g works

Existing CFFI wrappers (all done):

  • pstrain.lib.bw.BWTrainer - BW training

  • pstrain.lib.split.split_gaussians() - Gaussian splitting

  • pstrain.lib.mdef.create_untied_mdef() - Triphone mdef generation

  • pstrain.lib.dtree.make_quests() - Question generation

  • pstrain.lib.dtree.build_tree() - Tree building

  • pstrain.lib.dtree.prune_tree() - Tree pruning

  • pstrain.lib.dtree.tie_states() - State tying

Each step function is ~50-100 lines of orchestration code.

Step 0: Train/Test Split — ✅ DONE

Library: pstrain.lib.corpus.train_test_split. Pipeline: registered as the split task in pstrain/lib/pipeline/tasks.py (target name split).

  • Input: etc/all.transcription.

  • Outputs: experiments/{experiment}/etc/{train,test}.{fileids,transcription}.

  • Default split: 95% train, seed 42. Override via a split: block in etc/configs.yaml.

  • A complete pre-existing set of the four output files is instead treated as an authoritative external split. It is validated for exact ordered pairing, disjoint and complete corpus membership, transcript equality, and audio presence, then preserved byte-for-byte.

The CLI (pstrain split) and the pipeline task both call the same library function. The pipeline’s feature-extraction fan-out is keyed on audio_fileids() (corpus-wide, derived from audio/*.wav), so the extract tasks and the split task are parallel branches in the DAG that join at the model tasks. A content-hash sidecar distinguishes untouched automatically generated files from supplied or subsequently edited files. The pipeline runner’s completion-marker suffix hashes the task name, not file contents; split files are therefore explicit inputs to the validation task and its sentinel is an explicit input to downstream consumers, so an edit invalidates the protected work by mtime.

Step 1: BW Training (steps/train.py) — FIRST PRIORITY

Generic BW training that works for both CI and CD models:

# pstrain/lib/steps/train.py (~80 lines)

def run_bw_training(
    model_dir: Path,
    output_dir: Path,
    features_dir: Path,
    train_ctl: Path,
    transcription: Path,
    dictionary: Path,
    filler_dict: Path,
    n_iter: int = 10,
    convergence_ratio: float = 0.001,
) -> dict:
    """Run BW training iterations.

    Uses:
      - pstrain.lib.bw.BWTrainer (already implemented)
      - pstrain.lib.transcription.read_transcription()
      - pstrain.lib._cffi.io.read_mfc() or similar

    Returns:
      {"iterations": n, "final_likelihood": x, "converged": bool}
    """
    from pstrain.lib.bw import BWTrainer

    prev_likelihood = float("-inf")

    for i in range(n_iter):
        trainer = BWTrainer(
            mdef_path=model_dir / "mdef",
            means_path=model_dir / "means",
            vars_path=model_dir / "variances",
            mixw_path=model_dir / "mixture_weights",
            tmat_path=model_dir / "transition_matrices",
        )
        trainer.set_dict(str(dictionary), str(filler_dict))

        # Process all utterances
        for fileid in read_ctl(train_ctl):
            features = load_features(features_dir / f"{fileid}.mfc")
            transcript = lookup_transcript(transcription, fileid)
            trainer.process_utterance_text(features, transcript)

        trainer.normalize()
        trainer.save(output_dir / "means", output_dir / "variances",
                     output_dir / "mixture_weights", output_dir / "transition_matrices")

        stats = trainer.get_stats()
        if check_convergence(stats.avg_log_prob, prev_likelihood, convergence_ratio):
            return {"iterations": i + 1, "converged": True, ...}

        prev_likelihood = stats.avg_log_prob
        model_dir = output_dir  # next iteration reads from output

    return {"iterations": n_iter, "converged": False, ...}

Enables: ci-1g target, and all subsequent training

Forward pruning can occasionally remove the transcript’s final state. The live training driver treats that native status as recoverable: it retries the utterance exactly once with a forward beam widened by retry_beam_factor (default 1e10, so 1e-90 becomes 1e-100), restores the stage beam, and only then includes a second failure in the normal skip count and limit. Other Baum-Welch failures are never retried. retry_beam_factor may instead be an ascending list of factors, each relative to the stage beam: the rungs run in order, the first success ends the ladder, and each rung’s attempts and recoveries are reported per pass, in the stage summary, and in telemetry. An utterance whose audio is too short for its transcript runs no rung; that is measured once per failure, before the first rung.

BASIS / deliberate deviation: upstream SphinxTrain has no retry and silently skips an utterance when forward pruning loses the final state. The wider beam changes that utterance’s lattice and posteriors for the iteration, but recovers training evidence that upstream discards. Pstrain therefore permits a bounded retry (one by default, or the configured rungs), and accumulator counts come only from a successful pass; the failed passes contribute no counts.

Step 2: Gaussian Split (steps/split.py)

Split Gaussians to double density (works for CI and CD):

# pstrain/lib/steps/split.py (~40 lines)

def run_split(
    input_model_dir: Path,
    output_model_dir: Path,
) -> None:
    """Double Gaussian density.

    Uses:
      - pstrain.lib.split.split_gaussians() (already implemented)

    Also copies mdef, tmat (unchanged) to output.
    """
    from pstrain.lib.split import split_gaussians
    import shutil

    output_model_dir.mkdir(parents=True, exist_ok=True)

    # Split means, variances, mixture weights
    split_gaussians(
        in_mean_path=input_model_dir / "means",
        in_var_path=input_model_dir / "variances",
        in_mixw_path=input_model_dir / "mixture_weights",
        out_mean_path=output_model_dir / "means",
        out_var_path=output_model_dir / "variances",
        out_mixw_path=output_model_dir / "mixture_weights",
    )

    # Copy unchanged files
    shutil.copy(input_model_dir / "mdef", output_model_dir / "mdef")
    shutil.copy(input_model_dir / "transition_matrices",
                output_model_dir / "transition_matrices")

Enables: ci-2g, ci-4g, ci-8g (and later cd-2g, etc.)

Step 3: CD Untied (steps/cd_untied.py)

Generate triphone mdef and train untied CD model:

# pstrain/lib/steps/cd_untied.py (~60 lines)

def run_cd_untied(
    ci_model_dir: Path,
    output_dir: Path,
    features_dir: Path,
    train_ctl: Path,
    transcription: Path,
    dictionary: Path,
    filler_dict: Path,
    phoneset: Path,
) -> None:
    """Create and train untied CD model.

    Uses:
      - pstrain.lib.mdef.create_untied_mdef() (already implemented)
      - run_bw_training() from step 1
    """
    from pstrain.lib.mdef import create_untied_mdef

    # 1. Generate untied triphone mdef from transcripts
    create_untied_mdef(
        phone_list=phoneset,
        dict_path=dictionary,
        filler_dict_path=filler_dict,
        transcript_path=transcription,
        output_path=output_dir / "mdef",
    )

    # 2. Initialize parameters from CI model (expanded for triphones)
    init_cd_params_from_ci(ci_model_dir, output_dir)

    # 3. Train with BW
    run_bw_training(output_dir, output_dir, features_dir, ...)

Enables: cd-untied intermediate model

Step 4: Build Trees (steps/trees.py)

Generate questions, build trees, prune to target senones:

# pstrain/lib/steps/trees.py (~80 lines)

def run_build_trees(
    model_dir: Path,
    output_dir: Path,
    phoneset: Path,
    n_senones: int,
) -> None:
    """Build and prune decision trees.

    Uses:
      - pstrain.lib.dtree.make_quests() (already implemented)
      - pstrain.lib.dtree.build_tree() (already implemented)
      - pstrain.lib.dtree.prune_tree() (already implemented)
    """
    from pstrain.lib.dtree import make_quests, build_tree, prune_tree

    # 1. Generate phonetic questions
    make_quests(
        mdef_path=model_dir / "mdef",
        mixw_path=model_dir / "mixture_weights",
        mean_path=model_dir / "means",
        var_path=model_dir / "variances",
        output_path=output_dir / "questions",
    )

    # 2. Build tree for each phone and state
    phones = read_phoneset(phoneset)
    for phone in phones:
        for state in range(n_states):
            build_tree(
                mdef_path=model_dir / "mdef",
                mixw_path=model_dir / "mixture_weights",
                pset_path=output_dir / "questions",
                output_path=output_dir / f"{phone}.{state}.tree",
                phone=phone,
                state=state,
            )

    # 3. Prune to target senones
    prune_tree(
        mdef_path=model_dir / "mdef",
        pset_path=output_dir / "questions",
        input_tree_dir=output_dir,
        output_tree_dir=output_dir / "pruned",
        n_seno_target=n_senones,
    )

Enables: Decision trees for state clustering

Step 5: Tie States (steps/tiestate.py)

Apply trees to create tied mdef and initialize tied model:

# pstrain/lib/steps/tiestate.py (~50 lines)

def run_tie_states(
    untied_model_dir: Path,
    trees_dir: Path,
    output_dir: Path,
    phoneset: Path,
) -> None:
    """Apply decision trees to create tied model.

    Uses:
      - pstrain.lib.dtree.tie_states() (already implemented)
    """
    from pstrain.lib.dtree import tie_states

    # 1. Create tied mdef by applying trees
    tie_states(
        input_mdef_path=untied_model_dir / "mdef",
        output_mdef_path=output_dir / "mdef",
        tree_dir=trees_dir / "pruned",
        pset_path=trees_dir / "questions",
        allphones=True,
    )

    # 2. Initialize tied model parameters from untied
    init_tied_params(untied_model_dir, output_dir)

Enables: cd-1g (tied CD model), then training + splitting for cd-2g…cd-32g

Implementation Order

Phase 1: CI Pipeline (Already have CFFI, need step functions)
├── Implement steps/ci_hmm.py - orchestrate BWTrainer
├── Implement steps/ci_split.py - orchestrate split_gaussians
├── Wire tasks in pstrain/lib/pipeline/tasks.py to call step functions
├── Test: pstrain build ci-8g produces working model
└── Validate: decode with PocketSphinx

Phase 2: CD Pipeline (Already have CFFI, need step functions)
├── Implement steps/cd_untied.py - triphone mdef + BW
├── Implement steps/trees.py - questions + trees + prune
├── Implement steps/tiestate.py - apply trees
├── Test: pstrain build cd-8g works end-to-end
└── Validate: cd-8g achieves target WER

File Locations

pstrain/
├── lib/
│   ├── bw.py           ✅ BWTrainer CFFI wrapper
│   ├── split.py        ✅ split_gaussians() CFFI wrapper
│   ├── dtree.py        ✅ build_tree(), tie_states(), make_quests() CFFI
│   ├── mdef.py         ✅ create_ci_mdef(), create_untied_mdef() CFFI
│   ├── flat.py         ✅ init_flat_model() CFFI
│   ├── features.py     ✅ FeatureExtractor CFFI
│   ├── _cffi/          ✅ All CFFI bindings
│   └── steps/
│       ├── ci_hmm.py   ⏳ Orchestrate CI training
│       ├── ci_split.py ⏳ Orchestrate CI split
│       ├── cd_untied.py⏳ Orchestrate CD untied
│       ├── trees.py    ⏳ Orchestrate tree building
│       └── tiestate.py ⏳ Orchestrate state tying
├── lib/pipeline/
│   ├── runner.py       ✅ Task, Pipeline, staleness, executor
│   ├── context.py      ✅ PipelineContext + config loading
│   └── tasks.py        ✅ All tasks + TARGETS registry
├── etc/
│   └── configs.yaml    ✅ Named configs
└── cli/
    └── build.py        ✅ Build command (drives the runner)

Success Criteria

  • [ ] pstrain build ci-1g trains a working CI model

  • [ ] pstrain build ci-8g trains CI with Gaussian splitting

  • [ ] pstrain build cd-8g trains full CD pipeline

  • [ ] Models decode with PocketSphinx

  • [ ] cd-8g on CMU Arctic: WER < 10%

  • [ ] Training time: < 30 min for cd-8g on Arctic

  • [ ] No shell-outs - everything via CFFI