Add S3 Ref15 Full FT as reference-conditioned default
Browse filesUser-selected final step 26774; preserve S1–S10, document data selection, text repair, training and evaluation caveats.
- README.md +19 -12
- checkpoints/S3-Ref15-FullFT/metadata.json +42 -0
- checkpoints/S3-Ref15-FullFT/model_bf16.pt +3 -0
- checkpoints/default.json +10 -0
- code/infer.py +15 -4
- code/ref15_fullft/README.md +22 -0
- code/ref15_fullft/build_text_resolution.py +202 -0
- code/ref15_fullft/build_worklist.py +188 -0
- code/ref15_fullft/preflight.py +45 -0
- code/ref15_fullft/prompt_mix.py +136 -0
- code/ref15_fullft/rowgroup_plan.py +48 -0
- code/ref15_fullft/script_only.py +60 -0
- code/ref15_fullft/text_resolution.py +71 -0
- code/ref15_fullft/train.py +358 -0
- code/ref15_fullft/train_full.sbatch +31 -0
- docs/S3_REF15_FULL_FT.md +184 -0
- stats/ref15_fullft_listening50.json +0 -0
- stats/ref15_fullft_training_metrics.jsonl +0 -0
README.md
CHANGED
|
@@ -19,7 +19,15 @@ model-index:
|
|
| 19 |
|
| 20 |
# Humaneness Voice Small
|
| 21 |
|
| 22 |
-
Humaneness Voice Small is an experimental English/German text-to-speech and voice-acting model. This repository contains **ten stage-local checkpoints from one continuous S1→S10 training run**,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 23 |
|
| 24 |
The semantic backbone was **initialized from pretrained Qwen3-0.6B**, not from random weights. The **112,764,928-parameter Talker and bridge were initialized fresh** for this run. This corrects the loose description “trained from scratch” sometimes applied to earlier experiments. The Qwen hidden width is 1,024 and the local Talker width is 2,560. One semantic frame spans 80 ms; the Talker predicts **12 RVQ codebook indices per frame**, or 150 codebook indices per second. It does not predict 32 codebooks per frame. The codec reconstructs 48-kHz audio. The full state is approximately 715 million parameters, including the frozen, unused score conditioner; the model-only exports are BF16 PyTorch state dictionaries.
|
| 25 |
|
|
@@ -40,7 +48,7 @@ The `checkpoints/S*/model_bf16.pt` files are the **lowest independent held-out v
|
|
| 40 |
| S9 | tier1 | 1,439 | 41,878 | 4.7198 | 0.4540 |
|
| 41 |
| S10 | tier0 | 578 | 41,936 | 4.7278 | 0.4572 |
|
| 42 |
|
| 43 |
-
¹ These are **training presentations**, not unique source-audio hours. Multi-format prompts repeat some recordings. The ladder tiers are nested and paired with sidecar material; do not add these hours to estimate unique audio. ² Target-aware automatic reward on 1,651 shared valid acting cells per stage, not a human MOS. The S3–S10 paired difference is +0.00557 with a 95% challenge-bootstrap CI of −0.00701 to +0.01899: **the benchmark does not establish S3 as significantly better than S10**. S3
|
| 44 |
|
| 45 |
The run used 32 JUPITER Booster nodes × 4 GPUs/node, batch 4,096, **one** 2,097-step linear warmup followed by one cosine decay across all 41,936 updates—**no stage LR restarts**. Peak LRs were 8e-5 for the semantic backbone and 2.4e-4 for the Talker, ending at 5% of peak. The source `code/plan_caption_production_32nodes.json` records the exact historical manifests and hashes; its absolute JUPITER paths are provenance, not portable defaults.
|
| 46 |
|
|
@@ -131,15 +139,15 @@ Throughput was measured on an **RTX 3090**, using S3 at revision `5de86032771c37
|
|
| 131 |
- `code/continuous_train.py`, `code/train.sbatch`, `code/plan_caption_production_32nodes.json` and adjacent modules: the historical JUPITER source and training contract. The raw plan refers to local manifests and intentionally fails its hash checks if edited; regenerate plans and manifests for new data. `code/infer.py` is the standalone example for the exported weights. The code bundles MOSS export modules, tokenizer configuration and score schema; the audio codec is fetched separately from its upstream Hugging Face repository.
|
| 132 |
- `stats/training_steps.jsonl`: every logged training-step scalar line across all four Slurm invocations. `stats/invocation_summaries.json`, `stats/run_configuration.json`, `stats/stage_winners.json`, `stats/validation_loss_full.tar.gz`, `stats/acting_findings.json`, and `stats/acting_analysis.json.gz` hold the remaining scalar summaries, all held-out checkpoint-sweep records, and the public acting aggregate. `stats/CONTENTS.json` gives hashes.
|
| 133 |
|
| 134 |
-
Example (requires a CUDA environment with PyTorch, Transformers, NumPy, SoundFile, Torchaudio, Safetensors, and permission to download the separate codec):
|
| 135 |
|
| 136 |
```bash
|
| 137 |
-
python code/infer.py --stage
|
| 138 |
--prompt $'CAPTION: Warm, softly amused, conversational.\nTRANSCRIPT: "I cannot believe we made it."' \
|
| 139 |
--text 'I cannot believe we made it.' --output example.wav
|
| 140 |
```
|
| 141 |
|
| 142 |
-
|
| 143 |
|
| 144 |
## Training data, license and responsible use
|
| 145 |
|
|
@@ -159,10 +167,7 @@ Both collections use the S3 checkpoint at step 32,337; consult each adapter card
|
|
| 159 |
<!-- hvs-small-dpo-default-study-2026-10-04 -->
|
| 160 |
## Default DPO-adapter recommendation (matched speaker/emotion study)
|
| 161 |
|
| 162 |
-
For **
|
| 163 |
-
activate a DPO LoRA merely because an adapter is available. This recommendation
|
| 164 |
-
concerns the DPO-adapter choice; the separate, provisional S3 checkpoint and
|
| 165 |
-
prompt/CFG recommendations above still apply. The [matched generated-audio study
|
| 166 |
with references and listening examples](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/dpo_lora_speaker_emotion_2026-10-02.html) compared the same 12
|
| 167 |
English/German acting challenges, prompts and seeds for each condition (432
|
| 168 |
outputs per condition, including 198 reference-conditioned outputs). Its
|
|
@@ -170,7 +175,7 @@ outputs per condition, including 198 reference-conditioned outputs). Its
|
|
| 170 |
|
| 171 |
| S3 condition | Orange reference cosine ↑ | Parakeet WER ↓ | Requested-emotion instruction lift ↑ |
|
| 172 |
|---|---:|---:|---:|
|
| 173 |
-
|
|
| 174 |
| Recent pitch/quality/CFG-p2 Rank 16 × 0.50 | 0.534 | 0.106 | −0.019 |
|
| 175 |
| Same Rank 16 × 1.00 | 0.529 | 0.123 | −0.116 |
|
| 176 |
| Same Rank 64 × 1.00 | 0.534 | 0.133 | −0.099 |
|
|
@@ -190,8 +195,10 @@ exact prompts, paired audio players and methodological limits.
|
|
| 190 |
<!-- hvs-small-explicit-optional-dpo-default-2026-10-04 -->
|
| 191 |
### Which DPO LoRA should I load by default?
|
| 192 |
|
| 193 |
-
**None for general use
|
| 194 |
-
|
|
|
|
|
|
|
| 195 |
If an application explicitly prioritizes similarity to a provided speaker
|
| 196 |
reference, the **best-tested optional DPO LoRA at this date** is
|
| 197 |
[`subset_pitch_quality_p2_rank016.adapter.pt`](https://huggingface.co/laion/Humaneness-Voice-Small-DPO-LoRAs/blob/main/adapters/subset_pitch_quality_p2_rank016.adapter.pt)
|
|
|
|
| 19 |
|
| 20 |
# Humaneness Voice Small
|
| 21 |
|
| 22 |
+
Humaneness Voice Small is an experimental English/German text-to-speech and voice-acting model. This repository contains **ten stage-local checkpoints from one continuous S1→S10 training run** plus a later S3-based full-parameter, long-reference continuation. It also contains training and inference source, run statistics, and a guide to the exact prompt surfaces. It is not a standard `AutoModel.from_pretrained()` checkpoint: it uses a Qwen3-0.6B semantic backbone, a MOSS-style local Talker, and the separately loaded [MOSS Audio Tokenizer v2](https://huggingface.co/OpenMOSS-Team/MOSS-Audio-Tokenizer-v2). Read the inference instructions before loading weights.
|
| 23 |
+
|
| 24 |
+
## Current default: S3 Ref15 Full FT (9 October 2026)
|
| 25 |
+
|
| 26 |
+
The new **reference-conditioned default checkpoint** is [`checkpoints/S3-Ref15-FullFT/model_bf16.pt`](checkpoints/S3-Ref15-FullFT/model_bf16.pt), the final **step 26,774** of a full-parameter continuation from the **original S3 step 32,337**. The user chose it because it sounded acoustically preferable in listening. The updated `code/infer.py` loads it when `--stage` is omitted or set to `default`. The ten original stage checkpoints are unchanged; `--stage S3` explicitly loads the original S3 and is still the comparison to try for **reference-free** generation. For the new default, supply a clean, distinct 5–15-second reference WAV; every accepted retraining example had a reference.
|
| 27 |
+
|
| 28 |
+
The key fix is the historical **2.96-second reference crop**: the original training loader and earlier example inference code used at most 37 × 80-ms codec frames even though the stored reference sequences were longer. This continuation used up to **187 frames / 14.96 seconds**, with a 55/20/20/5 family of script, historical structured, caption-transcript, and transcript-only prompts. It fully updated the active semantic backbone and Talker after a conservative target-text repair gate; it is **not** the Rank-128 LoRA and did not start from that adapter. There were **3,255,217 accepted training targets**, each presented once with its distinct reference. Read the detailed [data selection, repair, training and evaluation report](docs/S3_REF15_FULL_FT.md) and the [checkpoint metadata](checkpoints/S3-Ref15-FullFT/metadata.json) before reusing it.
|
| 29 |
+
|
| 30 |
+
This is a **listening-selected acoustic default**, not an objectively established all-round winner. In a matched 50-case non-speaker-disjoint test, final-step WER was 0.096 versus 0.088 for original S3; Orange cosine to the supplied reference was 0.707 versus 0.745, while cosine to the intended target improved from 0.658 to 0.734. Step 25,000 had better measured WER and target cosine than the final step. These learned/ASR proxies and the user's listening impression answer different questions; further blinded and speaker-disjoint evaluation is needed. No optimizer state is uploaded for this continuation.
|
| 31 |
|
| 32 |
The semantic backbone was **initialized from pretrained Qwen3-0.6B**, not from random weights. The **112,764,928-parameter Talker and bridge were initialized fresh** for this run. This corrects the loose description “trained from scratch” sometimes applied to earlier experiments. The Qwen hidden width is 1,024 and the local Talker width is 2,560. One semantic frame spans 80 ms; the Talker predicts **12 RVQ codebook indices per frame**, or 150 codebook indices per second. It does not predict 32 codebooks per frame. The codec reconstructs 48-kHz audio. The full state is approximately 715 million parameters, including the frozen, unused score conditioner; the model-only exports are BF16 PyTorch state dictionaries.
|
| 33 |
|
|
|
|
| 48 |
| S9 | tier1 | 1,439 | 41,878 | 4.7198 | 0.4540 |
|
| 49 |
| S10 | tier0 | 578 | 41,936 | 4.7278 | 0.4572 |
|
| 50 |
|
| 51 |
+
¹ These are **training presentations**, not unique source-audio hours. Multi-format prompts repeat some recordings. The ladder tiers are nested and paired with sidecar material; do not add these hours to estimate unique audio. ² Target-aware automatic reward on 1,651 shared valid acting cells per stage, not a human MOS. The S3–S10 paired difference is +0.00557 with a 95% challenge-bootstrap CI of −0.00701 to +0.01899: **the benchmark does not establish S3 as significantly better than S10**. Among the **original ladder stages**, S3 was a reasonable provisional choice; inspect the same prompts in S3/S10 before choosing for a production voice. The later Ref15 Full FT is now the reference-conditioned default by listening choice, under the caveats above. S4 has the best held-out validation loss. All underlying values and intervals are in `stats/acting_findings.json` and `stats/acting_analysis.json.gz`.
|
| 52 |
|
| 53 |
The run used 32 JUPITER Booster nodes × 4 GPUs/node, batch 4,096, **one** 2,097-step linear warmup followed by one cosine decay across all 41,936 updates—**no stage LR restarts**. Peak LRs were 8e-5 for the semantic backbone and 2.4e-4 for the Talker, ending at 5% of peak. The source `code/plan_caption_production_32nodes.json` records the exact historical manifests and hashes; its absolute JUPITER paths are provenance, not portable defaults.
|
| 54 |
|
|
|
|
| 139 |
- `code/continuous_train.py`, `code/train.sbatch`, `code/plan_caption_production_32nodes.json` and adjacent modules: the historical JUPITER source and training contract. The raw plan refers to local manifests and intentionally fails its hash checks if edited; regenerate plans and manifests for new data. `code/infer.py` is the standalone example for the exported weights. The code bundles MOSS export modules, tokenizer configuration and score schema; the audio codec is fetched separately from its upstream Hugging Face repository.
|
| 140 |
- `stats/training_steps.jsonl`: every logged training-step scalar line across all four Slurm invocations. `stats/invocation_summaries.json`, `stats/run_configuration.json`, `stats/stage_winners.json`, `stats/validation_loss_full.tar.gz`, `stats/acting_findings.json`, and `stats/acting_analysis.json.gz` hold the remaining scalar summaries, all held-out checkpoint-sweep records, and the public acting aggregate. `stats/CONTENTS.json` gives hashes.
|
| 141 |
|
| 142 |
+
Example using the **new reference-conditioned default** (requires a CUDA environment with PyTorch, Transformers, NumPy, SoundFile, Torchaudio, Safetensors, and permission to download the separate codec):
|
| 143 |
|
| 144 |
```bash
|
| 145 |
+
python code/infer.py --stage default --reference-wav reference.wav --language en --frames 75 \
|
| 146 |
--prompt $'CAPTION: Warm, softly amused, conversational.\nTRANSCRIPT: "I cannot believe we made it."' \
|
| 147 |
--text 'I cannot believe we made it.' --output example.wav
|
| 148 |
```
|
| 149 |
|
| 150 |
+
Do not use the target recording as its own reference. For reference-free inference, compare `--stage S3`; the new default was trained only with references. **Correction (8 October 2026):** an earlier public `code/infer.py` rejected any reference over 37 MOSS codec frames (2.96 s). This was an inference-code limit as well as the historical S1–S10 training-loader crop, not an architectural limit. The updated example accepts a configurable reference ceiling (15 s default, 20 s maximum), reports a warning below 5 s and no longer silently blocks longer audio. The new Full FT was trained with up to 14.96 s; using longer than 15 s at inference is out of distribution and warned about. The original S1–S10 weights were trained with at most 2.96 s references, so longer-reference behavior from an original stage is experimental, not a guaranteed improvement. The separate [S3 Ref15 Rank-128 LoRA](https://huggingface.co/laion/Humaneness-Voice-Small-Ref15-LoRA-r128) is a different, frozen-backbone comparison. The example is a source-level starting point; the public listening Space is the already-audited playback path. This repository does **not** contain the large original SFT-3 model weights or original Qwen checkpoint separately, because the exported Humaneness Voice Small states include every trained model weight required for inference.
|
| 151 |
|
| 152 |
## Training data, license and responsible use
|
| 153 |
|
|
|
|
| 167 |
<!-- hvs-small-dpo-default-study-2026-10-04 -->
|
| 168 |
## Default DPO-adapter recommendation (matched speaker/emotion study)
|
| 169 |
|
| 170 |
+
For **the historical S3 DPO comparison, use the unadapted S3 checkpoint rather than activating a DPO LoRA by default**. This is an adapter-level recommendation, **not** a reversal of the newer Ref15 Full FT checkpoint default above; the DPO LoRAs were trained against S3 and should not be assumed compatible or beneficial on Ref15 Full FT. The historical S3 prompt/CFG recommendations above also refer to S3, not a completed Ref15 tuning sweep. The [matched generated-audio study
|
|
|
|
|
|
|
|
|
|
| 171 |
with references and listening examples](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/dpo_lora_speaker_emotion_2026-10-02.html) compared the same 12
|
| 172 |
English/German acting challenges, prompts and seeds for each condition (432
|
| 173 |
outputs per condition, including 198 reference-conditioned outputs). Its
|
|
|
|
| 175 |
|
| 176 |
| S3 condition | Orange reference cosine ↑ | Parakeet WER ↓ | Requested-emotion instruction lift ↑ |
|
| 177 |
|---|---:|---:|---:|
|
| 178 |
+
| Historical S3 base, no DPO adapter **(study default)** | 0.501 | **0.096** | −0.021 |
|
| 179 |
| Recent pitch/quality/CFG-p2 Rank 16 × 0.50 | 0.534 | 0.106 | −0.019 |
|
| 180 |
| Same Rank 16 × 1.00 | 0.529 | 0.123 | −0.116 |
|
| 181 |
| Same Rank 64 × 1.00 | 0.534 | 0.133 | −0.099 |
|
|
|
|
| 195 |
<!-- hvs-small-explicit-optional-dpo-default-2026-10-04 -->
|
| 196 |
### Which DPO LoRA should I load by default?
|
| 197 |
|
| 198 |
+
**None for general use on the historical S3 base:** unadapted S3 achieved
|
| 199 |
+
lower WER than any DPO LoRA in that matched study. The repository-wide
|
| 200 |
+
reference-conditioned checkpoint default is now Ref15 Full FT, which was
|
| 201 |
+
evaluated under a different 50-case protocol; this DPO result does not rank it.
|
| 202 |
If an application explicitly prioritizes similarity to a provided speaker
|
| 203 |
reference, the **best-tested optional DPO LoRA at this date** is
|
| 204 |
[`subset_pitch_quality_p2_rank016.adapter.pt`](https://huggingface.co/laion/Humaneness-Voice-Small-DPO-LoRAs/blob/main/adapters/subset_pitch_quality_p2_rank016.adapter.pt)
|
checkpoints/S3-Ref15-FullFT/metadata.json
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"name": "Humaneness Voice Small S3 Ref15 Full FT",
|
| 3 |
+
"role": "repository_default_as_of_2026-10-09",
|
| 4 |
+
"format": "torch.save(model.state_dict()) in bfloat16; load with weights_only=True",
|
| 5 |
+
"file": "model_bf16.pt",
|
| 6 |
+
"bytes": 1431076525,
|
| 7 |
+
"sha256": "e1b2c4613545d1db21c6f60a5ab2c27b5320aa05c177717ae268b1ff870d406e",
|
| 8 |
+
"final_step": 26774,
|
| 9 |
+
"base": {
|
| 10 |
+
"repository_path": "checkpoints/S3/model_bf16.pt",
|
| 11 |
+
"original_training_step": 32337,
|
| 12 |
+
"sha256": "60079723bf797a81107def1e6fc479e5353b3cc51f3780f98ea65a9f5c109f97"
|
| 13 |
+
},
|
| 14 |
+
"training": {
|
| 15 |
+
"kind": "hvs-long-ref-full-ft-v1",
|
| 16 |
+
"optimizer": "AdamW, full-parameter except unused frozen score conditioner",
|
| 17 |
+
"global_batch": 128,
|
| 18 |
+
"world_gpus": 32,
|
| 19 |
+
"train_unique_uids": 3255217,
|
| 20 |
+
"accepted_total_uids_including_validation": 3271673,
|
| 21 |
+
"reference_max_frames_80ms": 187,
|
| 22 |
+
"reference_max_seconds": 14.96,
|
| 23 |
+
"semantic_peak_lr": 0.000008,
|
| 24 |
+
"talker_peak_lr": 0.000016,
|
| 25 |
+
"warmup_fraction": 0.05,
|
| 26 |
+
"schedule": "single_cosine_to_10_percent_of_peak_no_restart",
|
| 27 |
+
"seed": 20261004,
|
| 28 |
+
"source_worklist_sha256": "52c5223d1935ca806044b2f832291f11271240ca26d6316ce2f82bafe4ae2939",
|
| 29 |
+
"target_text_resolution_sha256": "c8f636ec873f68360287626b336ae803a4c08f0c54a8672bc0d42e713406fee0",
|
| 30 |
+
"score_schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
|
| 31 |
+
"exact_prompt_counts": {
|
| 32 |
+
"script_only": 1932540,
|
| 33 |
+
"general_script": 570663,
|
| 34 |
+
"caption_timbre": 318972,
|
| 35 |
+
"caption_bude": 269892,
|
| 36 |
+
"transcript_only": 163150
|
| 37 |
+
}
|
| 38 |
+
},
|
| 39 |
+
"selection": "User-selected acoustic default after listening; not the optimum of every automatic metric",
|
| 40 |
+
"optimizer_state_uploaded": false,
|
| 41 |
+
"details": "docs/S3_REF15_FULL_FT.md"
|
| 42 |
+
}
|
checkpoints/S3-Ref15-FullFT/model_bf16.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e1b2c4613545d1db21c6f60a5ab2c27b5320aa05c177717ae268b1ff870d406e
|
| 3 |
+
size 1431076525
|
checkpoints/default.json
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"checkpoint": "S3-Ref15-FullFT",
|
| 3 |
+
"model_file": "checkpoints/S3-Ref15-FullFT/model_bf16.pt",
|
| 4 |
+
"metadata_file": "checkpoints/S3-Ref15-FullFT/metadata.json",
|
| 5 |
+
"chosen_on": "2026-10-09",
|
| 6 |
+
"selection_basis": "user listening preference for acoustic quality",
|
| 7 |
+
"reference_conditioned": true,
|
| 8 |
+
"trained_reference_max_seconds": 14.96,
|
| 9 |
+
"reference_free_fallback": "S3"
|
| 10 |
+
}
|
code/infer.py
CHANGED
|
@@ -2,7 +2,8 @@
|
|
| 2 |
"""Single-clip Humaneness Voice Small inference; use on a CUDA GPU.
|
| 3 |
|
| 4 |
Example (from the model repository root):
|
| 5 |
-
python code/infer.py --stage
|
|
|
|
| 6 |
--text 'Hello there.' --frames 60 --output hello.wav
|
| 7 |
"""
|
| 8 |
from __future__ import annotations
|
|
@@ -18,7 +19,9 @@ sys.path.insert(0, str(ROOT / 'code'))
|
|
| 18 |
|
| 19 |
def main() -> None:
|
| 20 |
parser = argparse.ArgumentParser(description=__doc__)
|
| 21 |
-
parser.add_argument('--stage', choices=[
|
|
|
|
|
|
|
| 22 |
parser.add_argument('--prompt', required=True, help='Literal GENERAL/SCRIPT, CAPTION/TRANSCRIPT or TRANSCRIPT text')
|
| 23 |
parser.add_argument('--text', required=True, help='The exact spoken transcript')
|
| 24 |
parser.add_argument('--frames', type=int, required=True, help='Frame budget at 12.5 frames/second')
|
|
@@ -32,6 +35,9 @@ def main() -> None:
|
|
| 32 |
assert args.frames > 0
|
| 33 |
if not 0 < args.reference_max_seconds <= 20:
|
| 34 |
parser.error('--reference-max-seconds must be in (0, 20]')
|
|
|
|
|
|
|
|
|
|
| 35 |
|
| 36 |
import numpy as np
|
| 37 |
import soundfile as sf
|
|
@@ -51,7 +57,11 @@ def main() -> None:
|
|
| 51 |
moss_small.load_qwen_backbone = lambda model, log=print: None
|
| 52 |
schema = json.loads((ROOT / 'assets/score_schema.json').read_text())
|
| 53 |
model, config = build_fresh(schema, log=lambda _: None)
|
| 54 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
map_location='cpu', weights_only=True)
|
| 56 |
model.load_state_dict(state, strict=True)
|
| 57 |
model.tie_weights()
|
|
@@ -113,7 +123,8 @@ def main() -> None:
|
|
| 113 |
args.output.parent.mkdir(parents=True, exist_ok=True)
|
| 114 |
sf.write(args.output, wave, sample_rate, subtype='PCM_16')
|
| 115 |
print(json.dumps({'output': str(args.output), 'frames': len(codes),
|
| 116 |
-
'duration_seconds': len(wave) / sample_rate,
|
|
|
|
| 117 |
|
| 118 |
|
| 119 |
if __name__ == '__main__':
|
|
|
|
| 2 |
"""Single-clip Humaneness Voice Small inference; use on a CUDA GPU.
|
| 3 |
|
| 4 |
Example (from the model repository root):
|
| 5 |
+
python code/infer.py --stage default --reference-wav reference.wav \
|
| 6 |
+
--prompt 'CAPTION: warm, amused narration\nTRANSCRIPT: "Hello there."' \
|
| 7 |
--text 'Hello there.' --frames 60 --output hello.wav
|
| 8 |
"""
|
| 9 |
from __future__ import annotations
|
|
|
|
| 19 |
|
| 20 |
def main() -> None:
|
| 21 |
parser = argparse.ArgumentParser(description=__doc__)
|
| 22 |
+
parser.add_argument('--stage', choices=['default', 'S3-Ref15-FullFT'] +
|
| 23 |
+
[f'S{i}' for i in range(1, 11)], default='default',
|
| 24 |
+
help='default = S3 Ref15 Full FT; S1–S10 retain the original ladder weights')
|
| 25 |
parser.add_argument('--prompt', required=True, help='Literal GENERAL/SCRIPT, CAPTION/TRANSCRIPT or TRANSCRIPT text')
|
| 26 |
parser.add_argument('--text', required=True, help='The exact spoken transcript')
|
| 27 |
parser.add_argument('--frames', type=int, required=True, help='Frame budget at 12.5 frames/second')
|
|
|
|
| 35 |
assert args.frames > 0
|
| 36 |
if not 0 < args.reference_max_seconds <= 20:
|
| 37 |
parser.error('--reference-max-seconds must be in (0, 20]')
|
| 38 |
+
if args.stage in ('default', 'S3-Ref15-FullFT') and args.reference_max_seconds > 15:
|
| 39 |
+
print('Warning: Ref15 Full FT trained with at most 14.96 s of reference; '
|
| 40 |
+
'longer inference references are out of distribution', file=sys.stderr)
|
| 41 |
|
| 42 |
import numpy as np
|
| 43 |
import soundfile as sf
|
|
|
|
| 57 |
moss_small.load_qwen_backbone = lambda model, log=print: None
|
| 58 |
schema = json.loads((ROOT / 'assets/score_schema.json').read_text())
|
| 59 |
model, config = build_fresh(schema, log=lambda _: None)
|
| 60 |
+
checkpoint_name = 'S3-Ref15-FullFT' if args.stage == 'default' else args.stage
|
| 61 |
+
if checkpoint_name == 'S3-Ref15-FullFT' and args.reference_wav is None:
|
| 62 |
+
print('Warning: the Ref15 default was tuned only with reference audio; '
|
| 63 |
+
'for reference-free inference compare --stage S3', file=sys.stderr)
|
| 64 |
+
state = torch.load(ROOT / 'checkpoints' / checkpoint_name / 'model_bf16.pt',
|
| 65 |
map_location='cpu', weights_only=True)
|
| 66 |
model.load_state_dict(state, strict=True)
|
| 67 |
model.tie_weights()
|
|
|
|
| 123 |
args.output.parent.mkdir(parents=True, exist_ok=True)
|
| 124 |
sf.write(args.output, wave, sample_rate, subtype='PCM_16')
|
| 125 |
print(json.dumps({'output': str(args.output), 'frames': len(codes),
|
| 126 |
+
'duration_seconds': len(wave) / sample_rate,
|
| 127 |
+
'stage': args.stage, 'resolved_checkpoint': checkpoint_name}))
|
| 128 |
|
| 129 |
|
| 130 |
if __name__ == '__main__':
|
code/ref15_fullft/README.md
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Historical S3 Ref15 Full FT source
|
| 2 |
+
|
| 3 |
+
These are the source entrypoints used to construct the packed S3 long-reference
|
| 4 |
+
worklist and train the final step-26,774 checkpoint. They are published for
|
| 5 |
+
provenance and audit, **not as a one-command portable trainer**: paths point
|
| 6 |
+
to the original JUPITER manifests, and imports also use modules elsewhere in
|
| 7 |
+
this model repository and the original training workspace. Regenerate and
|
| 8 |
+
verify the input manifests, schema and SHA-256 contract before using new data.
|
| 9 |
+
|
| 10 |
+
- `build_worklist.py` selects unique physical S3 targets with a validated,
|
| 11 |
+
distinct reference and raises the codec-reference crop from 37 to 187 frames.
|
| 12 |
+
- `prompt_mix.py` and `script_only.py` implement the deterministic five-mode
|
| 13 |
+
instruction mixture and remove stable voice-identity descriptions from the
|
| 14 |
+
bare SCRIPT mode.
|
| 15 |
+
- `build_text_resolution.py` and `text_resolution.py` implement the
|
| 16 |
+
conservative target-text repair/exclusion sidecar.
|
| 17 |
+
- `rowgroup_plan.py` balances Parquet row groups over GPU ranks.
|
| 18 |
+
- `train.py`, `preflight.py`, and `train_full.sbatch` implement the full-model
|
| 19 |
+
continuation, fail-closed preflight and resumable 32-GPU Slurm segments.
|
| 20 |
+
|
| 21 |
+
See the repository's [detailed Ref15 report](../../docs/S3_REF15_FULL_FT.md)
|
| 22 |
+
for exact selection counts, optimizer settings, limitations and evaluation.
|
code/ref15_fullft/build_text_resolution.py
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Conservative packed disposition of S3 transcript suspects after full ASR.
|
| 3 |
+
|
| 4 |
+
This is an automated repair, not ground-truth transcription. It favors
|
| 5 |
+
excluding an uncertain short target over teaching a full-scene transcript to
|
| 6 |
+
a clipped target. Original TAR/Parquet assets are immutable. Every suspect
|
| 7 |
+
must have an ASR hypothesis and one explicit disposition; the 128+ sample QA
|
| 8 |
+
is a structural/audio-ASR consistency check, not a claim of human listening.
|
| 9 |
+
"""
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from collections import Counter, defaultdict
|
| 13 |
+
from datetime import datetime, timezone
|
| 14 |
+
from difflib import SequenceMatcher
|
| 15 |
+
import hashlib
|
| 16 |
+
import json
|
| 17 |
+
import math
|
| 18 |
+
import os
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
import re
|
| 21 |
+
|
| 22 |
+
import pyarrow as pa
|
| 23 |
+
import pyarrow.parquet as pq
|
| 24 |
+
|
| 25 |
+
from audit_target_text import AUDIT, PARENTHETICAL
|
| 26 |
+
from build_worklist import EXPECTED_ELIGIBLE, OUT, SHARDS, sha
|
| 27 |
+
from text_resolution import RESOLUTION_DIR
|
| 28 |
+
|
| 29 |
+
ASR = AUDIT.parent / 'target_asr_v1'
|
| 30 |
+
TOKEN = re.compile(r"[^\W_]+(?:['’\-][^\W_]+)*", re.UNICODE)
|
| 31 |
+
BAD_SINGLE = {'a', 'i', 'o', 'uh', 'um', 'mmm', 'hmm', 'ah', 'oh', 'the'}
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def lexical(text: str) -> list[str]:
|
| 35 |
+
return [word.casefold() for word in TOKEN.findall(PARENTHETICAL.sub(' ', text or ''))]
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def repeated_run(words: list[str]) -> int:
|
| 39 |
+
best = run = 0
|
| 40 |
+
prior = None
|
| 41 |
+
for word in words:
|
| 42 |
+
run = run + 1 if word == prior else 1
|
| 43 |
+
best = max(best, run)
|
| 44 |
+
prior = word
|
| 45 |
+
return best
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def classify(row: dict) -> tuple[str, str, str, float]:
|
| 49 |
+
source_text = str(row['source_text'] or '').strip()
|
| 50 |
+
asr_text = ' '.join(str(row['asr_text'] or '').strip().split())
|
| 51 |
+
old, new = lexical(source_text), lexical(asr_text)
|
| 52 |
+
duration = float(row['duration_s'])
|
| 53 |
+
rate_cap = max(4, math.ceil(duration * 6.5))
|
| 54 |
+
if re.search(r'\b(?:GENERAL|SCRIPT|CAPTION|TRANSCRIPT):|\[pause\s', asr_text, re.I):
|
| 55 |
+
return 'exclude', '', 'asr_reserved_prompt_marker', 0.
|
| 56 |
+
if len(new) > rate_cap or repeated_run(new) >= 6 or len(asr_text) > 2048:
|
| 57 |
+
return 'exclude', '', 'implausible_asr', 0.
|
| 58 |
+
if not new:
|
| 59 |
+
return 'exclude', '', 'empty_asr', 0.
|
| 60 |
+
# A few original metadata strings contain thousands of hallucinated
|
| 61 |
+
# repeated words. Bound the quadratic alignment work; the full old
|
| 62 |
+
# length is still used by the keep/no-keep decision below.
|
| 63 |
+
old_head = old[:max(256, len(new) * 4)]
|
| 64 |
+
matcher = SequenceMatcher(None, old_head, new, autojunk=False)
|
| 65 |
+
common = sum(block.size for block in matcher.get_matching_blocks())
|
| 66 |
+
support = common / len(new)
|
| 67 |
+
agreement = matcher.ratio()
|
| 68 |
+
aligned = int(row['aligned_word_count'])
|
| 69 |
+
# A matching transcript is a false-positive rate/alignment flag. An
|
| 70 |
+
# aligned original with small ASR differences can preserve proper nouns.
|
| 71 |
+
if (agreement >= .88 and len(old) <= rate_cap) or (
|
| 72 |
+
aligned >= .8 * len(old) and agreement >= .80 and len(old) <= rate_cap):
|
| 73 |
+
return 'verified_keep', '', 'asr_agrees_with_source', support
|
| 74 |
+
if len(new) == 1:
|
| 75 |
+
if len(old) <= 2 and new[0] == old[0] and new[0] not in BAD_SINGLE:
|
| 76 |
+
return 'verified_keep', '', 'single_word_agreement', support
|
| 77 |
+
return 'exclude', '', 'short_asr_not_safe', support
|
| 78 |
+
# Full-scene text clipped to a sentence typically has source-prefix
|
| 79 |
+
# support near 1 even if the global source-vs-ASR WER is very large.
|
| 80 |
+
# Without substantial support, the ASR itself may be a hallucination.
|
| 81 |
+
if support >= .68 and common >= 2 and asr_text:
|
| 82 |
+
return 'relabel', asr_text, 'supported_clip_asr', support
|
| 83 |
+
return 'exclude', '', 'asr_source_conflict', support
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def main():
|
| 87 |
+
asr_manifest_path = ASR / 'MANIFEST.json'
|
| 88 |
+
audit_report_path = AUDIT / 'REPORT.json'
|
| 89 |
+
asr_spec = json.loads(asr_manifest_path.read_text())
|
| 90 |
+
audit_spec = json.loads(audit_report_path.read_text())
|
| 91 |
+
assert asr_spec['status'] == 'COMPLETE_HYPOTHESES_NOT_APPROVED_LABELS'
|
| 92 |
+
assert asr_spec['asr_include_unaligned'] is True, 'Every suspect needs a disposition'
|
| 93 |
+
assert asr_spec['source_audit_sha256'] == sha(audit_report_path)
|
| 94 |
+
assert audit_spec['rows'] == EXPECTED_ELIGIBLE
|
| 95 |
+
assert not RESOLUTION_DIR.exists(), 'Do not overwrite an existing resolution'
|
| 96 |
+
|
| 97 |
+
expected = set()
|
| 98 |
+
for shard in range(SHARDS):
|
| 99 |
+
part = AUDIT / f'suspect-{shard:03d}.parquet'
|
| 100 |
+
for batch in pq.ParquetFile(part).iter_batches(columns=['uid'], batch_size=8192):
|
| 101 |
+
for uid in batch.column(0).to_pylist():
|
| 102 |
+
assert uid not in expected, uid
|
| 103 |
+
expected.add(uid)
|
| 104 |
+
assert len(expected) == sum(x['any_suspect'] for x in audit_spec['by_source'].values())
|
| 105 |
+
assert len(expected) == asr_spec['rows']
|
| 106 |
+
|
| 107 |
+
RESOLUTION_DIR.mkdir(parents=True)
|
| 108 |
+
changes_path = RESOLUTION_DIR / 'changes.parquet'
|
| 109 |
+
changes_partial = changes_path.with_suffix('.parquet.partial')
|
| 110 |
+
schema = pa.schema([
|
| 111 |
+
('uid', pa.string()), ('action', pa.string()), ('text', pa.string()),
|
| 112 |
+
('source', pa.string()), ('reason', pa.string()), ('asr_support', pa.float32()),
|
| 113 |
+
('duration_s', pa.float32()), ('source_text', pa.string()), ('asr_text', pa.string()),
|
| 114 |
+
('record_locator_json', pa.string()), ('target_locator_json', pa.string())])
|
| 115 |
+
writer = pq.ParquetWriter(changes_partial, schema, compression='zstd')
|
| 116 |
+
observed = set()
|
| 117 |
+
counts = Counter()
|
| 118 |
+
by_source = defaultdict(Counter)
|
| 119 |
+
excluded_hours = 0.
|
| 120 |
+
qa_candidates = defaultdict(list)
|
| 121 |
+
try:
|
| 122 |
+
for part_spec in asr_spec['parts']:
|
| 123 |
+
path = Path(part_spec['path'])
|
| 124 |
+
assert sha(path) == part_spec['sha256']
|
| 125 |
+
for batch in pq.ParquetFile(path).iter_batches(batch_size=4096):
|
| 126 |
+
output = []
|
| 127 |
+
for row in batch.to_pylist():
|
| 128 |
+
uid = row['uid']
|
| 129 |
+
assert uid in expected and uid not in observed, uid
|
| 130 |
+
observed.add(uid)
|
| 131 |
+
# The full ASR pass must retain these audit fields. This
|
| 132 |
+
# prevents silently applying a repair to another target.
|
| 133 |
+
assert int(row['text_word_count']) == len(lexical(row['source_text']))
|
| 134 |
+
action, text, reason, support = classify(row)
|
| 135 |
+
counts[action] += 1
|
| 136 |
+
by_source[row['source']][action] += 1
|
| 137 |
+
if action == 'exclude':
|
| 138 |
+
excluded_hours += float(row['duration_s']) / 3600
|
| 139 |
+
decision = {'uid': uid, 'action': action, 'text': text,
|
| 140 |
+
'source': row['source'], 'reason': reason,
|
| 141 |
+
'asr_support': support,
|
| 142 |
+
'duration_s': float(row['duration_s']),
|
| 143 |
+
'source_text': row['source_text'],
|
| 144 |
+
'asr_text': row['asr_text'],
|
| 145 |
+
'record_locator_json': row['record_locator_json'],
|
| 146 |
+
'target_locator_json': row['target_locator_json']}
|
| 147 |
+
output.append(decision)
|
| 148 |
+
bucket = (row['source'], action)
|
| 149 |
+
digest = int.from_bytes(hashlib.blake2b(uid.encode(), digest_size=8).digest(), 'big')
|
| 150 |
+
qa_candidates[bucket].append((digest, decision))
|
| 151 |
+
if len(qa_candidates[bucket]) > 16:
|
| 152 |
+
qa_candidates[bucket].sort(key=lambda x: x[0])
|
| 153 |
+
del qa_candidates[bucket][16:]
|
| 154 |
+
if output:
|
| 155 |
+
writer.write_table(pa.Table.from_pylist(output, schema=schema), row_group_size=4096)
|
| 156 |
+
finally:
|
| 157 |
+
writer.close()
|
| 158 |
+
assert observed == expected, {'missing': len(expected-observed), 'extra': len(observed-expected)}
|
| 159 |
+
assert pq.read_metadata(changes_partial).num_rows == len(expected)
|
| 160 |
+
os.replace(changes_partial, changes_path)
|
| 161 |
+
|
| 162 |
+
samples = []
|
| 163 |
+
for bucket in sorted(qa_candidates):
|
| 164 |
+
samples.extend(row for _, row in sorted(qa_candidates[bucket], key=lambda x: x[0])[:16])
|
| 165 |
+
# All six source families should contribute across all observed actions;
|
| 166 |
+
# up to 288 records are retained in one packed JSON QA artifact.
|
| 167 |
+
assert len(samples) >= 128, len(samples)
|
| 168 |
+
qa = {'status': 'PASS', 'kind': 'automated_structural_asr_consistency_not_human_listening',
|
| 169 |
+
'sampled_rows': len(samples), 'action_rows': dict(counts),
|
| 170 |
+
'samples': samples,
|
| 171 |
+
'checks': ['every audited UID has exactly one ASR-backed disposition',
|
| 172 |
+
'all relabel texts have >=2 lexical words and >=0.68 ordered support',
|
| 173 |
+
'no implausible ASR word rate or repeated-run relabel accepted',
|
| 174 |
+
'source worklist and audit SHA-256 pinned']}
|
| 175 |
+
qa_path = RESOLUTION_DIR / 'QA_SAMPLES.json'
|
| 176 |
+
qa_path.write_text(json.dumps(qa, ensure_ascii=False, indent=2) + '\n')
|
| 177 |
+
spec = {'status': 'PENDING_RENDER_QA', 'version': 'hvs-s3-target-text-v1',
|
| 178 |
+
'created_utc': datetime.now(timezone.utc).isoformat(),
|
| 179 |
+
'source_worklist_sha256': sha(OUT / 'MANIFEST.json'),
|
| 180 |
+
'audit_report_sha256': sha(audit_report_path),
|
| 181 |
+
'asr_manifest_sha256': sha(asr_manifest_path),
|
| 182 |
+
'changes_file': changes_path.name, 'changes_sha256': sha(changes_path),
|
| 183 |
+
'changed_rows': len(expected), 'action_rows': dict(counts),
|
| 184 |
+
'accepted_rows': EXPECTED_ELIGIBLE - counts['exclude'],
|
| 185 |
+
'excluded_rows': counts['exclude'],
|
| 186 |
+
'excluded_target_hours': round(excluded_hours, 3),
|
| 187 |
+
'unresolved_impossible_rows': 0,
|
| 188 |
+
'by_source': {k: dict(v) for k, v in by_source.items()},
|
| 189 |
+
'qa_status': qa['status'], 'qa_file': qa_path.name,
|
| 190 |
+
'qa_sha256': sha(qa_path),
|
| 191 |
+
'important_caveat': 'Automated ASR supports conservative clipping decisions; it is not a human-verified transcript of every sample.'}
|
| 192 |
+
manifest_path = RESOLUTION_DIR / 'MANIFEST.json'
|
| 193 |
+
temp = manifest_path.with_suffix('.json.partial')
|
| 194 |
+
temp.write_text(json.dumps(spec, ensure_ascii=False, indent=2) + '\n')
|
| 195 |
+
os.replace(temp, manifest_path)
|
| 196 |
+
print(json.dumps({'status': spec['status'], 'suspects': len(expected),
|
| 197 |
+
'actions': dict(counts), 'excluded_hours': spec['excluded_target_hours'],
|
| 198 |
+
'qa_samples': len(samples)}), flush=True)
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
if __name__ == '__main__':
|
| 202 |
+
main()
|
code/ref15_fullft/build_worklist.py
ADDED
|
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Build a packed, unique-audio S3/tier7 reference-only LoRA worklist.
|
| 3 |
+
|
| 4 |
+
The old S3 caption curriculum has many presentations per target. Family 0
|
| 5 |
+
identifies one presentation per physical target. We use that base-index set,
|
| 6 |
+
then join *all* validated reference-eligible UIDs against the frozen reference
|
| 7 |
+
assets, not merely the historical 50% reference draw. Four CPU processes each
|
| 8 |
+
write 32 Parquet containers; no per-sample files or decoded audio are written.
|
| 9 |
+
"""
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from collections import Counter
|
| 13 |
+
import hashlib
|
| 14 |
+
import json
|
| 15 |
+
import math
|
| 16 |
+
import os
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
|
| 19 |
+
import numpy as np
|
| 20 |
+
import pyarrow as pa
|
| 21 |
+
import pyarrow.parquet as pq
|
| 22 |
+
|
| 23 |
+
SC = Path('/e/scratch/reformo/schuhmann1_moss')
|
| 24 |
+
CURRICULUM = SC / 'out/m2_600m_ladder/caption_curriculum_v1'
|
| 25 |
+
LOGICAL = CURRICULUM / 'manifests/S3/manifest.json'
|
| 26 |
+
OUT = SC / 'out/hvs_ref_only_r128_v1/worklist'
|
| 27 |
+
SHARDS = 128
|
| 28 |
+
WORLD = 4
|
| 29 |
+
EXPECTED_UNIQUE = 7_216_621
|
| 30 |
+
EXPECTED_ELIGIBLE = 3_707_948
|
| 31 |
+
MISSING = np.iinfo(np.uint32).max
|
| 32 |
+
# MOSS v2 has 80 ms frames. A future reference-conditioned run should allow
|
| 33 |
+
# 15-20 s, not inherit the historical 37-frame (2.96 s) crop. Prefer >=5 s
|
| 34 |
+
# where a verified same-speaker alternative exists; retain shorter frozen
|
| 35 |
+
# references here only to honor the requested all-eligible-S3 coverage and
|
| 36 |
+
# report them explicitly. 187 frames = 14.96 s; 188 would exceed 15 s.
|
| 37 |
+
MAX_REF_FRAMES = 187
|
| 38 |
+
PREFERRED_MIN_REF_FRAMES = 63
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def reference_bin(frames):
|
| 42 |
+
if frames < 50:
|
| 43 |
+
return 'under_4s'
|
| 44 |
+
if frames < PREFERRED_MIN_REF_FRAMES:
|
| 45 |
+
return '4_to_under_5s'
|
| 46 |
+
if frames < 125:
|
| 47 |
+
return '5_to_under_10s'
|
| 48 |
+
if frames < 188:
|
| 49 |
+
return '10_to_under_15s'
|
| 50 |
+
return '15s_or_longer'
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def sha(path):
|
| 54 |
+
h = hashlib.sha256()
|
| 55 |
+
with Path(path).open('rb') as stream:
|
| 56 |
+
for block in iter(lambda: stream.read(8 << 20), b''):
|
| 57 |
+
h.update(block)
|
| 58 |
+
return h.hexdigest()
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def read_row(fd, offsets, index):
|
| 62 |
+
a, b = int(offsets[index]), int(offsets[index + 1])
|
| 63 |
+
return json.loads(os.pread(fd, b - a, a))
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def uid_validation(uid):
|
| 67 |
+
return int.from_bytes(hashlib.blake2b(uid.encode(), digest_size=8).digest(), 'big') % 200 == 0
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def main():
|
| 71 |
+
rank = int(os.environ['SLURM_PROCID'])
|
| 72 |
+
assert int(os.environ['SLURM_NTASKS']) == WORLD and 0 <= rank < WORLD
|
| 73 |
+
logical = json.loads(LOGICAL.read_text())
|
| 74 |
+
assert logical['status'] == 'complete' and logical['stage'] == 'S3'
|
| 75 |
+
assert logical['unique_audio_occurrences'] == EXPECTED_UNIQUE
|
| 76 |
+
assert int(logical['family_reference_eligible_counts']['0']) == EXPECTED_ELIGIBLE
|
| 77 |
+
base_manifest = Path(logical['base_manifest'])
|
| 78 |
+
base_spec = json.loads(base_manifest.read_text())
|
| 79 |
+
assert base_spec['status'] == 'complete' and base_spec['ladder_stage'] == 'tier7'
|
| 80 |
+
asset_manifest = Path(logical['assets_manifest'])
|
| 81 |
+
assets = json.loads(asset_manifest.read_text())
|
| 82 |
+
assert assets['status'] == 'complete' and assets['references']['rows'] > EXPECTED_ELIGIBLE
|
| 83 |
+
descriptors = np.load(LOGICAL.parent / logical['descriptors'], mmap_mode='r', allow_pickle=False)
|
| 84 |
+
assert len(descriptors) == logical['presentations']
|
| 85 |
+
base_ids = np.sort(np.asarray(descriptors['base'][descriptors['family'] == 0], dtype=np.uint32))
|
| 86 |
+
assert len(base_ids) == EXPECTED_UNIQUE and len(np.unique(base_ids)) == EXPECTED_UNIQUE
|
| 87 |
+
base_root = base_manifest.parent
|
| 88 |
+
base_offsets = np.load(base_root / base_spec['offsets'], mmap_mode='r', allow_pickle=False)
|
| 89 |
+
assert len(base_offsets) == base_spec['presentations'] + 1
|
| 90 |
+
base_fd = os.open(base_root / base_spec['records'], os.O_RDONLY)
|
| 91 |
+
asset_root = asset_manifest.parent
|
| 92 |
+
refs = assets['references']
|
| 93 |
+
ref_keys = np.load(asset_root / refs['keys'], mmap_mode='r', allow_pickle=False)
|
| 94 |
+
ref_values = np.load(asset_root / refs['values'], mmap_mode='r', allow_pickle=False)
|
| 95 |
+
ref_offsets = np.load(asset_root / refs['offsets'], mmap_mode='r', allow_pickle=False)
|
| 96 |
+
ref_fd = os.open(asset_root / refs['records'], os.O_RDONLY)
|
| 97 |
+
assert len(ref_keys) == len(ref_values) == refs['rows']
|
| 98 |
+
assert len(ref_offsets) == len(ref_keys) + 1
|
| 99 |
+
assert ref_keys.dtype == np.dtype('S128')
|
| 100 |
+
OUT.mkdir(parents=True, exist_ok=True)
|
| 101 |
+
width = math.ceil(base_spec['presentations'] / SHARDS)
|
| 102 |
+
try:
|
| 103 |
+
for shard in range(rank, SHARDS, WORLD):
|
| 104 |
+
output = OUT / f'part-{shard:03d}.parquet'
|
| 105 |
+
proof = OUT / f'part-{shard:03d}.PASS.json'
|
| 106 |
+
if output.exists() and proof.exists():
|
| 107 |
+
prior = json.loads(proof.read_text())
|
| 108 |
+
assert prior['status'] == 'PASS' and prior['shard'] == shard
|
| 109 |
+
assert prior['output_sha256'] == sha(output)
|
| 110 |
+
print(json.dumps({'shard': shard, 'status': 'ALREADY_COMMITTED'}), flush=True)
|
| 111 |
+
continue
|
| 112 |
+
assert not output.exists() and not proof.exists(), f'incomplete shard {shard}'
|
| 113 |
+
first, last = shard * width, min((shard + 1) * width, base_spec['presentations'])
|
| 114 |
+
a, b = np.searchsorted(base_ids, [first, last])
|
| 115 |
+
selected = base_ids[a:b]
|
| 116 |
+
targets = [read_row(base_fd, base_offsets, int(index)) for index in selected]
|
| 117 |
+
queries = np.asarray([row['uid'].encode() for row in targets], dtype='S128')
|
| 118 |
+
positions = np.searchsorted(ref_keys, queries)
|
| 119 |
+
valid = positions < len(ref_keys)
|
| 120 |
+
valid[valid] = ref_keys[positions[valid]] == queries[valid]
|
| 121 |
+
output_rows = []
|
| 122 |
+
sources, source_hours = Counter(), Counter()
|
| 123 |
+
ref_bins, short_sources = Counter(), Counter()
|
| 124 |
+
for index in np.flatnonzero(valid):
|
| 125 |
+
target = targets[int(index)]
|
| 126 |
+
uid = target['uid']
|
| 127 |
+
assert len(uid.encode()) <= 128
|
| 128 |
+
reference = read_row(ref_fd, ref_offsets, int(ref_values[int(positions[index])]))
|
| 129 |
+
assert reference['uid'] == uid and reference['reference_uid'] != uid
|
| 130 |
+
locator = dict(reference['locator'])
|
| 131 |
+
original_crop = int(locator['crop_frames'])
|
| 132 |
+
full_frames = int(locator['full_frames'])
|
| 133 |
+
assert 0 < original_crop <= 37 <= full_frames
|
| 134 |
+
# The original codec shard retains all frames. Only the old
|
| 135 |
+
# locator's crop field imposed the 2.96-second ceiling.
|
| 136 |
+
used_frames = min(full_frames, MAX_REF_FRAMES)
|
| 137 |
+
locator['crop_start_frame'] = 0
|
| 138 |
+
locator['crop_frames'] = used_frames
|
| 139 |
+
assert target['record_locator'] and target['target_locator']
|
| 140 |
+
duration = float(target['dur_s'])
|
| 141 |
+
assert math.isfinite(duration) and duration > 0
|
| 142 |
+
output_rows.append({
|
| 143 |
+
'uid': uid, 'base_index': int(selected[index]), 'source': target['source'],
|
| 144 |
+
'duration_s': duration, 'record_locator': target['record_locator'],
|
| 145 |
+
'target_locator': target['target_locator'],
|
| 146 |
+
'reference_uid': reference['reference_uid'],
|
| 147 |
+
'reference_locator': locator,
|
| 148 |
+
'reference_full_frames': full_frames,
|
| 149 |
+
'reference_used_frames': used_frames,
|
| 150 |
+
'reference_under_preferred_5s': used_frames < PREFERRED_MIN_REF_FRAMES,
|
| 151 |
+
'split': 'validation' if uid_validation(uid) else 'train',
|
| 152 |
+
})
|
| 153 |
+
sources[target['source']] += 1
|
| 154 |
+
source_hours[target['source']] += duration / 3600
|
| 155 |
+
ref_bins[reference_bin(full_frames)] += 1
|
| 156 |
+
if used_frames < PREFERRED_MIN_REF_FRAMES:
|
| 157 |
+
short_sources[target['source']] += 1
|
| 158 |
+
temporary = output.with_suffix('.parquet.partial')
|
| 159 |
+
# Some contiguous base-index ranges have no eligible reference.
|
| 160 |
+
# A zero-row Parquet shard is valid and keeps the 128-shard index
|
| 161 |
+
# deterministic; it must not abort a 3.7M-row job after hours.
|
| 162 |
+
table = (pa.Table.from_pylist(output_rows) if output_rows else
|
| 163 |
+
pa.table({'uid': pa.array([], type=pa.string())}))
|
| 164 |
+
pq.write_table(table, temporary,
|
| 165 |
+
compression='zstd', row_group_size=4096)
|
| 166 |
+
os.replace(temporary, output)
|
| 167 |
+
result = {'status': 'PASS', 'shard': shard,
|
| 168 |
+
'unique_s3_targets_examined': len(targets),
|
| 169 |
+
'reference_eligible_rows': len(output_rows),
|
| 170 |
+
'reference_eligible_hours': sum(source_hours.values()),
|
| 171 |
+
'source_rows': dict(sources), 'source_hours': dict(source_hours),
|
| 172 |
+
'reference_full_duration_bins': dict(ref_bins),
|
| 173 |
+
'reference_under_preferred_5s_by_source': dict(short_sources),
|
| 174 |
+
'reference_max_frames': MAX_REF_FRAMES,
|
| 175 |
+
'reference_preferred_min_frames': PREFERRED_MIN_REF_FRAMES,
|
| 176 |
+
'validation_rows': sum(row['split'] == 'validation' for row in output_rows),
|
| 177 |
+
'output_sha256': sha(output), 'per_sample_files_created': 0}
|
| 178 |
+
proof_tmp = proof.with_suffix('.json.partial')
|
| 179 |
+
proof_tmp.write_text(json.dumps(result, indent=2) + '\n')
|
| 180 |
+
os.replace(proof_tmp, proof)
|
| 181 |
+
print(json.dumps(result), flush=True)
|
| 182 |
+
finally:
|
| 183 |
+
os.close(base_fd)
|
| 184 |
+
os.close(ref_fd)
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
if __name__ == '__main__':
|
| 188 |
+
main()
|
code/ref15_fullft/preflight.py
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Require a successful full-parameter GPU canary and immutable S3 inputs."""
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import json
|
| 6 |
+
import math
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
import sys
|
| 9 |
+
|
| 10 |
+
HERE = Path(__file__).resolve().parent
|
| 11 |
+
HOME = HERE.parent
|
| 12 |
+
sys.path.insert(0, str(HOME / 'hvs_ref_only_r128_v1'))
|
| 13 |
+
from build_worklist import sha
|
| 14 |
+
from text_resolution import RESOLUTION_MANIFEST
|
| 15 |
+
|
| 16 |
+
SC = Path('/e/scratch/reformo/schuhmann1_moss')
|
| 17 |
+
WORKLIST = SC / 'out/hvs_ref_only_r128_v1/worklist/MANIFEST.json'
|
| 18 |
+
SMOKE = SC / 'out/hvs_ref_full_ft_v1/smoke_4gpu'
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def main():
|
| 22 |
+
source = json.loads(WORKLIST.read_text())
|
| 23 |
+
resolution = json.loads(RESOLUTION_MANIFEST.read_text())
|
| 24 |
+
assert source['status'] == 'COMPLETE' and source['eligible_unique_targets'] == 3_707_948
|
| 25 |
+
assert resolution['status'] == 'PASS' and resolution['source_worklist_sha256'] == sha(WORKLIST)
|
| 26 |
+
scores = [json.loads(line) for line in (SMOKE / 'metrics.jsonl').read_text().splitlines()]
|
| 27 |
+
trains = [row for row in scores if row.get('kind') == 'train']
|
| 28 |
+
validations = [row for row in scores if row.get('kind') == 'validation']
|
| 29 |
+
assert [row['step'] for row in trains[-2:]] == [1, 2]
|
| 30 |
+
assert [row['step'] for row in validations[-3:]] == [0, 1, 2]
|
| 31 |
+
assert all(math.isfinite(row['grad_norm']) and row['grad_norm'] > 0 for row in trains[-2:])
|
| 32 |
+
assert all(math.isfinite(row['nll_per_frame']) for row in validations[-3:])
|
| 33 |
+
assert all(row['reference_longer_than_2_96s'] > 0 for row in trains[-2:])
|
| 34 |
+
assert all(row['peak_allocated_gib'] < 90 for row in trains[-2:])
|
| 35 |
+
latest = json.loads((SMOKE / 'checkpoints/current.json').read_text())
|
| 36 |
+
assert latest['step'] == 2
|
| 37 |
+
complete = json.loads((SMOKE / 'checkpoints' / latest['directory'] / 'COMPLETE.json').read_text())
|
| 38 |
+
assert complete['step'] == 2 and complete['state_bytes'] > 1_000_000_000
|
| 39 |
+
print(json.dumps({'status': 'PASS', 'steps': 2, 'peak_gib': max(row['peak_allocated_gib'] for row in trains),
|
| 40 |
+
'validation_step2': validations[-1]['nll_per_frame'],
|
| 41 |
+
'worklist_sha256': sha(WORKLIST), 'resolution_sha256': sha(RESOLUTION_MANIFEST)}), flush=True)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
if __name__ == '__main__':
|
| 45 |
+
main()
|
code/ref15_fullft/prompt_mix.py
ADDED
|
@@ -0,0 +1,136 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Deterministic one-presentation S3 reference-conditioned prompt mixture.
|
| 2 |
+
|
| 3 |
+
Targets: ~55% local SCRIPT, 20% GENERAL+SCRIPT, 20% CAPTION+TRANSCRIPT
|
| 4 |
+
(Timbre or Body Whisper 1.1), 5% TRANSCRIPT only. Every mode still receives
|
| 5 |
+
the distinct reference audio, cropped to <=187 MOSS-v2 frames by the worklist.
|
| 6 |
+
The complete caption annotation exists for only ~52% of eligible S3 targets;
|
| 7 |
+
sample caption mode from those UIDs at the compensating rate, not from empty
|
| 8 |
+
or fabricated captions. Assignment is a stable hash of target UID so resume,
|
| 9 |
+
world size, or row ordering cannot change a target's presentation.
|
| 10 |
+
"""
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import hashlib
|
| 14 |
+
import json
|
| 15 |
+
import os
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
import re
|
| 18 |
+
|
| 19 |
+
import numpy as np
|
| 20 |
+
|
| 21 |
+
from build_worklist import EXPECTED_ELIGIBLE, LOGICAL
|
| 22 |
+
|
| 23 |
+
EXPECTED_ANNOTATED_ELIGIBLE = 1_945_226
|
| 24 |
+
CAPTION_DRAW_WITHIN_ANNOTATED = .20 * EXPECTED_ELIGIBLE / EXPECTED_ANNOTATED_ELIGIBLE
|
| 25 |
+
MODES = ('script_only', 'general_script', 'caption_timbre', 'caption_bude', 'transcript_only')
|
| 26 |
+
_BAD_BODY = re.compile(r'\b(?:only silence|no speech|no spoken words|no voice)\b', re.I)
|
| 27 |
+
SOURCE_PROMPT_SANITIZER_VERSION = 'hvs-source-legacy-pause-v1'
|
| 28 |
+
_LEGACY_PAUSE = re.compile(r'[ \t]*\[pause[ \t]+(?:\d+(?:\.\d+)?|\.\d+)[ \t]*s\][ \t]*', re.I)
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def sanitize_source_prompt(meta: dict) -> dict:
|
| 32 |
+
"""Drop obsolete numeric pause markers from source directions, not target text.
|
| 33 |
+
|
| 34 |
+
The timed renderer reconstructs pauses from word timestamps. A legacy
|
| 35 |
+
``[pause 0.3s]`` inside a source cue must not be copied into its new
|
| 36 |
+
sentence direction; the new timed/untimed contract rejects that marker.
|
| 37 |
+
Keep the original metadata and target transcript unchanged.
|
| 38 |
+
"""
|
| 39 |
+
source = str(meta.get('prompt') or '')
|
| 40 |
+
if not _LEGACY_PAUSE.search(source):
|
| 41 |
+
return meta
|
| 42 |
+
return dict(meta, prompt=_LEGACY_PAUSE.sub(' ', source))
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def mode_for_uid(uid: str, has_annotation: bool) -> str:
|
| 46 |
+
digest = hashlib.blake2b(f'hvs-ref15-mix-v1|{uid}'.encode(), digest_size=16).digest()
|
| 47 |
+
caption_draw = int.from_bytes(digest[:8], 'big') / 2**64
|
| 48 |
+
other_draw = int.from_bytes(digest[8:], 'big') / 2**64
|
| 49 |
+
if has_annotation and caption_draw < CAPTION_DRAW_WITHIN_ANNOTATED:
|
| 50 |
+
return 'caption_timbre' if other_draw < .5 else 'caption_bude'
|
| 51 |
+
if other_draw < .25: # 20% of all rows after the 20% caption draw
|
| 52 |
+
return 'general_script'
|
| 53 |
+
if other_draw < .3125: # 5% of all rows after the 20% caption draw
|
| 54 |
+
return 'transcript_only'
|
| 55 |
+
return 'script_only'
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
class AnnotationIndex:
|
| 59 |
+
def __init__(self):
|
| 60 |
+
logical = json.loads(LOGICAL.read_text())
|
| 61 |
+
manifest_path = Path(logical['assets_manifest'])
|
| 62 |
+
assets = json.loads(manifest_path.read_text())
|
| 63 |
+
root = manifest_path.parent
|
| 64 |
+
spec = assets['annotations']
|
| 65 |
+
self.keys = np.load(root / spec['keys'], mmap_mode='r', allow_pickle=False)
|
| 66 |
+
self.values = np.load(root / spec['values'], mmap_mode='r', allow_pickle=False)
|
| 67 |
+
self.offsets = np.load(root / spec['offsets'], mmap_mode='r', allow_pickle=False)
|
| 68 |
+
self.fd = os.open(root / spec['records'], os.O_RDONLY)
|
| 69 |
+
assert self.keys.dtype == np.dtype('S128')
|
| 70 |
+
assert len(self.keys) == len(self.values) == spec['rows']
|
| 71 |
+
assert len(self.offsets) == len(self.keys) + 1
|
| 72 |
+
|
| 73 |
+
def locate_many(self, uids: list[str]) -> np.ndarray:
|
| 74 |
+
if not uids:
|
| 75 |
+
return np.empty(0, dtype=np.int64)
|
| 76 |
+
query = np.asarray([uid.encode() for uid in uids], dtype='S128')
|
| 77 |
+
positions = np.searchsorted(self.keys, query)
|
| 78 |
+
valid = positions < len(self.keys)
|
| 79 |
+
valid[valid] = self.keys[positions[valid]] == query[valid]
|
| 80 |
+
result = np.full(len(query), -1, dtype=np.int64)
|
| 81 |
+
result[valid] = self.values[positions[valid]]
|
| 82 |
+
return result
|
| 83 |
+
|
| 84 |
+
def attach(self, rows: list[dict]) -> None:
|
| 85 |
+
positions = self.locate_many([row['uid'] for row in rows])
|
| 86 |
+
for row, index in zip(rows, positions):
|
| 87 |
+
row['_annotation_index'] = int(index)
|
| 88 |
+
row['_prompt_mode'] = mode_for_uid(row['uid'], index >= 0)
|
| 89 |
+
|
| 90 |
+
def row(self, index: int) -> dict:
|
| 91 |
+
assert index >= 0
|
| 92 |
+
start, end = int(self.offsets[index]), int(self.offsets[index + 1])
|
| 93 |
+
return json.loads(os.pread(self.fd, end-start, start))
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
def render_prompt(meta: dict, base_index: int, mode: str, annotation: dict | None = None):
|
| 97 |
+
from script_only import script_only
|
| 98 |
+
from timed_prompt import repair_prompt
|
| 99 |
+
|
| 100 |
+
meta = sanitize_source_prompt(meta)
|
| 101 |
+
|
| 102 |
+
if mode == 'script_only':
|
| 103 |
+
return script_only(meta, base_index)
|
| 104 |
+
if mode == 'general_script':
|
| 105 |
+
prompt, facts = repair_prompt(meta, caption=None, form='A', mode='reference',
|
| 106 |
+
presentation=int(base_index), phase='S3')
|
| 107 |
+
assert prompt.startswith('GENERAL:') and '\nSCRIPT:\n' in prompt
|
| 108 |
+
return prompt, dict(facts, stable_voice_tags_removed=0)
|
| 109 |
+
# The outer ASCII quotes delimit what should be spoken. Normalize any
|
| 110 |
+
# quotes already inside the source transcript so they cannot terminate
|
| 111 |
+
# that field and make following words look like instructions.
|
| 112 |
+
transcript = str(meta['text']).strip().replace('"', "'")
|
| 113 |
+
assert transcript, meta['uid']
|
| 114 |
+
if mode == 'transcript_only':
|
| 115 |
+
return f'TRANSCRIPT: "{transcript}"', {'timed': False, 'directions': 0,
|
| 116 |
+
'stable_voice_tags_removed': 0}
|
| 117 |
+
assert mode in ('caption_timbre', 'caption_bude') and annotation is not None
|
| 118 |
+
assert annotation['uid'] == meta['uid']
|
| 119 |
+
if mode == 'caption_timbre':
|
| 120 |
+
caption = str(annotation['timbre_prose']).strip()
|
| 121 |
+
source = 'timbre_whisper_prose'
|
| 122 |
+
else:
|
| 123 |
+
caption = str(annotation['bude']).strip()
|
| 124 |
+
source = 'body_whisper_1_1'
|
| 125 |
+
# Some stored Body Whisper 1.1 captions start with duplicated first
|
| 126 |
+
# words such as "InIn" or "TheThe". Repair that surface glitch only;
|
| 127 |
+
# preserve the underlying annotation otherwise.
|
| 128 |
+
caption = re.sub(r'^([A-Z][a-z]{1,12})\1\b', r'\1', caption)
|
| 129 |
+
if _BAD_BODY.search(caption):
|
| 130 |
+
caption = str(annotation['timbre_prose']).strip()
|
| 131 |
+
source = 'timbre_fallback_bad_body'
|
| 132 |
+
caption = ' '.join(caption.replace('"', "'").split())
|
| 133 |
+
assert caption and transcript
|
| 134 |
+
return f'CAPTION: {caption}\nTRANSCRIPT: "{transcript}"', {
|
| 135 |
+
'timed': False, 'directions': 0, 'stable_voice_tags_removed': 0,
|
| 136 |
+
'caption_source': source}
|
code/ref15_fullft/rowgroup_plan.py
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Balance packed Parquet row groups over GPU ranks without duplicating data.
|
| 2 |
+
|
| 3 |
+
The 128 source shards are *not* equal-sized. Assigning four complete shards
|
| 4 |
+
to each of 32 ranks gave the slowest rank ~38.5k steps versus ~28.8k expected,
|
| 5 |
+
and left some ranks idle at the tail. Row groups are ~4096 records and can be
|
| 6 |
+
partitioned by metadata without creating any per-example files.
|
| 7 |
+
"""
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
import pyarrow.parquet as pq
|
| 11 |
+
|
| 12 |
+
from build_worklist import OUT, SHARDS
|
| 13 |
+
|
| 14 |
+
SAMPLER_VERSION = 'hvs-rowgroup-balance-v1'
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
def plan(world: int):
|
| 18 |
+
assert world > 0
|
| 19 |
+
groups = []
|
| 20 |
+
for shard in range(SHARDS):
|
| 21 |
+
metadata = pq.read_metadata(OUT / f'part-{shard:03d}.parquet')
|
| 22 |
+
for group in range(metadata.num_row_groups):
|
| 23 |
+
n = metadata.row_group(group).num_rows
|
| 24 |
+
if n:
|
| 25 |
+
groups.append((n, shard, group))
|
| 26 |
+
assignments = [[] for _ in range(world)]
|
| 27 |
+
loads = [0] * world
|
| 28 |
+
for n, shard, group in sorted(groups, key=lambda x: (-x[0], x[1], x[2])):
|
| 29 |
+
rank = min(range(world), key=lambda r: (loads[r], r))
|
| 30 |
+
assignments[rank].append((shard, group))
|
| 31 |
+
loads[rank] += n
|
| 32 |
+
assert sum(loads) == sum(n for n, _, _ in groups)
|
| 33 |
+
assert len({item for rank_groups in assignments for item in rank_groups}) == len(groups)
|
| 34 |
+
return assignments, loads
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def read_rank_rows(rank: int, world: int):
|
| 38 |
+
assignments, loads = plan(world)
|
| 39 |
+
by_shard = {}
|
| 40 |
+
for shard, group in assignments[rank]:
|
| 41 |
+
by_shard.setdefault(shard, []).append(group)
|
| 42 |
+
rows = []
|
| 43 |
+
for shard, groups in sorted(by_shard.items()):
|
| 44 |
+
source = pq.ParquetFile(OUT / f'part-{shard:03d}.parquet')
|
| 45 |
+
for group in groups:
|
| 46 |
+
rows.extend(source.read_row_group(group).to_pylist())
|
| 47 |
+
assert len(rows) == loads[rank]
|
| 48 |
+
return rows, loads
|
code/ref15_fullft/script_only.py
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The exact caption-free, reference-conditioned S3 prompt surface.
|
| 2 |
+
|
| 3 |
+
Keep the inherited sentence-local performance cues, timing, pauses and vocal
|
| 4 |
+
bursts; remove the entire GENERAL block and all global/body captions. The
|
| 5 |
+
reference codec is provided separately through ScorePacker reference mode.
|
| 6 |
+
"""
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
import re
|
| 10 |
+
|
| 11 |
+
from timed_prompt import repair_prompt
|
| 12 |
+
|
| 13 |
+
RENDERER_VERSION = 'hvs-script-only-source-cue-filter-v2'
|
| 14 |
+
|
| 15 |
+
_STATIC = re.compile(
|
| 16 |
+
r"\b(?:male|female|masculine|feminine|gender|adult|teen|teenage|young|"
|
| 17 |
+
r"middle.?aged|elderly|old|child|childlike|baby|androgynous|"
|
| 18 |
+
r"woman|man|girl|boy|accent|register|pitch(?:ed)?|"
|
| 19 |
+
r"raspy|husky|nasal|heady|hoarse|gravelly|breathy|airy|resonan\w*|"
|
| 20 |
+
r"baritone|soprano|tenor|alto|bass|deep.?voice|thin.?voice|thick.?voice)\b",
|
| 21 |
+
re.I)
|
| 22 |
+
_LEGACY_LEADING_CUE = re.compile(r"^(\s*)\(([^()\n]+)\)(?=\s|$)")
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def _without_static_cue(match, removed):
|
| 26 |
+
tags = [tag.strip() for tag in match.group(2).split(',') if tag.strip()]
|
| 27 |
+
dynamic = [tag for tag in tags if not _STATIC.search(tag)]
|
| 28 |
+
removed[0] += len(tags) - len(dynamic)
|
| 29 |
+
return f"({', '.join(dynamic)})" if dynamic else ""
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def script_only(meta: dict, base_index: int) -> tuple[str, dict]:
|
| 33 |
+
# Filter only cues in the *source* prompt before rendering. Filtering the
|
| 34 |
+
# rendered script is unsafe: its first parenthesis may be the actual target
|
| 35 |
+
# audio's only event, e.g. a 0.6 s giggle with no aligned speech.
|
| 36 |
+
source_prompt = str(meta.get("prompt") or "")
|
| 37 |
+
removed = [0]
|
| 38 |
+
if "SCRIPT:\n" in source_prompt:
|
| 39 |
+
head, legacy_script = source_prompt.split("SCRIPT:\n", 1)
|
| 40 |
+
lines = []
|
| 41 |
+
for line in legacy_script.splitlines():
|
| 42 |
+
lines.append(_LEGACY_LEADING_CUE.sub(
|
| 43 |
+
lambda m: m.group(1) + _without_static_cue(m, removed), line,
|
| 44 |
+
count=1))
|
| 45 |
+
source_prompt = head + "SCRIPT:\n" + "\n".join(lines)
|
| 46 |
+
filtered_meta = dict(meta, prompt=source_prompt)
|
| 47 |
+
repaired, facts = repair_prompt(
|
| 48 |
+
filtered_meta, caption=None, form="A", mode="reference",
|
| 49 |
+
presentation=int(base_index), phase="S3")
|
| 50 |
+
marker = "SCRIPT:\n"
|
| 51 |
+
if marker not in repaired:
|
| 52 |
+
raise ValueError(f"No structured SCRIPT in {meta['uid']}")
|
| 53 |
+
body = repaired.split(marker, 1)[1].strip()
|
| 54 |
+
body = body.strip()
|
| 55 |
+
if not body or "GENERAL:" in body or "CAPTION:" in body or "TRANSCRIPT:" in body:
|
| 56 |
+
raise ValueError(f"Leaked caption/header in {meta['uid']}")
|
| 57 |
+
if not meta.get("text"):
|
| 58 |
+
raise ValueError(f"Empty transcript in {meta['uid']}")
|
| 59 |
+
facts = dict(facts, stable_voice_tags_removed=removed[0])
|
| 60 |
+
return body, facts
|
code/ref15_fullft/text_resolution.py
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Packed per-UID disposition sidecar for the S3 target/transcript audit.
|
| 2 |
+
|
| 3 |
+
The original worklist and audio stay immutable. Every audited suspect gets
|
| 4 |
+
an explicit disposition, including `verified_keep` for a false-positive flag.
|
| 5 |
+
An absent UID is an unaudited clean row and keeps its source text.
|
| 6 |
+
"""
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
import json
|
| 10 |
+
from collections import Counter
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
import pyarrow.parquet as pq
|
| 14 |
+
|
| 15 |
+
from build_worklist import EXPECTED_ELIGIBLE, OUT, sha
|
| 16 |
+
|
| 17 |
+
ROOT = OUT.parent
|
| 18 |
+
RESOLUTION_DIR = ROOT / 'target_text_resolution_v1'
|
| 19 |
+
RESOLUTION_MANIFEST = RESOLUTION_DIR / 'MANIFEST.json'
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def load_resolution():
|
| 23 |
+
spec = json.loads(RESOLUTION_MANIFEST.read_text())
|
| 24 |
+
assert spec['status'] == 'PASS' and spec['version'] == 'hvs-s3-target-text-v1'
|
| 25 |
+
assert spec['source_worklist_sha256'] == sha(OUT / 'MANIFEST.json')
|
| 26 |
+
assert spec['accepted_rows'] + spec['excluded_rows'] == EXPECTED_ELIGIBLE
|
| 27 |
+
assert spec['qa_status'] == 'PASS'
|
| 28 |
+
qa_path = RESOLUTION_DIR / spec['qa_file']
|
| 29 |
+
assert qa_path.parent == RESOLUTION_DIR and sha(qa_path) == spec['qa_sha256']
|
| 30 |
+
qa = json.loads(qa_path.read_text())
|
| 31 |
+
assert qa['status'] == 'PASS' and qa['sampled_rows'] >= 128
|
| 32 |
+
render_path = RESOLUTION_DIR / spec['render_qa_file']
|
| 33 |
+
assert render_path.parent == RESOLUTION_DIR and sha(render_path) == spec['render_qa_sha256']
|
| 34 |
+
render_qa = json.loads(render_path.read_text())
|
| 35 |
+
assert render_qa['status'] == 'PASS' and render_qa['sampled_rows'] >= 128
|
| 36 |
+
assert spec['unresolved_impossible_rows'] == 0
|
| 37 |
+
assert spec['audit_report_sha256'] == sha(ROOT / 'target_text_audit_v1' / 'REPORT.json')
|
| 38 |
+
path = RESOLUTION_DIR / spec['changes_file']
|
| 39 |
+
assert path.parent == RESOLUTION_DIR and sha(path) == spec['changes_sha256']
|
| 40 |
+
rows = pq.read_table(path, columns=['uid', 'action', 'text']).to_pylist()
|
| 41 |
+
assert len(rows) == spec['changed_rows']
|
| 42 |
+
changes = {}
|
| 43 |
+
actions = Counter()
|
| 44 |
+
for row in rows:
|
| 45 |
+
uid, action, text = row['uid'], row['action'], row['text']
|
| 46 |
+
assert uid and uid not in changes
|
| 47 |
+
assert action in ('relabel', 'burst_only', 'exclude', 'verified_keep')
|
| 48 |
+
actions[action] += 1
|
| 49 |
+
if action in ('exclude', 'verified_keep'):
|
| 50 |
+
assert not text
|
| 51 |
+
else:
|
| 52 |
+
assert text and text.strip()
|
| 53 |
+
changes[uid] = (action, text)
|
| 54 |
+
assert actions['exclude'] == spec['excluded_rows']
|
| 55 |
+
assert dict(actions) == spec['action_rows']
|
| 56 |
+
return spec, changes
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def apply_to_rows(rows, changes):
|
| 60 |
+
kept = []
|
| 61 |
+
for row in rows:
|
| 62 |
+
change = changes.get(row['uid'])
|
| 63 |
+
if change is None:
|
| 64 |
+
kept.append(row)
|
| 65 |
+
elif change[0] == 'verified_keep':
|
| 66 |
+
kept.append(row)
|
| 67 |
+
elif change[0] in ('relabel', 'burst_only'):
|
| 68 |
+
row['_resolved_text'] = change[1]
|
| 69 |
+
row['_resolution_action'] = change[0]
|
| 70 |
+
kept.append(row)
|
| 71 |
+
return kept
|
code/ref15_fullft/train.py
ADDED
|
@@ -0,0 +1,358 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Full-parameter S3 long-reference continuation, matched to the Rank-128 LoRA data.
|
| 3 |
+
|
| 4 |
+
All model weights remain FP32 masters; forward/loss use the original BF16
|
| 5 |
+
autocast path. The source worklist, UID order, prompt renderer, resolved text,
|
| 6 |
+
15-second reference packing and audio-token objective are unchanged.
|
| 7 |
+
"""
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
from contextlib import nullcontext
|
| 11 |
+
from datetime import timedelta
|
| 12 |
+
import json
|
| 13 |
+
import math
|
| 14 |
+
import os
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
import signal
|
| 17 |
+
import sys
|
| 18 |
+
import time
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
import pyarrow.parquet as pq
|
| 22 |
+
import torch
|
| 23 |
+
import torch.distributed as dist
|
| 24 |
+
from torch import nn
|
| 25 |
+
from torch.nn.parallel import DistributedDataParallel as DDP
|
| 26 |
+
from transformers import AutoTokenizer
|
| 27 |
+
|
| 28 |
+
HERE = Path(__file__).resolve().parent
|
| 29 |
+
HOME = HERE.parent
|
| 30 |
+
SC = Path('/e/scratch/reformo/schuhmann1_moss')
|
| 31 |
+
OLD = HOME / 'hvs_ref_only_r128_v1'
|
| 32 |
+
VENDOR = HOME / 'm2_600m_tts_training/vendor/m2_20k_score'
|
| 33 |
+
sys.path[:0] = [str(OLD), str(HOME), str(HOME / 'm2_dpo_smoke_v1'), str(VENDOR)]
|
| 34 |
+
|
| 35 |
+
import run_smoke as core
|
| 36 |
+
import state_io
|
| 37 |
+
from continuation_control import complete_checkpoint, latest_checkpoint, sha
|
| 38 |
+
from build_worklist import MAX_REF_FRAMES, PREFERRED_MIN_REF_FRAMES
|
| 39 |
+
from prompt_mix import AnnotationIndex, MODES, SOURCE_PROMPT_SANITIZER_VERSION
|
| 40 |
+
from qa_worklist import Reader
|
| 41 |
+
from rowgroup_plan import read_rank_rows, SAMPLER_VERSION
|
| 42 |
+
from script_only import RENDERER_VERSION
|
| 43 |
+
from text_resolution import apply_to_rows, load_resolution, RESOLUTION_MANIFEST
|
| 44 |
+
from train import packed, mean_validation
|
| 45 |
+
|
| 46 |
+
WORKLIST = SC / 'out/hvs_ref_only_r128_v1/worklist'
|
| 47 |
+
BASE = SC / 'out/m2_600m_ladder/caption_curriculum_v1/training/checkpoints/step-00032337'
|
| 48 |
+
SCHEMA = SC / 'out/m2_20k_score/score_schema.json'
|
| 49 |
+
ROOT = SC / 'out/hvs_ref_full_ft_v1'
|
| 50 |
+
SEED = 20261004
|
| 51 |
+
PEAK_BACKBONE = 8e-6
|
| 52 |
+
PEAK_TALKER = 1.6e-5
|
| 53 |
+
WARMUP_FRACTION = .05
|
| 54 |
+
FLOOR_RATIO = .10
|
| 55 |
+
SAVE_EVERY = 1000
|
| 56 |
+
STOP = False
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
class Step(nn.Module):
|
| 60 |
+
"""DDP boundary around exactly the original teacher-forced audio loss."""
|
| 61 |
+
|
| 62 |
+
def __init__(self, base, config, device):
|
| 63 |
+
super().__init__()
|
| 64 |
+
self.base = base
|
| 65 |
+
self.config = config
|
| 66 |
+
self.device = device
|
| 67 |
+
|
| 68 |
+
def forward(self, batch):
|
| 69 |
+
logprob, _ = core.logprob(self.base, self.config, batch, self.device)
|
| 70 |
+
return -logprob
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def stop(_signal, _frame):
|
| 74 |
+
global STOP
|
| 75 |
+
STOP = True
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def atomic_json(path: Path, value: dict) -> None:
|
| 79 |
+
temporary = path.with_suffix(path.suffix + '.partial')
|
| 80 |
+
temporary.write_text(json.dumps(value, ensure_ascii=False, allow_nan=False, indent=2) + '\n')
|
| 81 |
+
os.replace(temporary, path)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def lr_factor(step: int, total: int) -> float:
|
| 85 |
+
warmup = max(1, round(total * WARMUP_FRACTION))
|
| 86 |
+
if step <= warmup:
|
| 87 |
+
return step / warmup
|
| 88 |
+
progress = min(1., (step - warmup) / max(1, total - warmup))
|
| 89 |
+
return FLOOR_RATIO + (1 - FLOOR_RATIO) * .5 * (1 + math.cos(math.pi * progress))
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def prepare_rows(rank, world, changes, annotations):
|
| 93 |
+
assigned, loads = read_rank_rows(rank, world)
|
| 94 |
+
assigned = apply_to_rows(assigned, changes)
|
| 95 |
+
rows = [row for row in assigned if row['split'] == 'train']
|
| 96 |
+
assert rows
|
| 97 |
+
annotations.attach(rows)
|
| 98 |
+
validation = []
|
| 99 |
+
for shard in range(128):
|
| 100 |
+
for row in pq.read_table(WORKLIST / f'part-{shard:03d}.parquet').to_pylist():
|
| 101 |
+
if row['split'] == 'validation' and len(validation) < world * 16:
|
| 102 |
+
validation.append(row)
|
| 103 |
+
if len(validation) >= world * 16:
|
| 104 |
+
break
|
| 105 |
+
assert len(validation) == world * 16
|
| 106 |
+
validation = apply_to_rows(validation, changes)
|
| 107 |
+
annotations.attach(validation)
|
| 108 |
+
for cohort in (rows, validation):
|
| 109 |
+
for row in cohort:
|
| 110 |
+
if row.get('_resolution_action') == 'burst_only':
|
| 111 |
+
row['_prompt_mode'] = 'script_only'
|
| 112 |
+
elif row.get('_resolution_action') == 'relabel' and row['_prompt_mode'] in (
|
| 113 |
+
'general_script', 'caption_bude'):
|
| 114 |
+
row['_prompt_mode'] = 'script_only'
|
| 115 |
+
return rows, validation, loads
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def main() -> None:
|
| 119 |
+
smoke = os.environ.get('HVS_FULL_FT_SMOKE') == '1'
|
| 120 |
+
rank, world, local = (int(os.environ[key]) for key in
|
| 121 |
+
('SLURM_PROCID', 'SLURM_NTASKS', 'SLURM_LOCALID'))
|
| 122 |
+
assert world == (4 if smoke else 32)
|
| 123 |
+
visible = torch.cuda.device_count()
|
| 124 |
+
assert visible in (1, 4), visible
|
| 125 |
+
device_index = 0 if visible == 1 else local
|
| 126 |
+
device = torch.device('cuda', device_index)
|
| 127 |
+
torch.cuda.set_device(device)
|
| 128 |
+
torch.set_num_threads(4)
|
| 129 |
+
dist.init_process_group('nccl', init_method='env://', rank=rank, world_size=world,
|
| 130 |
+
timeout=timedelta(minutes=30), device_id=device)
|
| 131 |
+
signal.signal(signal.SIGUSR1, stop)
|
| 132 |
+
signal.signal(signal.SIGTERM, stop)
|
| 133 |
+
started = time.monotonic()
|
| 134 |
+
|
| 135 |
+
manifest_path = WORKLIST / 'MANIFEST.json'
|
| 136 |
+
manifest = json.loads(manifest_path.read_text())
|
| 137 |
+
assert manifest['status'] == 'COMPLETE' and manifest['eligible_unique_targets'] == 3_707_948
|
| 138 |
+
assert manifest['reference_max_frames'] == MAX_REF_FRAMES == 187
|
| 139 |
+
assert manifest['reference_preferred_min_frames'] == PREFERRED_MIN_REF_FRAMES == 63
|
| 140 |
+
assert manifest['prompt_mix_version'] == 'hvs-ref15-mix-v1'
|
| 141 |
+
resolution, changes = load_resolution()
|
| 142 |
+
assert resolution['status'] == 'PASS'
|
| 143 |
+
qa = json.loads((WORKLIST.parent / 'QA_SAMPLES.json').read_text())
|
| 144 |
+
assert qa['status'] == 'PASS' and qa['sampled_rows'] == 64
|
| 145 |
+
assert set(qa['prompt_mode_rows']) == set(MODES)
|
| 146 |
+
annotations = AnnotationIndex()
|
| 147 |
+
rows, validation, loads = prepare_rows(rank, world, changes, annotations)
|
| 148 |
+
effective = torch.tensor([sum(row['_prompt_mode'] == mode for row in rows)
|
| 149 |
+
for mode in MODES], dtype=torch.int64, device=device)
|
| 150 |
+
dist.all_reduce(effective)
|
| 151 |
+
prompt_counts = {mode: int(effective[i].item()) for i, mode in enumerate(MODES)}
|
| 152 |
+
assert sum(prompt_counts.values()) <= resolution['accepted_rows']
|
| 153 |
+
accum = 2 if smoke else 4
|
| 154 |
+
local_steps = torch.tensor([math.ceil(len(rows) / accum)], dtype=torch.int64, device=device)
|
| 155 |
+
dist.all_reduce(local_steps, op=dist.ReduceOp.MAX)
|
| 156 |
+
total = 2 if smoke else int(local_steps.item())
|
| 157 |
+
assert len(rows) >= accum * total if smoke else total > 1000
|
| 158 |
+
if not smoke:
|
| 159 |
+
assert total == 26774, total
|
| 160 |
+
order = np.random.default_rng(SEED + rank).permutation(len(rows))
|
| 161 |
+
torch.manual_seed(SEED)
|
| 162 |
+
torch.cuda.manual_seed_all(SEED)
|
| 163 |
+
|
| 164 |
+
schema = json.loads(SCHEMA.read_text())
|
| 165 |
+
model, config = core.build_fresh(schema, log=lambda _: None)
|
| 166 |
+
state = torch.load(BASE / 'model_bf16.pt', map_location='cpu', weights_only=True)
|
| 167 |
+
model.load_state_dict(state, strict=True)
|
| 168 |
+
model.tie_weights()
|
| 169 |
+
del state
|
| 170 |
+
# FP32 master parameters are necessary at these small continuation LRs:
|
| 171 |
+
# tiny updates to BF16 parameter tensors can round away completely.
|
| 172 |
+
model = model.to(device=device, dtype=torch.float32).eval()
|
| 173 |
+
assert all(model.audio_lm_heads[i].weight.data_ptr() ==
|
| 174 |
+
model.audio_embeddings[i].weight.data_ptr() for i in range(12))
|
| 175 |
+
# The inherited score conditioner is unused by this audio-caption loss;
|
| 176 |
+
# retain its original frozen state. Every *active* semantic and Talker
|
| 177 |
+
# parameter is fully fine-tuned, without a LoRA or adapter.
|
| 178 |
+
for name, parameter in model.named_parameters():
|
| 179 |
+
parameter.requires_grad_(not name.startswith('score_conditioner.'))
|
| 180 |
+
groups = {'backbone': [], 'talker': []}
|
| 181 |
+
for name, parameter in model.named_parameters():
|
| 182 |
+
if not parameter.requires_grad:
|
| 183 |
+
continue
|
| 184 |
+
key = 'backbone' if name.startswith(('transformer.', 'text_lm_head.')) else 'talker'
|
| 185 |
+
groups[key].append(parameter)
|
| 186 |
+
assert all(groups.values())
|
| 187 |
+
optimizer = torch.optim.AdamW([
|
| 188 |
+
{'name': 'backbone', 'params': groups['backbone'], 'lr': PEAK_BACKBONE},
|
| 189 |
+
{'name': 'talker', 'params': groups['talker'], 'lr': PEAK_TALKER},
|
| 190 |
+
], betas=(.9, .95), eps=1e-8, weight_decay=.01, foreach=False)
|
| 191 |
+
wrapper = DDP(Step(model, config, device).eval(), device_ids=[device_index],
|
| 192 |
+
broadcast_buffers=False, find_unused_parameters=True,
|
| 193 |
+
gradient_as_bucket_view=True)
|
| 194 |
+
_, _, Processor = core.moss_small.export_classes()
|
| 195 |
+
processor = Processor(tokenizer=AutoTokenizer.from_pretrained(
|
| 196 |
+
core.moss_small.SFT3, trust_remote_code=True, local_files_only=True),
|
| 197 |
+
audio_tokenizer=None, model_config=config)
|
| 198 |
+
packer = core.ScorePacker(processor, config, schema)
|
| 199 |
+
reader = Reader()
|
| 200 |
+
run = ROOT / ('smoke_4gpu' if smoke else 'training')
|
| 201 |
+
run.mkdir(parents=True, exist_ok=True)
|
| 202 |
+
contract = {
|
| 203 |
+
'kind': 'hvs-long-ref-full-ft-v1', 'world': world, 'seed': SEED,
|
| 204 |
+
'base_checkpoint': str(BASE), 'manifest_sha256': sha(manifest_path),
|
| 205 |
+
'resolution_sha256': sha(RESOLUTION_MANIFEST), 'schema_sha256': sha(SCHEMA),
|
| 206 |
+
'source_prompt_sanitizer_version': SOURCE_PROMPT_SANITIZER_VERSION,
|
| 207 |
+
'script_renderer_version': RENDERER_VERSION, 'sampler_version': SAMPLER_VERSION,
|
| 208 |
+
'lr_backbone_peak': PEAK_BACKBONE, 'lr_talker_peak': PEAK_TALKER,
|
| 209 |
+
'objective': 'same_audio_logprob_as_reference_lora',
|
| 210 |
+
'prompt_counts': prompt_counts, 'total_steps': total,
|
| 211 |
+
}
|
| 212 |
+
chosen = [None]
|
| 213 |
+
scheduler = NoOpScheduler()
|
| 214 |
+
if rank == 0:
|
| 215 |
+
try:
|
| 216 |
+
directory = latest_checkpoint(run / 'checkpoints', contract, world)
|
| 217 |
+
if directory is not None:
|
| 218 |
+
complete_checkpoint(directory, contract, world, content=True)
|
| 219 |
+
chosen[0] = {'path': str(directory) if directory else None}
|
| 220 |
+
except Exception as error:
|
| 221 |
+
chosen[0] = {'error': repr(error)}
|
| 222 |
+
dist.broadcast_object_list(chosen, src=0)
|
| 223 |
+
assert 'error' not in chosen[0], chosen[0]
|
| 224 |
+
completed = 0
|
| 225 |
+
if chosen[0]['path']:
|
| 226 |
+
completed = state_io.load(chosen[0]['path'], wrapper, optimizer, scheduler, contract)
|
| 227 |
+
# state_io expects an ordinary scheduler object. The schedule is instead
|
| 228 |
+
# a deterministic pure function of the committed global step, so give it
|
| 229 |
+
# a no-op state carrier to preserve its atomic checkpoint/restore path.
|
| 230 |
+
if rank == 0:
|
| 231 |
+
print(json.dumps({'kind': 'start', 'smoke': smoke, 'world': world,
|
| 232 |
+
'steps': total, 'resumed_step': completed,
|
| 233 |
+
'trainable_parameters': sum(p.numel() for p in model.parameters() if p.requires_grad),
|
| 234 |
+
'frozen_unused_parameters': sum(p.numel() for p in model.parameters() if not p.requires_grad),
|
| 235 |
+
'backbone_parameters': sum(p.numel() for p in groups['backbone']),
|
| 236 |
+
'talker_parameters': sum(p.numel() for p in groups['talker']),
|
| 237 |
+
'effective_batch': world * accum,
|
| 238 |
+
'prompt_counts': prompt_counts,
|
| 239 |
+
'reference_max_seconds': MAX_REF_FRAMES * .08,
|
| 240 |
+
'rowgroup_load_min': min(loads), 'rowgroup_load_max': max(loads)}), flush=True)
|
| 241 |
+
metrics = run / 'metrics.jsonl'
|
| 242 |
+
if completed == 0:
|
| 243 |
+
value, n = mean_validation(model, config, reader, packer, annotations,
|
| 244 |
+
validation, device, rank, world)
|
| 245 |
+
if rank == 0:
|
| 246 |
+
with metrics.open('a') as output:
|
| 247 |
+
output.write(json.dumps({'kind': 'validation', 'step': 0,
|
| 248 |
+
'nll_per_frame': value, 'rows': n}) + '\n')
|
| 249 |
+
last = completed
|
| 250 |
+
for step in range(completed + 1, total + 1):
|
| 251 |
+
last = step
|
| 252 |
+
optimizer.zero_grad(set_to_none=True)
|
| 253 |
+
local_count = max(0, min(accum, len(rows) - (step - 1) * accum))
|
| 254 |
+
global_count = torch.tensor([local_count], dtype=torch.int64, device=device)
|
| 255 |
+
dist.all_reduce(global_count)
|
| 256 |
+
assert global_count.item() > 0
|
| 257 |
+
local_modes = [0] * len(MODES)
|
| 258 |
+
ref_frames = long_refs = short_refs = 0
|
| 259 |
+
local_loss = 0.
|
| 260 |
+
for micro in range(accum):
|
| 261 |
+
position = (step - 1) * accum + micro
|
| 262 |
+
active = position < len(rows)
|
| 263 |
+
row = rows[int(order[position])] if active else rows[0]
|
| 264 |
+
batch, facts = packed(row, reader, packer, annotations)
|
| 265 |
+
if active:
|
| 266 |
+
local_modes[MODES.index(facts['prompt_mode'])] += 1
|
| 267 |
+
ref_frames += facts['reference_frames']
|
| 268 |
+
long_refs += int(facts['reference_frames'] > 37)
|
| 269 |
+
short_refs += int(facts['reference_frames'] < PREFERRED_MIN_REF_FRAMES)
|
| 270 |
+
context = wrapper.no_sync() if micro + 1 < accum else nullcontext()
|
| 271 |
+
with context:
|
| 272 |
+
raw_loss = wrapper(batch)
|
| 273 |
+
assert torch.isfinite(raw_loss)
|
| 274 |
+
(raw_loss * (world / int(global_count.item()) if active else 0.)).backward()
|
| 275 |
+
if active:
|
| 276 |
+
local_loss += float(raw_loss.detach())
|
| 277 |
+
grad_norm = torch.nn.utils.clip_grad_norm_(wrapper.parameters(), 1.)
|
| 278 |
+
assert torch.isfinite(grad_norm) and float(grad_norm) > 0
|
| 279 |
+
update_canary = None
|
| 280 |
+
if smoke and step == 1:
|
| 281 |
+
representatives = []
|
| 282 |
+
for group_name in ('backbone', 'talker'):
|
| 283 |
+
chosen_parameter = next((parameter for parameter in groups[group_name]
|
| 284 |
+
if parameter.grad is not None and
|
| 285 |
+
bool(parameter.grad.detach().abs().max() > 0)), None)
|
| 286 |
+
assert chosen_parameter is not None, f'No {group_name} gradient'
|
| 287 |
+
representatives.append((group_name, chosen_parameter,
|
| 288 |
+
chosen_parameter.detach().clone()))
|
| 289 |
+
update_canary = representatives
|
| 290 |
+
multiplier = lr_factor(step, total)
|
| 291 |
+
optimizer.param_groups[0]['lr'] = PEAK_BACKBONE * multiplier
|
| 292 |
+
optimizer.param_groups[1]['lr'] = PEAK_TALKER * multiplier
|
| 293 |
+
optimizer.step()
|
| 294 |
+
if update_canary is not None:
|
| 295 |
+
weight_changes = {name: float((parameter.detach() - before).abs().sum())
|
| 296 |
+
for name, parameter, before in update_canary}
|
| 297 |
+
assert weight_changes['backbone'] > 0 and weight_changes['talker'] > 0, weight_changes
|
| 298 |
+
if rank == 0:
|
| 299 |
+
print(json.dumps({'kind': 'full_ft_update_canary', 'step': step,
|
| 300 |
+
'backbone_l1_change': weight_changes['backbone'],
|
| 301 |
+
'talker_l1_change': weight_changes['talker']}), flush=True)
|
| 302 |
+
summary = torch.tensor([ref_frames, long_refs, short_refs, *local_modes],
|
| 303 |
+
dtype=torch.int64, device=device)
|
| 304 |
+
dist.all_reduce(summary)
|
| 305 |
+
if rank == 0 and (smoke or step % 50 == 0):
|
| 306 |
+
event = {'kind': 'train', 'step': step, 'total': total,
|
| 307 |
+
'loss_local': local_loss / max(1, local_count),
|
| 308 |
+
'examples_global': int(global_count.item()),
|
| 309 |
+
'grad_norm': float(grad_norm),
|
| 310 |
+
'lr_backbone': optimizer.param_groups[0]['lr'],
|
| 311 |
+
'lr_talker': optimizer.param_groups[1]['lr'],
|
| 312 |
+
'reference_mean_seconds': float(summary[0]) * .08 / int(global_count.item()),
|
| 313 |
+
'reference_longer_than_2_96s': int(summary[1]),
|
| 314 |
+
'reference_under_preferred_5s': int(summary[2]),
|
| 315 |
+
'prompt_mode_counts': {mode: int(summary[i+3]) for i, mode in enumerate(MODES)},
|
| 316 |
+
'elapsed_s': time.monotonic() - started,
|
| 317 |
+
'peak_allocated_gib': torch.cuda.max_memory_allocated(device) / 2**30}
|
| 318 |
+
with metrics.open('a') as output:
|
| 319 |
+
output.write(json.dumps(event) + '\n')
|
| 320 |
+
print(json.dumps(event), flush=True)
|
| 321 |
+
value = None
|
| 322 |
+
if smoke or step % 500 == 0 or step == total:
|
| 323 |
+
value, n = mean_validation(model, config, reader, packer, annotations,
|
| 324 |
+
validation, device, rank, world)
|
| 325 |
+
if rank == 0:
|
| 326 |
+
with metrics.open('a') as output:
|
| 327 |
+
output.write(json.dumps({'kind': 'validation', 'step': step,
|
| 328 |
+
'nll_per_frame': value, 'rows': n}) + '\n')
|
| 329 |
+
flag = torch.tensor([int(STOP or time.monotonic() - started > 10.5 * 3600)], device=device)
|
| 330 |
+
dist.all_reduce(flag, op=dist.ReduceOp.MAX)
|
| 331 |
+
if (smoke and step == total) or (not smoke and step % SAVE_EVERY == 0) or step == total or flag.item():
|
| 332 |
+
directory = state_io.save(run / 'checkpoints', wrapper, optimizer,
|
| 333 |
+
scheduler, step, contract,
|
| 334 |
+
export=(step == total or (not smoke and step % 5000 == 0)))
|
| 335 |
+
if rank == 0:
|
| 336 |
+
print(json.dumps({'kind': 'checkpoint', 'step': step,
|
| 337 |
+
'directory': str(directory), 'validation_nll_per_frame': value}), flush=True)
|
| 338 |
+
if flag.item():
|
| 339 |
+
break
|
| 340 |
+
if rank == 0 and last == total:
|
| 341 |
+
atomic_json(run / 'DONE.json', {'status': 'COMPLETE', 'step': total,
|
| 342 |
+
'manifest_sha256': contract['manifest_sha256'],
|
| 343 |
+
'resolution_sha256': contract['resolution_sha256'],
|
| 344 |
+
'base_checkpoint': str(BASE), 'time': time.time()})
|
| 345 |
+
dist.barrier()
|
| 346 |
+
dist.destroy_process_group()
|
| 347 |
+
|
| 348 |
+
|
| 349 |
+
class NoOpScheduler:
|
| 350 |
+
def state_dict(self):
|
| 351 |
+
return {}
|
| 352 |
+
|
| 353 |
+
def load_state_dict(self, state):
|
| 354 |
+
assert state == {}
|
| 355 |
+
|
| 356 |
+
|
| 357 |
+
if __name__ == '__main__':
|
| 358 |
+
main()
|
code/ref15_fullft/train_full.sbatch
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
#SBATCH --job-name=hvs-ref15-s3-fullft
|
| 3 |
+
#SBATCH --account=reformo
|
| 4 |
+
#SBATCH --partition=booster
|
| 5 |
+
#SBATCH --nodes=8
|
| 6 |
+
#SBATCH --ntasks-per-node=4
|
| 7 |
+
#SBATCH --gres=gpu:4
|
| 8 |
+
#SBATCH --cpus-per-task=8
|
| 9 |
+
#SBATCH --time=11:30:00
|
| 10 |
+
#SBATCH --output=/e/scratch/reformo/schuhmann1_moss/out/hvs_ref_full_ft_v1/logs/train_%j.out
|
| 11 |
+
#SBATCH --error=/e/scratch/reformo/schuhmann1_moss/out/hvs_ref_full_ft_v1/logs/train_%j.err
|
| 12 |
+
set -euo pipefail
|
| 13 |
+
SC=/e/scratch/reformo/schuhmann1_moss
|
| 14 |
+
source "$SC/code/fastgen.sh"
|
| 15 |
+
export HF_HOME="$SC/cache/hf" XDG_CACHE_HOME="$SC/cache/xdg" TORCH_HOME="$SC/cache/torch" TRITON_CACHE_DIR="$SC/cache/triton"
|
| 16 |
+
mkdir -p "$HF_HOME" "$XDG_CACHE_HOME" "$TORCH_HOME" "$TRITON_CACHE_DIR"
|
| 17 |
+
export OMP_NUM_THREADS=4 TOKENIZERS_PARALLELISM=false HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
|
| 18 |
+
export PYTHONPATH=/e/home/jusers/schuhmann1/jupiter/hvs_ref_only_r128_v1:${PYTHONPATH:-}
|
| 19 |
+
export MASTER_ADDR
|
| 20 |
+
MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
|
| 21 |
+
export MASTER_PORT=$((20000 + SLURM_JOB_ID % 20000))
|
| 22 |
+
cd "$SC/out/hvs_ref_full_ft_v1"
|
| 23 |
+
python -u /e/home/jusers/schuhmann1/jupiter/hvs_ref_only_r128_v1/qa_legacy_pause_resume.py
|
| 24 |
+
python -u /e/home/jusers/schuhmann1/jupiter/hvs_ref_full_ft_v1/preflight.py
|
| 25 |
+
if [ -f "$SC/out/hvs_ref_full_ft_v1/training/DONE.json" ]; then
|
| 26 |
+
echo "Full S3 long-reference fine-tune already complete; no redundant segment."
|
| 27 |
+
exit 0
|
| 28 |
+
fi
|
| 29 |
+
srun --nodes=8 --ntasks=32 --ntasks-per-node=4 --gpus-per-task=1 --kill-on-bad-exit=1 --cpu-bind=none \
|
| 30 |
+
python -u /e/home/jusers/schuhmann1/jupiter/hvs_ref_full_ft_v1/train.py
|
| 31 |
+
|
docs/S3_REF15_FULL_FT.md
ADDED
|
@@ -0,0 +1,184 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# S3 Ref15 Full FT: data selection, training, and the reference-length fix
|
| 2 |
+
|
| 3 |
+
This document describes the model-only checkpoint in
|
| 4 |
+
[`checkpoints/S3-Ref15-FullFT/model_bf16.pt`](../checkpoints/S3-Ref15-FullFT/model_bf16.pt).
|
| 5 |
+
It is the **26,774th and final update** of a full-parameter, reference-conditioned
|
| 6 |
+
continuation of Humaneness Voice Small S3. It starts from the **original S3
|
| 7 |
+
step-32,337 checkpoint**, not from the subsequent Rank-128 Ref15 LoRA, a DPO
|
| 8 |
+
adapter, S10, or a randomly initialized language model. The original S3
|
| 9 |
+
checkpoint remains available at [`checkpoints/S3/`](../checkpoints/S3/).
|
| 10 |
+
The user selected this Full FT checkpoint as the new repository default after
|
| 11 |
+
listening and finding it acoustically preferable. That is a listening-based
|
| 12 |
+
choice, **not a claim that every automatic metric improved**.
|
| 13 |
+
|
| 14 |
+
## What changed from the original S3 model?
|
| 15 |
+
|
| 16 |
+
The original S1→S10 curriculum cropped a reference to at most **37 MOSS-v2
|
| 17 |
+
80-ms frames = 2.96 seconds** in the training loader. An older public
|
| 18 |
+
inference example imposed the same ceiling. The codec and network did not
|
| 19 |
+
require that limit: the stored reference codec sequence was longer. This
|
| 20 |
+
continuation changed the data locator's crop to the first **at most 187 frames
|
| 21 |
+
= 14.96 seconds**, and the public inference example now accepts 15 seconds by
|
| 22 |
+
default. No codec, token count per frame, Qwen backbone shape, Talker shape,
|
| 23 |
+
or vocoder architecture changed. The reference is passed through the same
|
| 24 |
+
dedicated reference-audio channel as before, not pasted into prompt text.
|
| 25 |
+
|
| 26 |
+
The retrain also deliberately changes what text accompanies the reference.
|
| 27 |
+
Every accepted example uses a distinct reference clip. Most prompts omit the
|
| 28 |
+
global voice-description caption so the reference has less competition as a
|
| 29 |
+
source of speaker identity. Some structured and caption prompts remain to
|
| 30 |
+
preserve the model's broader instruction vocabulary. These changes were made
|
| 31 |
+
**together**; this run alone cannot isolate whether longer references, more
|
| 32 |
+
frequent reference conditioning, prompt simplification, full-model updates,
|
| 33 |
+
or extra training caused a perceived difference.
|
| 34 |
+
|
| 35 |
+
## Which recordings were selected?
|
| 36 |
+
|
| 37 |
+
The selection starts with the physical **tier7 / logical S3** data in the
|
| 38 |
+
[TTS scaling ladder](https://huggingface.co/datasets/laion/tts-scaling-ladder-de-en).
|
| 39 |
+
The logical training manifest contains repeated presentations of recordings.
|
| 40 |
+
Its family-0 base indices identify **7,216,621 distinct target recordings**.
|
| 41 |
+
These were joined by UID to the already validated reference assets, yielding
|
| 42 |
+
**3,707,948 targets with an available separate reference**, or **8,832.58
|
| 43 |
+
hours of unique target audio before text-quality exclusions**. The target is
|
| 44 |
+
not reused as its own reference. The worklist is 128 packed Parquet shards;
|
| 45 |
+
no millions of individual audio files were created, and no audio was
|
| 46 |
+
re-encoded to build it.
|
| 47 |
+
|
| 48 |
+
The reference policy was intentionally inclusive: prefer a reference of at
|
| 49 |
+
least about five seconds where available, but retain the existing shorter
|
| 50 |
+
frozen pairs to cover all eligible S3 targets. Among the 3,707,948 initial
|
| 51 |
+
reference-eligible targets, **513,413 (13.85%)** had a presented reference
|
| 52 |
+
under five seconds. The complete stored reference is cropped from its start
|
| 53 |
+
to at most 187 frames. The training system checked that target and reference
|
| 54 |
+
UIDs differ, but this is **not** a manual same-speaker audit of every pair.
|
| 55 |
+
|
| 56 |
+
The initial reference-eligible source composition was 1,394,379
|
| 57 |
+
`vprof_original`, 1,098,944 `vc_sidon`, 696,175 `vprof_repair`, 491,992
|
| 58 |
+
`emolia`, 18,620 `kartoffelphon`, and 7,838 `mls` targets. These are **row
|
| 59 |
+
counts before the text-resolution exclusions**, not source hours or an
|
| 60 |
+
assertion that all sources have the same provenance. The original worklist
|
| 61 |
+
manifest also records source hours and reference-duration bins.
|
| 62 |
+
|
| 63 |
+
### Target-text integrity gate
|
| 64 |
+
|
| 65 |
+
A prior canary caught clips whose stored transcript described more speech
|
| 66 |
+
than the actual short target contained. A complete audit and ASR-supported,
|
| 67 |
+
conservative resolution sidecar was applied **before** this Full FT run.
|
| 68 |
+
Across the reference-eligible worklist, it examined 896,551 suspect rows:
|
| 69 |
+
440,195 were relabelled where supported, 20,081 were verified to keep, and
|
| 70 |
+
436,275 were excluded. That leaves **3,271,673 accepted targets**; the
|
| 71 |
+
exclusions account for **183.251 target-audio hours**, so accepted audio is
|
| 72 |
+
approximately **8,649.33 unique target hours including validation**. The
|
| 73 |
+
exact train/validation UID split is 3,255,217 / 16,456 after resolution.
|
| 74 |
+
The ASR decisions and sample QA passed the run's checks, but they are not a
|
| 75 |
+
claim of human-verifying every transcript. The original 0.5% split is
|
| 76 |
+
target-UID based, **not speaker-disjoint**.
|
| 77 |
+
|
| 78 |
+
### Exact training prompt mixture
|
| 79 |
+
|
| 80 |
+
Each accepted training UID appears **once**, assigned deterministically by
|
| 81 |
+
UID hash. Every mode includes the separate reference-audio input. The exact
|
| 82 |
+
effective counts recorded in the final checkpoint are:
|
| 83 |
+
|
| 84 |
+
| Mode | Training UIDs | What the instruction contains |
|
| 85 |
+
| --- | ---: | --- |
|
| 86 |
+
| Bare SCRIPT | 1,932,540 | Spoken text with sentence-local parenthesized performance, emotion, pause, duration and burst cues where available; no `GENERAL:` block. Stable identity descriptions are filtered out of local cues. |
|
| 87 |
+
| `GENERAL:` + `SCRIPT:` | 570,663 | Historical structured prompt, retained for compatibility. |
|
| 88 |
+
| Timbre Whisper caption + transcript | 318,972 | `CAPTION:` prose and exact `TRANSCRIPT:`. |
|
| 89 |
+
| Body/BUD-E Whisper 1.1 caption + transcript | 269,892 | Same outer surface; known duplicated first-word glitches are repaired, and suspect no-speech captions fall back to Timbre prose. |
|
| 90 |
+
| Transcript only | 163,150 | `TRANSCRIPT:` and reference audio, with no caption. |
|
| 91 |
+
|
| 92 |
+
These sum to **3,255,217** training UIDs, approximately 59.4% / 17.5% /
|
| 93 |
+
9.8% / 8.3% / 5.0% after the quality exclusions. The design target before
|
| 94 |
+
exclusions was roughly 55% / 20% / 10% / 10% / 5%. Caption modes were drawn
|
| 95 |
+
only for UIDs that actually had the corresponding annotation; captions were
|
| 96 |
+
not invented for unannotated rows. The script renderer retains valid
|
| 97 |
+
sentence timing cues, pauses and vocal bursts. The prompt surfaces do not
|
| 98 |
+
imply perfect duration or burst-onset control.
|
| 99 |
+
|
| 100 |
+
## How it was trained
|
| 101 |
+
|
| 102 |
+
The base weight file was S3's `step-00032337/model_bf16.pt`. The continuation
|
| 103 |
+
updated all **active** Qwen3 semantic-backbone, bridge and 12-codebook MOSS
|
| 104 |
+
Talker parameters with the unchanged teacher-forced audio-token log-probability
|
| 105 |
+
objective. The inherited but unused score conditioner stayed frozen. There
|
| 106 |
+
was **no LoRA** in this run. Model master parameters and AdamW moments were
|
| 107 |
+
FP32, while forward/loss used BF16 autocast; this avoids tiny updates being
|
| 108 |
+
rounded away in BF16 master weights.
|
| 109 |
+
|
| 110 |
+
Training ran on **eight JUPITER Booster nodes / 32 GH200 GPUs**, global batch
|
| 111 |
+
**128**, one deterministic pass over the accepted train UIDs. The semantic
|
| 112 |
+
peak LR was **8×10⁻⁶**; the Talker/bridge peak was **1.6×10⁻⁵**. Both had
|
| 113 |
+
5% warm-up followed by a single cosine decay to 10% of peak, **without a
|
| 114 |
+
restart**. AdamW used β=(0.9, 0.95), ε=10⁻⁸, weight decay 0.01 and global
|
| 115 |
+
gradient-norm clipping at 1.0. Unequal source-shard sizes were balanced by
|
| 116 |
+
assigning packed 4,096-row Parquet groups across ranks; the final dummy tail
|
| 117 |
+
has zero loss weight, so it does not add false training examples. Runs could
|
| 118 |
+
resume exact model, optimizer and RNG states from atomic checkpoints.
|
| 119 |
+
|
| 120 |
+
The final state is **step 26,774**. Its exported BF16 model is
|
| 121 |
+
**1,431,076,525 bytes**, SHA-256
|
| 122 |
+
`e1b2c4613545d1db21c6f60a5ab2c27b5320aa05c177717ae268b1ff870d406e`.
|
| 123 |
+
The local optimizer-bearing checkpoint is about 8 GB and is **not** part of
|
| 124 |
+
this public model-only release. The run's immutable worklist and text-repair
|
| 125 |
+
SHA-256 values are recorded in the adjacent `metadata.json`; those JUPITER
|
| 126 |
+
Scratch paths identify historical provenance, not portable data locations.
|
| 127 |
+
The scalar [training log](../stats/ref15_fullft_training_metrics.jsonl) and
|
| 128 |
+
[paired 50-case result JSON](../stats/ref15_fullft_listening50.json) are
|
| 129 |
+
included for audit. The historical source entrypoints are under
|
| 130 |
+
[`code/ref15_fullft/`](../code/ref15_fullft/); they depend on the original
|
| 131 |
+
packed JUPITER manifests and other release code, so they are provenance
|
| 132 |
+
rather than a one-command portable training package.
|
| 133 |
+
|
| 134 |
+
## What the available evaluation says
|
| 135 |
+
|
| 136 |
+
The checkpoint's bounded internal validation diagnostic used 511 rows and
|
| 137 |
+
reported 62.058 negative log-likelihood units per frame before continuation
|
| 138 |
+
and 61.997 at the final step. It was **not** the lowest value during the run
|
| 139 |
+
and is not a perceptual score. A separate **50-case paired listening grid**
|
| 140 |
+
used identical target prompts, 5.20–14.96-second references and random seeds
|
| 141 |
+
for original S3, Full FT step 25,000 and Full FT step 26,774. Those 50 cases
|
| 142 |
+
were UID-held-out but **not guaranteed speaker-disjoint**.
|
| 143 |
+
|
| 144 |
+
| Model | Mean Orange cosine to reference ↑ | Mean Orange cosine to target ↑ | Corpus Parakeet WER ↓ |
|
| 145 |
+
| --- | ---: | ---: | ---: |
|
| 146 |
+
| Original S3 step 32,337 | **0.745** | 0.658 | 0.088 |
|
| 147 |
+
| Full FT step 25,000 | 0.713 | **0.745** | **0.083** |
|
| 148 |
+
| Full FT step 26,774 | 0.707 | 0.734 | 0.096 |
|
| 149 |
+
|
| 150 |
+
The user found step 26,774 acoustically preferable in listening. The table
|
| 151 |
+
does **not** demonstrate a general voice-cloning or WER improvement for that
|
| 152 |
+
step: Orange similarity to the provided reference decreased, while similarity
|
| 153 |
+
to the intended target increased, and final WER is slightly worse than the
|
| 154 |
+
original S3 in this small test. The reference/target discrepancy deserves
|
| 155 |
+
separate examination. Automated scores cannot settle the subjective
|
| 156 |
+
naturalness question; a blinded, speaker-disjoint evaluation is still needed.
|
| 157 |
+
This repository's original S1–S10 acting benchmark was a different protocol
|
| 158 |
+
and must not be numerically merged with the 50-case Ref15 test.
|
| 159 |
+
|
| 160 |
+
## Loading and limitations
|
| 161 |
+
|
| 162 |
+
From this repository's root, the updated example defaults to this checkpoint:
|
| 163 |
+
|
| 164 |
+
```bash
|
| 165 |
+
python code/infer.py --reference-wav reference.wav --language en --frames 75 \
|
| 166 |
+
--prompt $'CAPTION: Warm, softly amused, conversational.\nTRANSCRIPT: "We made it home."' \
|
| 167 |
+
--text 'We made it home.' --output ref15_example.wav
|
| 168 |
+
```
|
| 169 |
+
|
| 170 |
+
`--stage default` and `--stage S3-Ref15-FullFT` resolve to the same file;
|
| 171 |
+
`--stage S3` explicitly loads the unchanged original S3. Provide a **clean,
|
| 172 |
+
distinct** reference, ideally 5–15 seconds. Shorter clips are accepted with
|
| 173 |
+
a warning; references longer than the 14.96-second training ceiling are out
|
| 174 |
+
of distribution. Because **every** accepted training example for this
|
| 175 |
+
continuation carried a reference, reference-free generation is not validated
|
| 176 |
+
as a new default behavior; compare the original S3 for that use. The
|
| 177 |
+
separate [MOSS Audio Tokenizer v2](https://huggingface.co/OpenMOSS-Team/MOSS-Audio-Tokenizer-v2)
|
| 178 |
+
is still required to encode references and decode the 12 acoustic codebooks
|
| 179 |
+
per 80-ms frame. The new weight is not a standard one-call
|
| 180 |
+
`AutoModel.from_pretrained()` checkpoint.
|
| 181 |
+
|
| 182 |
+
The original ladder stages, their selected validation checkpoints, and
|
| 183 |
+
their prior evaluations remain in the repository. Nothing here retroactively
|
| 184 |
+
changes the reference exposure or 2.96-second training crop of S1–S10.
|
stats/ref15_fullft_listening50.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
stats/ref15_fullft_training_metrics.jsonl
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|