code/README.md

# 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](https://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](https://huggingface.co/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

```bash
# 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

```bash
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

```bash
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

```bash
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`).