Inference and scoring code

This folder holds the code that trained, decoded and scored the models in this repository. It is the same as github.com/maxpicy/FYP: the training repository's model, training, inference and scoring code with the comments and docstrings removed (each file was checked to parse to the same program) and one header comment per file. The loaders accept the released model.safetensors as well as a training checkpoint. decode.sh decodes, train.sh trains.

Environments

The report ran everything in two Python 3.11 environments inside the pytorch/pytorch:2.6.0-cuda11.8-cudnn9-devel container on A100 GPUs:

Environment	Used for	Pins
model	training (`train.py`), text to codes (`scripts/gen_fish_cot_samples.py`), scoring (`scripts/eval_suite.py`)	`requirements.txt`
Fish	codes to wav (`scripts/decode_fish_codes.py`), the Path B renderer (`scripts/s1mini_render.py`)	`requirements-fish.txt`

Three pins carry the load in the model environment:

mamba_ssm==2.2.6.post3, because 2.3 needs triton's set_allocator.
transformers==4.45.2, because mamba_ssm imports a class that 5.x removed.
torch==2.6.0+cu118.

env/model_env_freeze.txt and env/fish_env_freeze.txt list every package exactly as installed.

You also need these, downloaded separately:

The Fish S2 codec is codec.pth from fishaudio/openaudio-s1-mini. The licence is CC BY-NC-SA 4.0 and the repo is gated. The report used the file with md5 d3a90dbe1d535e7d50e3aa81be05e652.
A fish-speech checkout for the codec code. The report used commit bc72d1f. Set FISH_SPEECH_REPO to it and FISH_CODEC_PTH to codec.pth.
The base models' configs and tokenizer come from the Hub when a model is built:
- state-spaces/mamba2-1.3b for the pure and hybrid models;
- EleutherAI/pythia-1.4b for the transformer;
- the GPT-NeoX-20B tokenizer.
The released weights replace every base weight.
For scoring: openai/whisper-large-v3 for WER, and UTMOS22 (strong) from torch.hub tarepan/SpeechMOS:v1.2.0.

Decode

# one model, one prompt list, the report's main cell (the model writes its own plan)
bash decode.sh ../two_stage_v2/pure_c3 prompts.jsonl out/pure_c3 selfplan
# the no-plan cell, and wav files through the Fish codec
FISH_PYTHON=/path/to/fish/env/bin/python FISH_SPEECH_REPO=/path/to/fish-speech FISH_CODEC_PTH=/path/to/codec.pth \
  bash decode.sh ../two_stage_v2/pure_c3 prompts.jsonl out/pure_c3 noplan

The settings are the report's: voice 4200, seed 1234, fp32, and the cached-state decode. Each row's sampling seed comes from its index, so ROWS=a:b decodes those rows as a full run would.

On the same GPU and software, a decode reproduces the report's codes on almost every row, but not on all of them. The Mamba kernels' floating-point results can vary slightly from run to run, and now and then that flips a near-tie in sampling. The release was checked against the banked cells on A100 GPUs:

35 of the first 36 rows (12 per C3 model) came back identical, think block and all ten codebooks.
The 36th row (pure model, row 10) came out three different ways across eight more runs. Both the released file and the original training checkpoint produced the banked codes in some runs and not in others. Its think block never changed.

Score

python scripts/eval_suite.py --arm selfplan out/pure_c3/selfplan.jsonl out/pure_c3/selfplan \
    --whisper openai/whisper-large-v3 --whisper_norm whisper --out scores.json

This is the field's protocol: whisper-large-v3 with Whisper's English text normaliser for WER, UTMOS for naturalness, and completion. SIM does not apply, because the models take a speaker id, not prompt audio.

Train

python scripts/fetch_corpus.py --config stage2_v2 --out_dir data
STAGE=2 ARM=pure TIER=full INIT=../stage1p_bases/pure/model.safetensors DATA=data/stage2_v2_train.jsonl bash train.sh

train.sh holds the report's Stage 1′ (STAGE=1p) and Stage 2 (STAGE=2) recipes. TIER=full trains C3 and TIER=tags+plan trains C2 from the same file. SEED=2 gives the second training seed. Stage 2 runs 40K steps on 4 A100s.

Compare cells

python scripts/paired_stats.py --pair A.json:selfplan B.json:selfplan

This gives the paired bootstrap interval, the exact sign test and Holm correction over the pairs given, as in ../eval/stats/numbers_stats_H.json.

Files

config.py, tokenizer.py, model.py, backbones.py, train.py (the training loop and the loader), dataset.py, delay_dataset.py, inference.py and codec.py: the model and its token layout.
train.sh: the Stage 1′ and Stage 2 recipes.
decode.sh: the report's decode.
scripts/gen_fish_cot_samples.py: the generator. It writes the think block, then the codes. It also covers the plan window, plan injection, the silence guard (scripts/test_depth_checkpoint.py) and the loop breaker.
scripts/plan_extract.py: the word-level plan measured from audio.
scripts/decode_fish_codes.py and scripts/fish_codec_util.py: codes to wav.
scripts/s1mini_render.py and scripts/gen_s1mini_synth.py: Path B, where codebook 0 is rendered by the fine-tuned s1-mini in ../path_b_renderer.
scripts/eval_suite.py, its metric helpers (scripts/*_metrics.py, scripts/score_naturalness_panel.py, scripts/fan_metric.py) and eval/: scoring.
scripts/paired_stats.py: paired contrasts between score files.
scripts/export_release_weights.py: how the released weights were made from the training checkpoints.
scripts/jsonl_to_parquet.py and scripts/fetch_corpus.py: how the corpus files were made, and how to read them back into the JSONL the trainer reads.
data/silence_codes_fish_s2.json: the codebook 0 pause codes the silence guard masks.
data/duration_model.json: the length window of the no-plan cell.
tests/: python -m pytest tests -q. Tests that need a GPU or model weights skip themselves when those are absent.

Licence

MIT (LICENSE).