Instruction:
## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  3) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  4) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


<permissions instructions>Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `workspace-write`: The sandbox permits reading files, and editing files in `cwd` and `writable_roots`. Editing files in other directories requires approval. Network access is restricted. Approvals are your mechanism to get user consent to run shell commands without the sandbox. `approval_policy` is `on-request`: Commands will be run in the sandbox by default, and you can specify in your tool call if you want to escalate a command to run without sandboxing. If the completing the task requires escalated permissions, Do not let these settings or the sandbox deter you from attempting to accomplish the user's task.

Here are scenarios where you'll need to request approval:
- You need to run a command that writes to a directory that requires it (e.g. running tests that write to /var)
- You need to run a GUI app (e.g., open/xdg-open/osascript) to open browsers or files.
- You are running sandboxed and need to run a command that requires network access (e.g. installing packages)
- If you run a command that is important to solving the user's query, but it fails because of sandboxing, rerun the command with approval. ALWAYS proceed to use the `sandbox_permissions` and `justification` parameters - do not message the user before requesting approval for the command.
- You are about to take a potentially destructive action such as an `rm` or `git reset` that the user did not explicitly ask for.

When requesting approval to execute a command that will require escalated privileges:
  - Provide the `sandbox_permissions` parameter with the value `"require_escalated"`
  - Include a short, 1 sentence explanation for why you need escalated permissions in the justification parameter The writable roots are `/home/ubuntu/maya3_prepare`, `/tmp`.</permissions instructions>


# AGENTS.md instructions for /home/ubuntu/maya3_prepare

<INSTRUCTIONS>
## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  3) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  4) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.
</INSTRUCTIONS>


<environment_context>
  <cwd>/home/ubuntu/maya3_prepare</cwd>
  <shell>bash</shell>
</environment_context>


<permissions instructions>Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `danger-full-access`: No filesystem sandboxing - all commands are permitted. Network access is enabled. Approvals are your mechanism to get user consent to run shell commands without the sandbox. `approval_policy` is `never`: This is a non-interactive mode where you may NEVER ask the user for approval to run commands. Instead, you must always persist and work around constraints to solve the task for the user. You MUST do your utmost best to finish the task and validate your work before yielding. If this mode is paired with `danger-full-access`, take advantage of it to deliver the best outcome for the user. Further, in this mode, your default testing philosophy is overridden: Even if you don't see local patterns for testing, you may add tests and scripts to validate your work. Just remove them before yielding.</permissions instructions>


# Context from my IDE setup:

## Open tabs:
- .env: .env
- benchmark_scaling.sh: scripts/benchmark_scaling.sh
- objective.md: objective.md
- CURRENT_STATUS.md: CURRENT_STATUS.md

## My request for Codex:
We are on a plan of data preparation for codec training. plan is at [strategy.md](strategy.md) and objective - [objective.md](objective.md)

i want you to go through entire plan, pipeline implementation, and see if everything is implemented and is working as intended. test with a video download from R2 directly. and test the segments and read everything right from audio file by validating everything from audio file as well and see if everything tallies with what we are trying to do. and then see if we are missing anything as per plan. i also need you to go through the codec training code in x-codec-2.0 aswell and then ensure our data as .tar will be valid with it completely, or we can fix those scripts. in short, i need to be completely sure and reliable that my data is completely alright and there is no issue or failures, or regressions or problems with our pipeline here. and it's completely optimal about what we are trying to do here. validate with 1 videoID from R2, you can pick the videos from Supabase's table, run a video on this pipeline, and then if pipeline is validated, we can plan for optimising and parallelisation and scaling horizontally by docker and other methods later. first finish your complete effort in validating the pipeline by checking everything manually and ensuring each and every segment is read, validated, and i can rely the pipeline is robust and is working perfect. we can talk about scaling horizontally later. original files are in `test` bucket and cleaned segmented flac .tar files segmented and metadata packed in videoID.tar are in `1-cleaned-data` bucket and note that there might be few prefixes. double check with supabase's r2_key column in Supabase. all credentials are in .env


Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  3) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  4) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


**Listing repo contents**


Output:
Exit code: 0
Wall time: 0 seconds
Output:
docs
objective.md
scripts
strategy.md
Summary.md
tts
X-Codec-2.0


Summary: auto


**Checking existing project docs**

I'm opening strategy.md and objective.md to review current plans and status before proceeding.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
# Codec Training Strategy (Indic, X-Codec-2.0 + Inworld SR)

## Current State

- **Source**: ~507k YouTube podcast videos across 12 Indic languages in Cloudflare R2 (`1-cleaned-data` bucket)
- **Format**: `videoID.tar` files containing `metadata.json` + `segments/*.flac` (original quality from YouTube WebM, compressed as FLAC)
- **Usable hours**: ~149k hours (after VAD/speaker segmentation)
- **Native sample rate**: 99.8% containers at 48kHz, but actual bandwidth varies (see bandwidth distribution below)
- **Note**: Segments are already at original quality; resampling to 16/24/32/44k happens during sharding

### Codebase References
- **X-Codec-2.0**: Full 16kHz codec (encoder + decoder + VQ) with W2V-BERT semantic features
  - Random 6s crops (`min_audio_length = 96000` @ 16kHz)
  - W2V-BERT frozen during training
  - Codebook: 16384 entries, dim=16, vq_dim=2048
- **Inworld TTS**: Decoder-only training from precomputed codes
  - 48kHz SR with `upsample_factors=[3,2]`, `hop_length=160`, `token_rate=50`
  - RMS loss for volume consistency (`lambda_rms=1.0`)
  - Two-stage 48kHz training: ≥32kHz → ≥44.1kHz

### Actual Bandwidth Distribution (from smoke test analysis)

Analysis of effective Nyquist frequency (actual audio bandwidth):
- **24-32k bandwidth**: ~62% of segments (effective_nyquist 12-16 kHz)
- **32-44k bandwidth**: ~30% of segments (effective_nyquist 16-22 kHz)
- **16-24k bandwidth**: ~7% of segments (effective_nyquist 8-12 kHz)
- **<16k bandwidth**: <1% of segments (effective_nyquist <8 kHz)
- **>=44k bandwidth (true wideband)**: Near zero (upsampled or low-passed)

**Key insight**: Container SR (48kHz) ≠ actual bandwidth. Most audio is bandwidth-limited to 12-18 kHz.

### Known Issues to Address
1. ~~**Language missing** from `metadata.json`~~ → FIXED: Fetched from Supabase
2. **Segment duration skew**: ~58% < 6s, ~27% < 3s → HANDLED: Random 6s crops during training
3. ~~**True sample rate untrusted**~~ → FIXED: Spectral analysis computes effective bandwidth
4. ~~**Missing quality metrics**~~ → FIXED: Loudness, clipping, SNR computed during sharding
5. ~~**Native SR buckets missing**~~ → FIXED: Multi-resolution buckets (16/24/32/44k) implemented

---

## 1) Data Validation Pipeline

### 1.1 Fetch Language Metadata from Supabase

Query Supabase using credentials in `.env` (scripts in `scripts/` for reference):

```
videoID → language
```

Store as `language_lookup.json` for fast local access during processing.

### 1.2 Compute Per-Segment Quality Metrics

For each segment in each `videoID.tar`, compute and store:

**Audio Quality Metrics (fast, essential):**
| Metric | Description | Purpose |
|--------|-------------|---------|
| `duration_sec` | Segment length | Already in metadata |
| `loudness_lufs` | Integrated loudness | RMS loss validation |
| `loudness_rms_db` | RMS in dB | Volume normalization |
| `clipping_fraction` | % samples \|x\| > 0.99 | Filter damaged audio |
| `dc_offset` | Mean of waveform | Catch broken pipelines |
| `snr_proxy` | Speech vs noise floor | Quality weighting |

**True Sample Rate Detection (critical for Phase C):**
| Metric | Description | Purpose |
|--------|-------------|---------|
| `spectral_rolloff_hz` | Freq below 85% energy | Detect bandwidth |
| `effective_nyquist_hz` | Detected cutoff frequency | True sample rate |
| `high_freq_energy_ratio` | Energy >8kHz / total | Wideband detection |
| `native_sr_bucket` | `<16k`, `16-24k`, `24-32k`, `32-44k`, `>=44k` | Bandwidth classification |
| `is_truly_wideband` | Boolean: effective_nyquist >= 12kHz | Higher-SR eligibility |

**Preserve Existing Metrics:**
- `music_stats` (already computed via PANN)
- `speaker` label
- `start`, `end`, `original_start`, `original_end`, `duration`

### 1.3 Native SR Classification + Resample Outputs

For each usable segment:
1. Load the segment FLAC from the tar (already at original quality from YouTube WebM).
2. Run spectral detection (rolloff, effective_nyquist, HF ratio).
3. Assign buckets and resample targets based on **actual bandwidth**, not container SR:
   - **16k bucket**: all usable segments (universal). Long segments (>=12s) are split into 6–10s windows.
   - **24k bucket**: `effective_nyquist >= 8000` Hz (narrowband to low-wideband)
   - **32k bucket**: `effective_nyquist >= 12000` Hz (mid-wideband)
   - **44k bucket**: `effective_nyquist >= 16000` Hz (high-wideband, true 44.1k sources)
4. Record `resample_target_sr` for each emitted sample.
5. Wideband flag: `is_truly_wideband = effective_nyquist >= 12000` (relaxed from previous 14k + HF ratio gate)

### 1.3 Composite Quality Score

Derive `quality_score` (0-1) for weighted sampling (store, do not filter):

```python
quality_score = (
    0.3 * (1 - clipping_fraction) +
    0.2 * normalize(snr_proxy, min=10, max=40) +
    0.2 * (1 if not music_stats.has_music else 0.5) +
    0.2 * normalize(loudness_lufs, min=-40, max=-14) +
    0.1 * (1 if abs(dc_offset) < 0.01 else 0.5)
)
```

---

## 2) Sharding Strategy

### 2.1 Segment Handling Decision: No Stitching

**Problem**: ~58% of segments are <6s, but X-Codec requires 6s crops.

**Decision**: Do NOT stitch segments. Reasons:
1. Speaker diarization is not 100% reliable - stitching risks mixing speakers
2. W2V-BERT semantic encoder expects consistent speaker within sample
3. Padding is recoverable (efficiency loss), data corruption is not

**Implemented approach**:
- **Drop segments <3s** (too short, mostly padding anyway - ~27% of segments)
- **Duration-weighted sampling** to favor longer segments
- **Chunk long segments for 16k**: For segments >=12s, split into non-overlapping 6–10s windows to boost sample count (no stitching across speakers).
- **Silence-aware chunking**: Only cut when a low-energy boundary exists near the target window; otherwise keep the full segment.
- **No pre-cropping needed for training**: Random 6s crops still happen during training (X-Codec data_module.py:104)
- **VAD regions**: Store full duration `[[0, duration_sec]]` as placeholder; actual cropping is random at train time

### 2.2 Output Structure

```
processed_data/
├── manifests/
│   ├── telugu_train.jsonl
│   ├── telugu_val.jsonl
│   └── ... (per language)
├── telugu_16/           # 16kHz - universal bucket
│   ├── shard-000000.tar
│   ├── shard-000000.jsonl
│   └── ...
├── telugu_24/           # 24kHz - effective_nyquist >= 8kHz
│   ├── shard-000000.tar
│   ├── shard-000000.jsonl
│   └── ...
├── telugu_32/           # 32kHz - effective_nyquist >= 12kHz
│   ├── shard-000000.tar
│   ├── shard-000000.jsonl
│   └── ...
├── telugu_44/           # 44.1kHz - effective_nyquist >= 16kHz
│   ├── shard-000000.tar
│   ├── shard-000000.jsonl
│   └── ...
└── stats/
    ├── telugu_stats.json
    └── ...
```

Note: Directory naming is `{language}_{sr_khz}` for easier R2 prefix management.

### 2.3 Shard Contents

Each shard (~5GB, multipart upload friendly) contains audio only, plus a single
JSONL sidecar file with per-sample metadata:

```
shard-000000.tar
shard-000000.jsonl
```

`shard-000000.jsonl` is one line per sample with:
```json
{
  "id": "videoID__SPEAKER_00_0001",
  "video_id": "---g_VL7ySo",
  "language": "bengali",
  "duration_sec": 4.11,
  "sample_rate": 16000,
  "resample_target_sr": 16000,
  "native_sr": 48000,
  "effective_nyquist_hz": 24000.0,
  "is_truly_wideband": true,
  "native_sr_bucket": ">=44k",
  "loudness_lufs": -23.5,
  "loudness_rms_db": -18.2,
  "clipping_fraction": 0.0,
  "snr_proxy": 28.5,
  "quality_score": 0.85,
  "music_prob": 0.045,
  "speaker": "SPEAKER_00",
  "vad_speech_regions": [[0.0, 4.11]],
  "duration_bucket": "3-6s",
  "sampling_weight": 0.3
}
```

Metadata storage:
- Keep per-shard JSONL locally for now.
- Also write per-language JSONL files for easy Supabase upsert later.

### 2.4 Filtering Rules (applied during sharding)


Output:
Exit code: 0
Wall time: 0 seconds
Output:
### 2.4 Filtering Rules (applied during sharding)

| Rule | Action | Rationale |
|------|--------|-----------|
| `duration_sec < 3.0` | **Drop** | Too short, mostly padding |
| `duration_sec > 60.0` | **Drop** | Too long, memory issues |
| `clipping_fraction > 0.1` | **Drop** | Severely clipped audio |
| `music_stats.has_music = true` | **Keep, lower weight** | May still have speech |
| `effective_nyquist < 8000` | **16k only** | Narrowband audio |
| `effective_nyquist < 12000` | **16k, 24k only** | Low-wideband audio |
| `effective_nyquist < 16000` | **Exclude 44k** | Not high-wideband |

### 2.5 Duration Bucket Weights

| Bucket | Duration | Weight | Rationale |
|--------|----------|--------|-----------|
| A | 3-6s | 0.3 | Padding overhead, keep for robustness |
| B | 6-12s | 1.0 | Optimal for 6s random crops |
| C | 12-24s | 0.8 | Good, multiple crops per segment |
| D | 24-60s | 0.5 | Memory intensive |

### 2.6 Train/Val Split

- Split by `video_id` (not segment) to avoid data leakage
- 99% train / 1% val per language
- Ensure val set covers all duration buckets
- Target: ~1000+ hours validation per major language

---

## 3) Training Strategy

### Phase A: 16kHz Fine-tune (Indic Adaptation)

**Goal**: Adapt semantics + acoustic reconstruction to Indic languages without destabilizing codebook.

**Data**: All languages from `shards_16k/`, weighted sampling by duration + quality

**Initialization**:
1. Load pretrained X-Codec-2.0 checkpoint (HKUST 16kHz)
2. Keep W2V-BERT semantic encoder frozen (default)
3. **First 5k steps**: Also freeze encoder + VQ codebook (stabilize)
4. **After 5k steps**: Unfreeze encoder, keep VQ learning rate at 0.1x

**Training Config**:
```yaml
batch_size: 8  # per GPU
min_audio_length: 96000  # 6s @ 16kHz
precision: 16-mixed
losses:
  - mel (lambda: 15.0)
  - adversarial (lambda: 1.0)
  - feature_matching (lambda: 1.0)
  - vq_commitment
  - semantic_reconstruction (lambda: 5.0)
  - rms (lambda: 1.0)  # ADD THIS from Inworld
```

**Curriculum (progressive scaling)**:
| Stage | Data | Steps | Notes |
|-------|------|-------|-------|
| A1 | 5k hrs, quality_score > 0.7 | 50k | Clean subset, stability check |
| A2 | 15k hrs, quality_score > 0.5 | 100k | Expand coverage |
| A3 | Full dataset, weighted sampling | 200k+ | All languages |

**Validation Metrics (per language)**:
- Mel L1 loss
- Multi-resolution STFT loss
- Speaker cosine similarity (WavLM, already in X-Codec)
- Codebook utilization (detect collapse: should use >90% of codes)
- SI-SDR (scale-invariant SDR)

### Phase B: Stabilize and Vectorize

**Goal**: Lock encoder, precompute codes for decoder-only training.

1. Take best Phase A checkpoint
2. Freeze encoder + VQ completely
3. Run short stabilization (10k steps) with frozen encoder
4. Use Inworld `data_vectorizer.py` to generate:
   - `train_codes.npy`, `train_codes_index.npy`, `train_samples.jsonl`
   - Per language, per resolution

**Vectorizer Modifications Needed (later, if/when used)**:
- Remove `filter_non_english` from `tts/tts/data/filtering.py`
- Add Indic language codes to `allowed_languages`

### Phase C1: 24kHz Decoder (Optional Low-bandwidth)

**Goal**: Handle low-bandwidth audio efficiently.

**Data**: `{lang}_24/` buckets (`effective_nyquist >= 8kHz`)

**Config** (token rate 50):
```python
sample_rate = 24000
hop_length = 240
upsample_factors = [2]
```

### Phase C2: 32kHz Decoder

**Goal**: Mid-range quality decoder for majority of data.

**Data**: `{lang}_32/` buckets (`effective_nyquist >= 12kHz`)

**Config** (token rate 50):
```python
sample_rate = 32000
hop_length = 320
upsample_factors = [2]
```

### Phase C3: 44.1kHz High-Fidelity Decoder

**Goal**: High-fidelity output for true wideband sources.

**Data**: `{lang}_44/` buckets (`effective_nyquist >= 16kHz`)

**Architecture**:
```python
sample_rate = 44100
hop_length = 147
upsample_factors = [3, 2]  # or [6] for single stage
# Verify: 44100 // 147 // 6 ≈ 50 tokens/sec
```

**Training Notes**:
- Encoder frozen (use codes from Phase B)
- Decoder backbone trainable, quantizer frozen
- Progressive training: start with 32k decoder, then fine-tune to 44k
- `audio_window_size` adjusted per sample rate

**Validation Metrics**:
- DNSMOS (perceptual quality)
- High-frequency energy preservation (effective_nyquist comparison)
- Spectral rolloff comparison (input vs output)
- Speaker similarity (WavLM embeddings)
- Mel spectrogram L1 distance

---

## 4) Implementation Plan

### Priority Order

| Step | Task | Output |
|------|------|--------|
| 1 | Fetch languages from Supabase | `language_lookup.json` |
| 2 | Compute quality metrics per segment | Updated metadata |
| 3 | Classify native SR + resample to 16/32/48 | Bucketed audio |
| 4 | Create sharding pipeline | WebDataset tars + JSONL |
| 4 | Generate train/val manifests | JSONL files |
| 5 | Upload to R2 | Organized buckets |

### Files to Create

| File | Purpose |
|------|---------|
| `scripts/fetch_languages.py` | Query Supabase for videoID → language |
| `scripts/compute_metrics.py` | Compute audio quality metrics per segment |
| `scripts/resample_and_bucket.py` | Slice from originals + emit 16/32/48 |
| `scripts/create_shards.py` | Pack segments into WebDataset tars + JSONL |
| `scripts/generate_manifests.py` | Create train/val JSONL files |

### Files to Modify (later, for training)

| File | Change |
|------|--------|
| `tts/tts/data/filtering.py` | Remove `filter_non_english` (if vectorizer is used) |
| `X-Codec-2.0/lightning_module.py` | Add RMS loss |

---

## 5) Validation Checkpoints

Before proceeding to training, verify:

- [x] All videoIDs have language assigned (from Supabase)
- [x] Quality metrics computed for all kept segments (loudness, clipping, SNR, spectral)
- [x] Shard sizes ~5GB and JSONL sidecars present
- [x] Train/val splits created per language
- [ ] Larger smoke test completed (50-100 videos per language)
- [ ] Bandwidth distribution validated across languages
- [x] Segments with `clipping_fraction > 0.1` dropped during filtering
- [ ] Upload to R2 verified (download and re-check integrity)
- [ ] At least 100 hours validation data per major language

---

## 6) Language Distribution (from Supabase)

| Language | Videos | Usable Hours | % of Total |
|----------|--------|--------------|------------|
| Telugu | 71,619 | 19,712 | 13.2% |
| Malayalam | 65,273 | 18,184 | 12.2% |
| English | 59,657 | 17,904 | 12.0% |
| Hindi | 58,824 | 17,710 | 11.9% |
| Punjabi | 55,761 | 18,065 | 12.1% |
| Tamil | 50,199 | 11,258 | 7.6% |
| Kannada | 45,689 | 13,578 | 9.1% |


**Reviewing pipeline implementation and scripts**


Output:
Exit code: 0
Wall time: 0 seconds
Output:
total 96
drwxr-xr-x 3 ubuntu ubuntu  4096 Jan 19 04:18 .
drwxrwxr-x 6 ubuntu ubuntu  4096 Jan 19 04:19 ..
-rw-rw-r-- 1 ubuntu ubuntu  7291 Jan 19 04:03 bulk_download.py
-rw-rw-r-- 1 ubuntu ubuntu  2127 Jan 19 04:03 fetch_languages.py
-rw-r--r-- 1 ubuntu ubuntu  6397 Jan 19 04:03 fill_missing_languages.py
-rw-r--r-- 1 ubuntu ubuntu  2169 Jan 19 04:03 normalize_languages.py
-rw-rw-r-- 1 ubuntu ubuntu 47849 Jan 19 04:03 prepare_shards.py
drwxrwxr-x 2 ubuntu ubuntu  4096 Jan 19 04:49 __pycache__
-rw-r--r-- 1 ubuntu ubuntu 11693 Jan 19 04:03 supabase_summary.py


**Scanning pipeline scripts**

I'm checking prepare_shards.py and related pipeline scripts to understand their role and see if they clarify the objective or workflow.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
Got it. Here’s a cluster-friendly sharding + training flow for an X-Codec2-style codec, plus the metadata you should precompute so training stays fast and you can do smart sampling/filters later.
1) How to shard the data for clusters
Goals
Sequential reads (no random tiny files)
Fast dataloader (workers stream chunks)
Deterministic sampling (repeatable epochs)
Easy filtering (by language, SNR, duration, speaker, etc.)
Recommended format (simple + scalable)
WebDataset shards: shard-000000.tar, shard-000001.tar, …
Each sample inside the tar:
sample_id.wav (or .flac)
sample_id.json (per-sample metadata) OR a global manifest.jsonl indexed by sample_id
Why WebDataset:
Streams well from object storage / NFS
Great for multi-node training
Avoids “millions of files” pain
Shard size rule of thumb
0.5–2 GB per tar
Aim for 2k–10k samples per shard
Keep shard durations roughly balanced (so workers don’t straggle)
Directory layout
data/
  shards/
    train/
      shard-000000.tar
      shard-000001.tar
      ...
    val/
      shard-000000.tar
      ...
  manifests/
    train.jsonl
    val.jsonl
  stats/
    speaker_embeds/
    asr_transcripts/
    ...
2) What to put in your manifest (metadata before training)
Create a train.jsonl where each line is one sample:
Key fields to include:
id (unique)
shard (which tar)
offset (optional if you do indexed tar reads; not required for pure streaming)
duration_sec
sr (sample rate)
language (hi/te/ta/kn/ml/bn/mr/gu/pa/ur/en etc.)
speaker_id (if known)
source (YT, audiobook, call, podcast)
text (if transcript exists)
split (train/val/test)
quality_flags (booleans like clipped, music, low_snr, overlap, etc.)
weights (sampling weight you can compute later)
3) Metrics to precompute (the ones that actually help)
A) Fast audio-quality metrics (super useful)
These let you filter garbage and do curriculum:
1. Duration
drop too short (e.g. < 0.8s) and too long (e.g. > 30s) early
2. Loudness (LUFS) or RMS
filter near-silent / extremely loud
3. Clipping %
count samples where |x| > 0.99
avoid clipped audio early
4. SNR estimate
even a rough SNR proxy helps for weighting
(speech energy vs noise floor percentiles)
5. Bandwidth / lowpass indicator
detect 8k telephone-like audio vs wideband
compute spectral roll-off or high-frequency energy ratio
6. DC offset
catches broken audio pipelines
These are cheap and help massively.
B) Speech-content/semantic metrics (for codec semantics)
These help you track “meaning preservation” later and build eval sets:
7. ASR transcript + confidence
• run a multilingual ASR (Whisper large-v3 / Indic ASR if you have)
• store:
transcript
avg token logprob / confidence
no-speech probability (Whisper)
• Use it to:
filter non-speech
build WER/CER evaluation sets
8. VAD speech ratio
percent of frames with speech
filter music/long silences
9. Language ID confidence
LID per clip (fasttext / whisper LID / langid)
helps avoid mislabeled language
10. Overlap / multi-speaker probability
diarization-lite: “single-speaker score”
optional but helpful for cloning-quality datasets
C) Speaker/style metrics (if you care about cloning)
11. Speaker embedding
• e.g., ECAPA-TDNN / pyannote embedding
• store vector or a hash + cluster id
• enables:
speaker-balanced sampling
speaker similarity eval later
D) Dataset management metrics
12. Duplicate / near-duplicate hash
audio fingerprinting (chromaprint) or embedding similarity
remove duplicates so model doesn’t overfit to repeated clips
4) Sharding workflow (end-to-end)
Step 0 — Normalize audio consistently
• Pick training SR: 16k or 24k (start here even if you’ll do 48k later)
• Convert all audio to:
mono
float32
target SR
peak normalize lightly (or loudness normalize if you want)
Step 1 — Segment
• Use VAD to split long recordings into clips:
target 3–12 seconds
keep some 15–25s clips for robustness (later phases)
Step 2 — Compute metadata (the metrics above)
Write a single manifest_raw.jsonl
Step 3 — Filter + label
• remove:
too silent
too clipped
too low speech ratio
• assign:
language
speaker cluster (optional)
• create train/val/test splits by speaker (important!)
Step 4 — Pack into shards
sort by duration (to reduce padding waste)
pack into tar shards of 0.5–2GB
write final train.jsonl with shard paths
Step 5 — Upload to storage
Object store (R2/S3) or shared filesystem
Keep shards immutable; write new versions for new filtering
5) Training flow for the codec (what you do on clusters)
Phase 0: smoke test
• train on 100–300h
• verify:
recon audio isn’t collapsing
losses decrease
RVQ usage isn’t dead (codebook collapse)
Phase 1: clean Indic core
train on 5k–15k hours clean
mix: 90–95% Indic + 5–10% English (optional)
Phase 2: scale + robustness
• train on full dataset with weighted sampling:
60–70% clean
20–30% mild noise / phone
5–10% hard cases
Always keep a fixed validation suite
by language
by duration bucket (short/med/long)
by noise bucket
by accent cluster (if you have it)
6) What metadata is most “handy later”
If I had to pick the top 8 that pay off the most:
1. duration
2. VAD speech ratio
3. clipping %
4. loudness (LUFS/RMS)
5. SNR proxy
6. language + LID confidence
7. ASR transcript + confidence
8. speaker embedding / speaker cluster id
These enable:
curriculum sampling
robust eval sets
debugging “why WER spiked”
filtering bad shards quickly
If you tell me where your shards will live (local NVMe on nodes vs S3/R2 vs NFS) and your target SR (16k/24k/48k), I’ll give you a specific recipe for:
shard size
dataloader settings (num_workers/prefetch/buffer/shuffle)
sampling weights per bucket
a clean “train/val suite” definition (WER/CER + DNSMOS + SI-SDR + speaker sim)
nikil
  8:20 PM
Perfect — random 6s crops + R2 (S3-compatible) + 16k now, 48k later is a very workable setup. Here’s how I’d structure it so training stays fast and your metadata actually helps.
1) Data format + shard strategy for R2
Use WebDataset tar shards (streaming-friendly)
r2://bucket/xcodec16/train/shard-000000.tar
r2://bucket/xcodec16/val/shard-000000.tar
Shard size (important for object storage):
512MB – 1GB per tar is a sweet spot on S3/R2
Too small → too many GETs
Too big → slow retries / hotspot objects
Inside each shard
For each sample uid store:
uid.flac (recommended) or uid.wav
uid.json (optional; can also use a global manifest)
Why FLAC?
Cuts bandwidth from R2 significantly
Decode cost is fine (CPU) vs network savings
2) Don’t pre-cut into 6s clips (keep variable-length originals)
Since your trainer randomly crops 6s:
• Store clips as natural segments (e.g., 2–30s after VAD)
• At load time:
if duration ≥ 6s → random crop
else → pad (as you’re doing)
But: too many <6s samples = lots of padding waste and worse training signal.
Rule of thumb
Try to make ≥80–90% of training samples ≥ 6s
Keep a small portion of shorter clips for robustness (but don’t let them dominate)


**Reviewing pipeline and prepare_shards.py**

I noted the objective.md has more instructions, so I'll review the pipeline and examine prepare_shards.py next to understand their roles and see how to test with R2.


sed -n '1,220p' scripts/prepare_shards.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
from __future__ import annotations

import argparse
import io
import json
import math
import os
import random
import tarfile
import threading
import time
import zlib
from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

import boto3
import numpy as np
import psycopg
import soundfile as sf
from boto3.s3.transfer import TransferConfig
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
from scipy import signal
from tqdm import tqdm

try:
    import pyloudnorm as pyln
except Exception:  # pragma: no cover - optional dependency
    pyln = None

ENV_PATH = Path(".env")
LANGUAGES = [
    "telugu",
    "malayalam",
    "english",
    "hindi",
    "punjabi",
    "tamil",
    "kannada",
    "gujarati",
    "bengali",
    "odia",
    "marathi",
    "assamese",
]

# Global timing stats (thread-safe)
_timing_lock = threading.Lock()
_timing_stats: dict[str, list[float]] = defaultdict(list)
_timing_enabled = False

@contextmanager
def timer(name: str):
    """Context manager for timing code blocks."""
    if not _timing_enabled:
        yield
        return
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        with _timing_lock:
            _timing_stats[name].append(elapsed)

def print_timing_stats():
    """Print aggregated timing statistics."""
    if not _timing_stats:
        return
    print("\n=== Timing Statistics ===")
    total_time = sum(sum(times) for times in _timing_stats.values())
    for name, times in sorted(_timing_stats.items()):
        count = len(times)
        total = sum(times)
        avg = total / count if count > 0 else 0
        percent = (total / total_time * 100) if total_time > 0 else 0
        print(f"{name:35s}: {count:6d} calls, {total:8.2f}s total, {avg*1000:8.2f}ms avg, {percent:5.1f}%")
    print(f"{'TOTAL':35s}:              {total_time:8.2f}s")
    print("=" * 80)

_window_cache: dict[int, np.ndarray] = {}
_freq_cache: dict[tuple[int, int], np.ndarray] = {}
_cache_lock = threading.Lock()
_lufs_local = threading.local()

def get_hann_window(length: int) -> np.ndarray:
    window = _window_cache.get(length)
    if window is not None:
        return window
    with _cache_lock:
        window = _window_cache.get(length)
        if window is None:
            window = np.hanning(length)
            _window_cache[length] = window
    return window

def get_rfftfreq(length: int, sr: int) -> np.ndarray:
    key = (length, sr)
    freqs = _freq_cache.get(key)
    if freqs is not None:
        return freqs
    with _cache_lock:
        freqs = _freq_cache.get(key)
        if freqs is None:
            freqs = np.fft.rfftfreq(length, d=1.0 / sr)
            _freq_cache[key] = freqs
    return freqs

def get_lufs_meter(sr: int) -> Optional[Any]:
    if pyln is None:
        return None
    meters = getattr(_lufs_local, "meters", None)
    if meters is None:
        meters = {}
        _lufs_local.meters = meters
    meter = meters.get(sr)
    if meter is None:
        meter = pyln.Meter(sr)
        meters[sr] = meter
    return meter

@dataclass(frozen=True)
class VideoRecord:
    youtube_id: str
    r2_tar_key: str

@dataclass(frozen=True)
class SegmentOutput:
    sample_rate: int
    audio_bytes: bytes
    metadata: dict[str, Any]

@dataclass(frozen=True)
class ChunkConfig:
    enabled: bool
    min_sec: float
    max_sec: float
    threshold_sec: float
    all_buckets: bool

@dataclass(frozen=True)
class SilenceConfig:
    enabled: bool
    frame_sec: float
    hop_sec: float
    search_sec: float
    percentile: float
    db_drop: float

def load_env(env_path: Path) -> dict[str, str]:
    env: dict[str, str] = {}
    for line in env_path.read_text().splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in stripped:
            continue
        if stripped.startswith("export "):
            stripped = stripped[len("export ") :]
        key, value = stripped.split("=", 1)
        env[key.strip()] = value.strip()
    return env

def make_r2_client(env: dict[str, str]):
    endpoint = env.get("R2_ENDPOINT_URL")
    access_key = env.get("R2_ACCESS_KEY_ID")
    secret_key = env.get("R2_SECRET_ACCESS_KEY")
    if not endpoint or not access_key or not secret_key:
        raise SystemExit("Missing R2 credentials in .env")
    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
    )

def resolve_r2_key(
    s3, bucket: str, r2_key: str, prefix: Optional[str]
) -> Optional[str]:
    if not r2_key:
        return None
    key = r2_key.lstrip("/")
    candidates = [key]

    if prefix:
        cleaned_prefix = prefix.strip("/")
        if not key.startswith(cleaned_prefix):
            candidates.append(f"{cleaned_prefix}/{key}")

    base = os.path.basename(key)
    if base != key:
        candidates.append(base)
        if prefix:
            candidates.append(f"{prefix.strip('/')}/{base}")

    for candidate in candidates:
        try:
            s3.head_object(Bucket=bucket, Key=candidate)
            return candidate
        except Exception:
            continue
    return None

def parse_source_url(source_url: Optional[str], default_bucket: Optional[str]) -> Optional[tuple[str, str]]:
    if not source_url or not default_bucket:
        return None


sed -n '220,440p' scripts/prepare_shards.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
        return None
    cleaned = source_url.lstrip("/")
    if cleaned.startswith(default_bucket + "/"):
        return default_bucket, cleaned[len(default_bucket) + 1 :]
    return default_bucket, cleaned

def fetch_video_records(
    db_url: str,
    language: str,
    limit: Optional[int],
    partition_count: int = 1,
    partition_index: int = 0,
) -> list[VideoRecord]:
    records: list[VideoRecord] = []
    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            query = """
                select youtube_id, r2_tar_key
                from videos
                where lower(language) = %s
                  and r2_tar_key is not null
                  and status = 'COMPLETED'
            """
            params = [language.lower()]
            if partition_count > 1:
                query += " and mod(abs(hashtext(youtube_id)), %s) = %s"
                params.extend([partition_count, partition_index])
            query += " order by random()"
            if limit:
                query += " limit %s"
                params.append(limit)
            cur.execute(query, params)
            for youtube_id, r2_tar_key in cur.fetchall():
                if youtube_id and r2_tar_key:
                    records.append(VideoRecord(youtube_id, r2_tar_key))
    return records

def normalize(value: float, min_val: float, max_val: float) -> float:
    if max_val <= min_val:
        return 0.0
    return max(0.0, min(1.0, (value - min_val) / (max_val - min_val)))

def compute_spectral_metrics(audio: np.ndarray, sr: int) -> tuple[float, float, float]:
    if audio.size == 0:
        return 0.0, 0.0, 0.0
    with timer("cpu_spectral_compute"):
        max_samples = min(audio.size, sr * 5)
        window = get_hann_window(max_samples)
        chunk = audio[:max_samples] * window
        spec = np.fft.rfft(chunk)
        power = np.abs(spec) ** 2
        freqs = get_rfftfreq(max_samples, sr)
        total_energy = float(np.sum(power)) + 1e-12
        cumulative = np.cumsum(power)
        rolloff_idx = int(np.searchsorted(cumulative, 0.85 * cumulative[-1]))
        rolloff_hz = float(freqs[min(rolloff_idx, len(freqs) - 1)])

        power_db = 10.0 * np.log10(power + 1e-12)
        threshold = power_db.max() - 60.0
        above = np.where(power_db >= threshold)[0]
        effective_nyquist_hz = float(freqs[above[-1]]) if above.size > 0 else 0.0

        high_freq_energy = float(np.sum(power[freqs >= 8000.0]))
        high_freq_ratio = high_freq_energy / total_energy
    return rolloff_hz, effective_nyquist_hz, high_freq_ratio

def classify_native_bucket(effective_nyquist_hz: float) -> str:
    """Classify audio bandwidth based on effective Nyquist frequency.

    These buckets represent actual audio bandwidth, not container sample rates.
    Thresholds are based on observed distribution in the dataset.
    """
    if effective_nyquist_hz < 8000:
        return "<16k"
    if effective_nyquist_hz < 12000:
        return "16-24k"
    if effective_nyquist_hz < 16000:
        return "24-32k"
    if effective_nyquist_hz < 22050:
        return "32-44k"
    return ">=44k"

def compute_audio_metrics(audio: np.ndarray, sr: int) -> dict[str, float]:
    with timer("audio_metrics_basic"):
        audio = audio.astype(np.float32, copy=False)
        rms = float(np.sqrt(np.mean(audio ** 2) + 1e-12))
        loudness_rms_db = 20.0 * math.log10(rms + 1e-12)
        clipping_fraction = float(np.mean(np.abs(audio) > 0.99))
        dc_offset = float(np.mean(audio))
        noise_floor = float(np.percentile(np.abs(audio), 10))
        snr_proxy = 20.0 * math.log10((rms + 1e-9) / (noise_floor + 1e-9))

    with timer("audio_metrics_lufs"):
        loudness_lufs = -100.0
        if pyln is not None:
            try:
                meter = get_lufs_meter(sr)
                if meter is not None:
                    loudness_lufs = float(meter.integrated_loudness(audio))
            except Exception:
                loudness_lufs = -100.0

    rolloff_hz, effective_nyquist_hz, high_freq_ratio = compute_spectral_metrics(audio, sr)

    return {
        "loudness_lufs": loudness_lufs,
        "loudness_rms_db": loudness_rms_db,
        "clipping_fraction": clipping_fraction,
        "dc_offset": dc_offset,
        "snr_proxy": snr_proxy,
        "spectral_rolloff_hz": rolloff_hz,
        "effective_nyquist_hz": effective_nyquist_hz,
        "high_freq_energy_ratio": high_freq_ratio,
    }

def compute_quality_score(metrics: dict[str, float], has_music: bool) -> float:
    score = (
        0.3 * (1.0 - metrics["clipping_fraction"])
        + 0.2 * normalize(metrics["snr_proxy"], 10, 40)
        + 0.2 * (1.0 if not has_music else 0.5)
        + 0.2 * normalize(metrics["loudness_lufs"], -40, -14)
        + 0.1 * (1.0 if abs(metrics["dc_offset"]) < 0.01 else 0.5)
    )
    return max(0.0, min(1.0, score))

def resample_audio(
    audio: np.ndarray,
    src_sr: int,
    dst_sr: int,
) -> np.ndarray:
    if src_sr == dst_sr:
        return audio
    with timer("cpu_resample"):
        gcd = math.gcd(src_sr, dst_sr)
        up = dst_sr // gcd
        down = src_sr // gcd
        return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)

def encode_flac(audio: np.ndarray, sr: int) -> bytes:
    with timer("flac_encode"):
        buf = io.BytesIO()
        sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
        return buf.getvalue()

def duration_bucket(duration_sec: float) -> Optional[tuple[str, float]]:
    if duration_sec < 3.0:
        return None
    if duration_sec < 6.0:
        return "3-6s", 0.3
    if duration_sec < 12.0:
        return "6-12s", 1.0
    if duration_sec < 24.0:
        return "12-24s", 0.8
    if duration_sec <= 60.0:
        return "24-60s", 0.5
    return None

def compute_frame_rms(
    audio: np.ndarray, sr: int, frame_sec: float, hop_sec: float
) -> tuple[np.ndarray, np.ndarray]:
    frame_len = max(1, int(round(frame_sec * sr)))
    hop_len = max(1, int(round(hop_sec * sr)))
    if audio.size < frame_len:
        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

    if not audio.flags["C_CONTIGUOUS"]:
        audio = np.ascontiguousarray(audio)

    num_frames = 1 + (audio.size - frame_len) // hop_len
    if num_frames <= 0:
        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

    stride = audio.strides[0]
    frames = np.lib.stride_tricks.as_strided(
        audio,
        shape=(num_frames, frame_len),
        strides=(hop_len * stride, stride),
    )
    rms = np.sqrt(np.mean(frames * frames, axis=1) + 1e-12)
    times = (np.arange(num_frames, dtype=np.float32) * hop_len + (frame_len / 2)) / sr
    return times.astype(np.float32), rms.astype(np.float32)

def find_silence_boundary(
    times: np.ndarray,
    rms: np.ndarray,
    search_start: float,
    search_end: float,
    threshold_rms: float,
    median_db: float,
    db_drop: float,
) -> Optional[float]:
    if times.size == 0:
        return None
    mask = (times >= search_start) & (times <= search_end)
    if not np.any(mask):
        return None
    subset = rms[mask]
    min_idx = int(np.argmin(subset))
    min_rms = float(subset[min_idx])
    min_time = float(times[mask][min_idx])
    min_db = 20.0 * math.log10(min_rms + 1e-12)
    if min_rms <= threshold_rms or (median_db - min_db) >= db_drop:
        return min_time
    return None

def make_chunk_windows(
    audio: np.ndarray,
    sr: int,
    cfg: ChunkConfig,


sed -n '440,880p' scripts/prepare_shards.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
    cfg: ChunkConfig,
    silence_cfg: SilenceConfig,
    rng: random.Random,
) -> list[tuple[float, float]]:
    duration_sec = float(audio.shape[0] / sr)
    if not cfg.enabled or duration_sec < cfg.threshold_sec:
        return [(0.0, duration_sec)]

    if not silence_cfg.enabled:
        windows: list[tuple[float, float]] = []
        offset = 0.0
        while duration_sec - offset >= cfg.min_sec:
            remaining = duration_sec - offset
            if remaining <= cfg.max_sec:
                length = remaining
            else:
                length = rng.uniform(cfg.min_sec, cfg.max_sec)
                if remaining - length < cfg.min_sec:
                    length = remaining
            windows.append((offset, length))
            offset += length
        return windows

    times, rms = compute_frame_rms(
        audio, sr, silence_cfg.frame_sec, silence_cfg.hop_sec
    )
    if times.size == 0:
        return [(0.0, duration_sec)]

    threshold_rms = float(np.quantile(rms, silence_cfg.percentile / 100.0))
    median_db = 20.0 * math.log10(float(np.median(rms)) + 1e-12)

    windows: list[tuple[float, float]] = []
    offset = 0.0
    while duration_sec - offset > cfg.max_sec:
        remaining = duration_sec - offset
        latest_cut = duration_sec - cfg.min_sec
        if offset + cfg.min_sec >= latest_cut:
            break
        target_len = rng.uniform(cfg.min_sec, cfg.max_sec)
        if remaining - target_len < cfg.min_sec:
            target_len = remaining - cfg.min_sec
        target_end = offset + target_len
        search_low = max(offset + cfg.min_sec, target_end - silence_cfg.search_sec)
        search_high = min(offset + cfg.max_sec, target_end + silence_cfg.search_sec, latest_cut)
        if search_high <= search_low:
            search_low = offset + cfg.min_sec
            search_high = min(offset + cfg.max_sec, latest_cut)
        boundary = find_silence_boundary(
            times,
            rms,
            search_low,
            search_high,
            threshold_rms,
            median_db,
            silence_cfg.db_drop,
        )
        if boundary is None:
            if not windows:
                return [(0.0, duration_sec)]
            windows.append((offset, duration_sec - offset))
            return windows
        windows.append((offset, boundary - offset))
        offset = boundary

    windows.append((offset, duration_sec - offset))
    return windows

def make_segment_seed(base_seed: int, video_id: str, speaker: str, index: int) -> int:
    key = f"{video_id}-{speaker}-{index}-{base_seed}"
    return zlib.adler32(key.encode("utf-8")) & 0xFFFFFFFF

def parse_segment_filename(filename: str) -> Optional[tuple[str, float]]:
    """Parse segment filename to extract speaker and start time.

    Filename format: segments/SPEAKER_XX_NNNN_start-end.flac
    Example: segments/SPEAKER_00_0007_35.10-37.49.flac

    Returns (speaker, start_time) for matching against metadata.
    """
    base = filename.split("/")[-1]
    if not base.endswith(".flac"):
        return None

    # Remove .flac extension
    base = base[:-5]

    # Try to extract time range (last part after final underscore)
    parts = base.split("_")
    if len(parts) < 3:
        return None

    # Speaker is everything except last two parts (index and time range)
    speaker = "_".join(parts[:-2])

    # Time range is the last part: "35.10-37.49"
    time_range = parts[-1]
    if "-" not in time_range:
        return None

    try:
        start_str = time_range.split("-")[0]
        start_time = float(start_str)
    except (ValueError, IndexError):
        return None

    return speaker, start_time

def process_segment(
    video_id: str,
    language: str,
    segment_bytes: bytes,
    segment_meta: dict[str, Any],
    speaker: str,
    index: int,
    chunk_cfg: ChunkConfig,
    silence_cfg: SilenceConfig,
    rng_seed: int,
) -> tuple[list[SegmentOutput], dict[str, int]]:
    stats = {
        "dropped_too_short": 0,
        "dropped_too_long": 0,
        "dropped_clipped": 0,
        "chunks_emitted": 0,
        "segments_chunked": 0,
    }
    outputs: list[SegmentOutput] = []

    with timer("flac_decode"):
        try:
            audio, sr = sf.read(io.BytesIO(segment_bytes), dtype="float32")
        except Exception:
            return outputs, stats

    if audio.ndim > 1:
        audio = np.mean(audio, axis=1)

    rng = random.Random(rng_seed)
    windows = make_chunk_windows(audio, sr, chunk_cfg, silence_cfg, rng)

    music_stats = segment_meta.get("music_stats") or {}
    has_music = bool(music_stats.get("has_music", False))

    def emit_outputs(
        audio_slice: np.ndarray,
        metrics: dict[str, float],
        sample_id: str,
        segment_index: int,
        chunk_index: Optional[int],
        start: Optional[float],
        end: Optional[float],
        original_start: Optional[float],
        original_end: Optional[float],
        target_srs: list[int],
    ) -> None:
        slice_duration = float(audio_slice.shape[0] / sr)
        bucket = duration_bucket(slice_duration)
        if bucket is None:
            if slice_duration < 3.0:
                stats["dropped_too_short"] += 1
            else:
                stats["dropped_too_long"] += 1
            return
        if metrics["clipping_fraction"] > 0.1:
            stats["dropped_clipped"] += 1
            return

        duration_bucket_name, base_weight = bucket
        sampling_weight = base_weight * (0.5 if has_music else 1.0)
        quality_score = compute_quality_score(metrics, has_music)
        effective_nyquist = metrics["effective_nyquist_hz"]

        base_metadata = {
            "id": sample_id,
            "video_id": video_id,
            "language": language,
            "segment_index": segment_index,
            "chunk_index": chunk_index,
            "start": start,
            "end": end,
            "original_start": original_start,
            "original_end": original_end,
            "duration_sec": slice_duration,
            "native_sr": sr,
            "effective_nyquist_hz": effective_nyquist,
            "spectral_rolloff_hz": metrics["spectral_rolloff_hz"],
            "high_freq_energy_ratio": metrics["high_freq_energy_ratio"],
            "is_truly_wideband": bool(effective_nyquist >= 12000),
            "native_sr_bucket": classify_native_bucket(effective_nyquist),
            "loudness_lufs": metrics["loudness_lufs"],
            "loudness_rms_db": metrics["loudness_rms_db"],
            "clipping_fraction": metrics["clipping_fraction"],
            "snr_proxy": metrics["snr_proxy"],
            "quality_score": quality_score,
            "music_prob": float(music_stats.get("music_mean", 0.0)),
            "music_stats": music_stats,
            "speaker": speaker,
            # VAD regions: Full duration placeholder; random 6s crops happen during training.
            "vad_speech_regions": [[0.0, slice_duration]],
            "duration_bucket": duration_bucket_name,
            "sampling_weight": sampling_weight,
        }

        for target_sr in target_srs:
            resampled = resample_audio(audio_slice, sr, target_sr)
            resampled = np.clip(resampled, -1.0, 1.0)
            audio_bytes = encode_flac(resampled, target_sr)
            meta = dict(base_metadata)
            meta["sample_rate"] = target_sr
            meta["resample_target_sr"] = target_sr
            outputs.append(SegmentOutput(target_sr, audio_bytes, meta))

        if target_srs == [16000]:
            stats["chunks_emitted"] += 1

    sample_id_base = f"{video_id}__{speaker}_{index:04d}"
    use_chunk_suffix = len(windows) > 1
    if use_chunk_suffix:
        stats["segments_chunked"] += 1

    segment_start = segment_meta.get("start")
    segment_end = segment_meta.get("end")
    original_start = segment_meta.get("original_start")
    original_end = segment_meta.get("original_end")

    # Always emit 16k outputs; for long segments, emit chunked windows to boost sample count.
    for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
        start_idx = int(round(offset_sec * sr))
        end_idx = int(round((offset_sec + chunk_len_sec) * sr))
        audio_slice = audio[start_idx:end_idx]
        if audio_slice.size == 0:
            continue
        slice_duration = float(audio_slice.shape[0] / sr)
        slice_metrics = compute_audio_metrics(audio_slice, sr)

        chunk_start = float(segment_start) + offset_sec if segment_start is not None else None
        chunk_end = chunk_start + slice_duration if chunk_start is not None else None
        chunk_original_start = (
            float(original_start) + offset_sec if original_start is not None else None
        )
        chunk_original_end = (
            chunk_original_start + slice_duration
            if chunk_original_start is not None
            else None
        )

        sample_id = sample_id_base
        chunk_index = None
        if use_chunk_suffix:
            sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
            chunk_index = chunk_idx

        emit_outputs(
            audio_slice,
            slice_metrics,
            sample_id,
            index,
            chunk_index,
            chunk_start,
            chunk_end,
            chunk_original_start,
            chunk_original_end,
            [16000],
        )

    # Optional high-rate buckets: use full segment unless configured otherwise.
    full_metrics = compute_audio_metrics(audio, sr)
    effective_nyquist = full_metrics["effective_nyquist_hz"]
    targets_high: list[int] = []
    if effective_nyquist >= 8000:
        targets_high.append(24000)
    if effective_nyquist >= 12000:
        targets_high.append(32000)
    if effective_nyquist >= 16000:
        targets_high.append(44100)

    if targets_high:
        if chunk_cfg.all_buckets and use_chunk_suffix:
            for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
                start_idx = int(round(offset_sec * sr))
                end_idx = int(round((offset_sec + chunk_len_sec) * sr))
                audio_slice = audio[start_idx:end_idx]
                if audio_slice.size == 0:
                    continue
                slice_duration = float(audio_slice.shape[0] / sr)
                slice_metrics = compute_audio_metrics(audio_slice, sr)
                slice_effective_nyquist = slice_metrics["effective_nyquist_hz"]
                slice_targets: list[int] = []
                if slice_effective_nyquist >= 8000:
                    slice_targets.append(24000)
                if slice_effective_nyquist >= 12000:
                    slice_targets.append(32000)
                if slice_effective_nyquist >= 16000:
                    slice_targets.append(44100)
                if not slice_targets:
                    continue

                chunk_start = (
                    float(segment_start) + offset_sec if segment_start is not None else None
                )
                chunk_end = (
                    chunk_start + slice_duration if chunk_start is not None else None
                )
                chunk_original_start = (
                    float(original_start) + offset_sec if original_start is not None else None
                )
                chunk_original_end = (
                    chunk_original_start + slice_duration
                    if chunk_original_start is not None
                    else None
                )

                sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
                emit_outputs(
                    audio_slice,
                    slice_metrics,
                    sample_id,
                    index,
                    chunk_idx,
                    chunk_start,
                    chunk_end,
                    chunk_original_start,
                    chunk_original_end,
                    slice_targets,
                )
        else:
            emit_outputs(
                audio,
                full_metrics,
                sample_id_base,
                index,
                None,
                segment_start,
                segment_end,
                original_start,
                original_end,
                targets_high,
            )

    return outputs, stats

class ShardWriter:
    def __init__(self, output_dir: Path, prefix: str, target_size_bytes: int) -> None:
        self.output_dir = output_dir
        self.prefix = prefix
        self.target_size_bytes = target_size_bytes
        self.shard_index = 0
        self.current_tar: Optional[tarfile.TarFile] = None
        self.current_jsonl = None
        self.current_size = 0
        self.current_count = 0
        self.current_name = ""
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def _open_new_shard(self) -> None:
        if self.current_tar:
            self.close()
        shard_name = f"{self.prefix}-{self.shard_index:06d}"
        tar_path = self.output_dir / f"{shard_name}.tar"
        jsonl_path = self.output_dir / f"{shard_name}.jsonl"
        self.current_tar = tarfile.open(tar_path, "w")
        self.current_jsonl = jsonl_path.open("w", encoding="utf-8")
        self.current_size = 0
        self.current_count = 0
        self.current_name = shard_name
        self.shard_index += 1

    def add_sample(self, sample_id: str, audio_bytes: bytes, metadata: dict[str, Any]) -> str:
        if self.current_tar is None:
            self._open_new_shard()
        if (
            self.current_count > 0
            and self.current_size + len(audio_bytes) > self.target_size_bytes
        ):
            self._open_new_shard()

        assert self.current_tar is not None
        assert self.current_jsonl is not None

        arcname = f"{sample_id}.flac"
        info = tarfile.TarInfo(name=arcname)
        info.size = len(audio_bytes)
        info.mtime = int(time.time())
        self.current_tar.addfile(info, io.BytesIO(audio_bytes))
        self.current_jsonl.write(json.dumps(metadata) + "\n")

        self.current_size += len(audio_bytes)
        self.current_count += 1
        return self.current_name

    def close(self) -> None:
        if self.current_tar:
            self.current_tar.close()
        if self.current_jsonl:
            self.current_jsonl.close()
        self.current_tar = None
        self.current_jsonl = None

def create_prefix_placeholders(
    s3,
    bucket: str,
    languages: list[str],
    rates: list[int],
    extra_prefixes: list[str],
    subprefix: Optional[str],
) -> None:
    prefixes = []
    for lang in languages:
        for rate in rates:
            base = f"{lang}_{rate}"
            prefixes.append(base)
            if subprefix:
                prefixes.append(f"{base}/{subprefix.strip('/')}")
    prefixes.extend(extra_prefixes)
    if subprefix:
        prefixes.extend([f"{p}/{subprefix.strip('/')}" for p in extra_prefixes])

    for prefix in prefixes:
        key = f"{prefix.strip('/')}/.keep"
        try:
            s3.put_object(Bucket=bucket, Key=key, Body=b"")
        except Exception:
            continue

def upload_directory(
    s3,
    bucket: str,
    local_dir: Path,
    remote_prefix: str,
    max_concurrency: int,
) -> list[str]:
    uploaded = []
    config = TransferConfig(
        multipart_threshold=8 * 1024 * 1024,
        multipart_chunksize=64 * 1024 * 1024,


sed -n '880,1320p' scripts/prepare_shards.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
        multipart_chunksize=64 * 1024 * 1024,
        max_concurrency=max_concurrency,
        use_threads=True,
    )
    for path in sorted(local_dir.glob("*")):
        if not path.is_file():
            continue
        key = f"{remote_prefix}/{path.name}"
        s3.upload_file(str(path), bucket, key, Config=config)
        uploaded.append(key)
    return uploaded

def verify_uploaded_shard(
    s3,
    bucket: str,
    remote_prefix: str,
    shard_name: str,
    sample_rate: int,
    tmp_dir: Path,
) -> dict[str, Any]:
    tar_key = f"{remote_prefix}/{shard_name}.tar"
    jsonl_key = f"{remote_prefix}/{shard_name}.jsonl"
    tar_path = tmp_dir / f"{shard_name}.tar"
    jsonl_path = tmp_dir / f"{shard_name}.jsonl"

    s3.download_file(bucket, tar_key, str(tar_path))
    s3.download_file(bucket, jsonl_key, str(jsonl_path))

    with tarfile.open(tar_path, "r") as tf:
        tar_names = set(tf.getnames())
        with jsonl_path.open("r", encoding="utf-8") as fh:
            lines = [json.loads(line) for line in fh if line.strip()]

        missing = 0
        checked = 0
        for entry in lines[:10]:
            checked += 1
            audio_name = f"{entry['id']}.flac"
            if audio_name not in tar_names:
                missing += 1
                continue
            member = tf.getmember(audio_name)
            audio_bytes = tf.extractfile(member).read()
            audio, sr = sf.read(io.BytesIO(audio_bytes))
            if sr != sample_rate:
                missing += 1

    return {
        "checked_samples": checked,
        "missing_or_mismatch": missing,
        "total_entries": len(lines),
    }

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Prepare sharded datasets from cleaned tars.")
    parser.add_argument("--language", required=True)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--output-dir", type=Path, default=Path("processed_data"))
    parser.add_argument("--tmp-dir", type=Path, default=Path("/tmp/xcodec_work"))
    parser.add_argument("--target-shard-size-mb", type=int, default=5120)
    parser.add_argument("--download-workers", type=int, default=16)
    parser.add_argument("--segment-workers", type=int, default=8)
    parser.add_argument(
        "--segment-max-inflight",
        type=int,
        default=0,
        help="Limit number of in-flight segment tasks (0=auto).",
    )
    parser.add_argument(
        "--skip-downloads",
        action="store_true",
        help="Use only local tar files; do not download from R2.",
    )
    parser.add_argument(
        "--local-tars-dir",
        type=Path,
        default=None,
        help="Directory containing pre-downloaded tar files for this language.",
    )
    parser.add_argument(
        "--partition-count",
        type=int,
        default=1,
        help="Total number of partitions for parallel processing.",
    )
    parser.add_argument(
        "--partition-index",
        type=int,
        default=0,
        help="Partition index for this run (0-based).",
    )
    parser.add_argument("--val-ratio", type=float, default=0.01)
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument(
        "--chunk-long",
        dest="chunk_long",
        action="store_true",
        default=True,
        help="Split long segments into 6-10s windows for the 16k bucket.",
    )
    parser.add_argument(
        "--no-chunk-long",
        dest="chunk_long",
        action="store_false",
        help="Disable chunking long segments.",
    )
    parser.add_argument("--chunk-min-sec", type=float, default=6.0)
    parser.add_argument("--chunk-max-sec", type=float, default=10.0)
    parser.add_argument("--chunk-threshold-sec", type=float, default=12.0)
    parser.add_argument(
        "--chunk-all-buckets",
        action="store_true",
        help="Also chunk high-rate buckets (default: keep full segments).",
    )
    parser.add_argument(
        "--silence-snap",
        dest="silence_snap",
        action="store_true",
        default=True,
        help="Snap chunk boundaries to low-energy points when possible.",
    )
    parser.add_argument(
        "--no-silence-snap",
        dest="silence_snap",
        action="store_false",
        help="Disable silence-aware chunk boundaries.",
    )
    parser.add_argument("--silence-frame-ms", type=float, default=50.0)
    parser.add_argument("--silence-hop-ms", type=float, default=10.0)
    parser.add_argument("--silence-search-sec", type=float, default=0.5)
    parser.add_argument("--silence-percentile", type=float, default=20.0)
    parser.add_argument("--silence-db-drop", type=float, default=20.0)
    parser.add_argument("--upload", action="store_true")
    parser.add_argument("--verify-upload", action="store_true")
    parser.add_argument("--create-prefixes", action="store_true")
    parser.add_argument("--r2-subprefix", default="")
    parser.add_argument(
        "--profile-timing",
        action="store_true",
        help="Enable detailed timing instrumentation for profiling.",
    )
    return parser.parse_args()

def main() -> None:
    global _timing_enabled
    args = parse_args()
    _timing_enabled = args.profile_timing
    env = load_env(ENV_PATH)
    db_url = env.get("DATABASE_URL")
    if not db_url:
        raise SystemExit("DATABASE_URL not found in .env")

    cleaned_bucket = env.get("R2_ClEANED_BUCKET") or env.get("R2_CLEANED_BUCKET")
    shard_bucket = env.get("R2_BUCKET_SHARDS")
    prefix = env.get("R2_PREFIX")
    if not cleaned_bucket:
        raise SystemExit("R2 cleaned bucket not found in .env")

    if args.upload and not shard_bucket:
        raise SystemExit("R2_BUCKET_SHARDS not found in .env")

    if args.chunk_min_sec <= 0 or args.chunk_max_sec <= 0:
        raise SystemExit("Chunk durations must be positive.")
    if args.chunk_min_sec > args.chunk_max_sec:
        raise SystemExit("chunk-min-sec must be <= chunk-max-sec.")
    if args.chunk_threshold_sec < args.chunk_min_sec:
        raise SystemExit("chunk-threshold-sec must be >= chunk-min-sec.")
    if args.silence_frame_ms <= 0 or args.silence_hop_ms <= 0:
        raise SystemExit("Silence frame/hop must be positive.")
    if args.silence_search_sec <= 0:
        raise SystemExit("silence-search-sec must be positive.")
    if not 0.0 < args.silence_percentile < 100.0:
        raise SystemExit("silence-percentile must be between 0 and 100.")
    if args.silence_db_drop <= 0:
        raise SystemExit("silence-db-drop must be positive.")

    if args.partition_count < 1:
        raise SystemExit("partition-count must be >= 1.")
    if not 0 <= args.partition_index < args.partition_count:
        raise SystemExit("partition-index must be between 0 and partition-count - 1.")

    partition_tag = ""
    if args.partition_count > 1:
        partition_tag = f"p{args.partition_index:02d}"

    chunk_cfg = ChunkConfig(
        enabled=args.chunk_long,
        min_sec=args.chunk_min_sec,
        max_sec=args.chunk_max_sec,
        threshold_sec=args.chunk_threshold_sec,
        all_buckets=args.chunk_all_buckets,
    )
    silence_cfg = SilenceConfig(
        enabled=args.silence_snap,
        frame_sec=args.silence_frame_ms / 1000.0,
        hop_sec=args.silence_hop_ms / 1000.0,
        search_sec=args.silence_search_sec,
        percentile=args.silence_percentile,
        db_drop=args.silence_db_drop,
    )
    language = args.language.strip().lower()
    records = fetch_video_records(
        db_url,
        language,
        args.limit,
        args.partition_count,
        args.partition_index,
    )
    if not records:
        raise SystemExit(f"No records found for language={language}")

    s3 = make_r2_client(env)
    if args.create_prefixes and args.upload:
        create_prefix_placeholders(
            s3,
            shard_bucket,
            LANGUAGES,
            [16, 24, 32, 44],
            ["manifests", "stats", "reports", "checks"],
            args.r2_subprefix or None,
        )

    tmp_dir = args.tmp_dir
    output_dir = args.output_dir
    if partition_tag:
        if tmp_dir.name != partition_tag:
            tmp_dir = tmp_dir / partition_tag
        if output_dir.name != partition_tag:
            output_dir = output_dir / partition_tag
    tmp_dir.mkdir(parents=True, exist_ok=True)
    download_dir = tmp_dir / "downloads"
    download_dir.mkdir(parents=True, exist_ok=True)

    config = TransferConfig(
        multipart_threshold=8 * 1024 * 1024,
        multipart_chunksize=64 * 1024 * 1024,
        max_concurrency=max(4, args.download_workers),
        use_threads=True,
    )

    def download_record(record: VideoRecord) -> Path:
        resolved = resolve_r2_key(s3, cleaned_bucket, record.r2_tar_key, prefix)
        if not resolved:
            raise FileNotFoundError(f"Missing tar in R2: {record.r2_tar_key}")
        local_path = download_dir / f"{record.youtube_id}.tar"
        if not local_path.exists():
            s3.download_file(cleaned_bucket, resolved, str(local_path), Config=config)
        return local_path

    download_paths: dict[str, Path] = {}
    if args.skip_downloads:
        local_dir = args.local_tars_dir or download_dir
        if not local_dir.exists():
            raise SystemExit(f"Local tar directory not found: {local_dir}")
        missing = 0
        for record in tqdm(records, desc="local"):
            local_path = local_dir / f"{record.youtube_id}.tar"
            if local_path.exists():
                download_paths[record.youtube_id] = local_path
            else:
                missing += 1
        if missing:
            print(f"Missing local tars: {missing} (skipping)")
    else:
        with ThreadPoolExecutor(max_workers=args.download_workers) as executor:
            futures = {
                executor.submit(download_record, record): record.youtube_id
                for record in records
            }
            for future in tqdm(as_completed(futures), total=len(futures), desc="download"):
                youtube_id = futures[future]
                try:
                    download_paths[youtube_id] = future.result()
                except Exception as exc:
                    print(f"download failed {youtube_id}: {exc}")

    if not download_paths:
        raise SystemExit("No downloads succeeded.")

    video_ids = list(download_paths.keys())
    random.seed(args.seed)
    val_count = max(1, int(len(video_ids) * args.val_ratio)) if len(video_ids) > 1 else 0
    val_ids = set(random.sample(video_ids, val_count)) if val_count > 0 else set()

    manifest_dir = output_dir / "manifests"
    stats_dir = output_dir / "stats"
    manifest_dir.mkdir(parents=True, exist_ok=True)
    stats_dir.mkdir(parents=True, exist_ok=True)

    shard_prefix = "shard" if not partition_tag else f"shard-{partition_tag}"
    writers = {
        16000: ShardWriter(output_dir / f"{language}_16", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
        24000: ShardWriter(output_dir / f"{language}_24", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
        32000: ShardWriter(output_dir / f"{language}_32", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
        44100: ShardWriter(output_dir / f"{language}_44", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
    }

    manifest_suffix = f"_{partition_tag}" if partition_tag else ""
    train_manifest = (manifest_dir / f"{language}_train{manifest_suffix}.jsonl").open(
        "w", encoding="utf-8"
    )
    val_manifest = (manifest_dir / f"{language}_val{manifest_suffix}.jsonl").open(
        "w", encoding="utf-8"
    )

    stats = {
        "videos": 0,
        "segments_total": 0,
        "segments_kept": {16000: 0, 24000: 0, 32000: 0, 44100: 0},
        "dropped_too_short": 0,
        "dropped_too_long": 0,
        "dropped_clipped": 0,
        "segments_chunked": 0,
        "chunks_emitted": 0,
        "native_sr_buckets": {},
    }

    segment_workers = max(1, args.segment_workers)
    max_inflight = (
        args.segment_max_inflight
        if args.segment_max_inflight > 0
        else max(128, segment_workers * 2)
    )
    pending: set[Any] = set()
    future_info: dict[Any, bool] = {}

    def handle_future(future) -> None:
        is_val = future_info.pop(future, False)
        outputs, seg_stats = future.result()
        stats["dropped_too_short"] += seg_stats["dropped_too_short"]
        stats["dropped_too_long"] += seg_stats["dropped_too_long"]
        stats["dropped_clipped"] += seg_stats["dropped_clipped"]
        stats["segments_chunked"] += seg_stats.get("segments_chunked", 0)
        stats["chunks_emitted"] += seg_stats.get("chunks_emitted", 0)
        for out in outputs:
            stats["segments_kept"][out.sample_rate] += 1
            if out.sample_rate == 16000:
                bucket = out.metadata.get("native_sr_bucket")
                if bucket:
                    stats["native_sr_buckets"][bucket] = (
                        stats["native_sr_buckets"].get(bucket, 0) + 1
                    )
            writer = writers[out.sample_rate]
            shard_name = writer.add_sample(out.metadata["id"], out.audio_bytes, out.metadata)
            manifest_entry = dict(out.metadata)
            manifest_entry["shard"] = f"{shard_name}.tar"
            manifest_entry["audio"] = f"{out.metadata['id']}.flac"
            line = json.dumps(manifest_entry) + "\n"
            if is_val:
                val_manifest.write(line)
            else:
                train_manifest.write(line)

    def drain_completed() -> None:
        nonlocal pending
        if not pending:
            return
        done, pending = wait(pending, return_when=FIRST_COMPLETED)
        for future in done:
            handle_future(future)

    with ThreadPoolExecutor(max_workers=segment_workers) as executor:
        for video_id, tar_path in tqdm(download_paths.items(), desc="process"):
            stats["videos"] += 1
            is_val = video_id in val_ids
            with tarfile.open(tar_path, "r") as tf:
                meta = json.loads(tf.extractfile(f"{video_id}/metadata.json").read())
                usable_segments = [
                    s for s in meta.get("segments", []) if s.get("status") == "usable"
                ]
                # Create lookup by (speaker, start) for robust matching
                # The filename contains 'start' timestamp (relative to intro-skipped audio),
                # not 'original_start' (absolute time in original video)
                segment_map = {
                    (seg.get("speaker"), round(seg.get("start", 0.0), 2)): seg
                    for seg in usable_segments
                }

                members = [m for m in tf.getmembers() if m.name.endswith(".flac")]
                stats["segments_total"] += len(members)

                for member in members:
                    parsed = parse_segment_filename(member.name)
                    if not parsed:
                        continue
                    speaker, start_time = parsed
                    # Match by (speaker, start) with rounding for float precision
                    seg_meta = segment_map.get((speaker, round(start_time, 2)))
                    if not seg_meta:
                        # Try fuzzy match for edge cases (tolerance 0.1s)
                        matches = [
                            s for s in usable_segments
                            if s.get("speaker") == speaker
                            and abs(s.get("start", 0.0) - start_time) < 0.1
                        ]
                        if not matches:
                            continue
                        seg_meta = matches[0]

                    while len(pending) >= max_inflight:
                        drain_completed()

                    segment_bytes = tf.extractfile(member).read()
                    # Extract index from filename for ID generation (cosmetic only)
                    filename_parts = member.name.split("_")
                    try:
                        index = int(filename_parts[-2])
                    except (ValueError, IndexError):
                        index = 0
                    rng_seed = make_segment_seed(args.seed, video_id, speaker, index)

                    future = executor.submit(
                        process_segment,
                        video_id,
                        language,
                        segment_bytes,
                        seg_meta,
                        speaker,
                        index,
                        chunk_cfg,
                        silence_cfg,
                        rng_seed,
                    )
                    future_info[future] = is_val
                    pending.add(future)

        while pending:
            drain_completed()

    for writer in writers.values():
        writer.close()
    train_manifest.close()
    val_manifest.close()

    stats_suffix = f"_{partition_tag}" if partition_tag else ""
    stats_path = stats_dir / f"{language}_stats{stats_suffix}.json"
    stats_path.write_text(json.dumps(stats, indent=2))


sed -n '1320,1760p' scripts/prepare_shards.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:

    if args.upload:
        r2_subprefix = args.r2_subprefix.strip("/")
        upload_map = {
            16000: f"{language}_16",
            24000: f"{language}_24",
            32000: f"{language}_32",
            44100: f"{language}_44",
        }
        for sr, local_prefix in upload_map.items():
            remote_prefix = local_prefix if not r2_subprefix else f"{local_prefix}/{r2_subprefix}"
            upload_directory(
                s3,
                shard_bucket,
                output_dir / local_prefix,
                remote_prefix,
                args.download_workers,
            )
        upload_directory(
            s3,
            shard_bucket,
            manifest_dir,
            "manifests" if not r2_subprefix else f"manifests/{r2_subprefix}",
            args.download_workers,
        )
        upload_directory(
            s3,
            shard_bucket,
            stats_dir,
            "stats" if not r2_subprefix else f"stats/{r2_subprefix}",
            args.download_workers,
        )

        if args.verify_upload:
            verify_dir = tmp_dir / "verify"
            verify_dir.mkdir(parents=True, exist_ok=True)
            for sr, local_prefix in upload_map.items():
                shard_dir = output_dir / local_prefix
                shard_files = sorted(p.stem for p in shard_dir.glob("*.tar"))
                if not shard_files:
                    continue
                shard_name = shard_files[0]
                remote_prefix = local_prefix if not r2_subprefix else f"{local_prefix}/{r2_subprefix}"
                result = verify_uploaded_shard(
                    s3,
                    shard_bucket,
                    remote_prefix,
                    shard_name,
                    sr,
                    verify_dir,
                )
                print(
                    f"verify {language} {sr}k: "
                    f"checked={result['checked_samples']} "
                    f"missing_or_mismatch={result['missing_or_mismatch']} "
                    f"entries={result['total_entries']}"
                )

    if _timing_enabled:
        print_timing_stats()

    print(f"Done. Stats written to {stats_path}")

if __name__ == "__main__":
    main()


**Verifying pipeline metadata usage**

I'm checking other scripts like fetch_languages for metadata references to confirm alignment with the current pipeline plan.


sed -n '1,200p' scripts/fetch_languages.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Optional

import psycopg

ENV_PATH = Path(".env")

def load_env_value(env_path: Path, key: str) -> Optional[str]:
    if not env_path.exists():
        return None
    for line in env_path.read_text().splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in stripped:
            continue
        if stripped.startswith("export "):
            stripped = stripped[len("export ") :]
        k, v = stripped.split("=", 1)
        if k.strip() == key:
            return v.strip()
    return None

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fetch youtube_id -> language mapping from Supabase."
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("language_lookup.json"),
        help="Output JSON path (default: language_lookup.json).",
    )
    parser.add_argument(
        "--only-completed",
        action="store_true",
        help="Only include videos with status = COMPLETED.",
    )
    return parser.parse_args()

def main() -> None:
    args = parse_args()
    db_url = load_env_value(ENV_PATH, "DATABASE_URL")
    if not db_url:
        raise SystemExit("DATABASE_URL not found in .env")

    where_clause = "where language is not null"
    if args.only_completed:
        where_clause += " and status = 'COMPLETED'"

    mapping: dict[str, str] = {}
    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            cur.execute(
                f"""
                select youtube_id, language
                from videos
                {where_clause};
                """
            )
            for youtube_id, language in cur.fetchall():
                if youtube_id and language:
                    mapping[youtube_id] = language.strip().lower()

    args.output.write_text(json.dumps(mapping, indent=2, sort_keys=True))
    print(f"Wrote {len(mapping):,} entries to {args.output}")

if __name__ == "__main__":
    main()


**Mapping pipeline data sources**

The pipeline uses direct DB queries per language, not a language lookup table, fetching records with prepare_shards. Next steps include testing the pipeline end-to-end on a single video from R2 via Supabase, carefully reading needed keys from the .env without exposing secrets, and querying the Supabase table using Python with provided credentials.


**Detailing supabase query and bucket usage**

The pipeline’s prepare_shards script downloads from the cleaned bucket set in the env, so ensuring R2_CLEANED_BUCKET is correctly configured is key. Next, I'll run a minimal psycopg query to fetch one video record with an R2 key, then test downloading and parsing its tar archive, validating metadata consistency against audio segments. I'll also inspect the training data loader in X-Codec-2.0 to confirm correct .tar and JSONL handling, likely with WebDataset.


rg -n "webdataset|wds|tar" -S X-Codec-2.0


Output:
Exit code: 0
Wall time: 0 seconds
Output:
X-Codec-2.0/train.py:46:        print("No checkpoint found, starting training from scratch.")
X-Codec-2.0/requirements.txt:99:rotary-embedding-torch==0.8.4
X-Codec-2.0/lightning_module.py:154:            semantic_target = self.semantic_model(feats[:,0,:,:])
X-Codec-2.0/lightning_module.py:156:            semantic_target = semantic_target.hidden_states[16]
X-Codec-2.0/lightning_module.py:157:            semantic_target = semantic_target.detach()
X-Codec-2.0/lightning_module.py:159:        semantic_target = semantic_target.transpose(1, 2)
X-Codec-2.0/lightning_module.py:160:        semantic_target_processed = self.SemanticEncoder_module(semantic_target)
X-Codec-2.0/lightning_module.py:162:        vq_emb = torch.cat([semantic_target_processed, vq_emb], dim=1)
X-Codec-2.0/lightning_module.py:192:            'semantic_recon_loss': F.mse_loss(semantic_recon, semantic_target),
X-Codec-2.0/inference.py:41:        if key.startswith('CodecEnc.'):
X-Codec-2.0/inference.py:45:        elif key.startswith('generator.'):
X-Codec-2.0/inference.py:49:        elif key.startswith('fc_post_a.'):
X-Codec-2.0/inference.py:53:        elif key.startswith('SemanticEncoder_module.'):
X-Codec-2.0/inference.py:57:        elif key.startswith('fc_prior.'):
X-Codec-2.0/inference.py:98:        target_wav_path = join(wav_dir, basename(wav_path))
X-Codec-2.0/inference.py:117:            semantic_target = semantic_model(feat[:,  :,:])
X-Codec-2.0/inference.py:119:            semantic_target = semantic_target.hidden_states[16]
X-Codec-2.0/inference.py:121:            semantic_target = semantic_target.transpose(1, 2)
X-Codec-2.0/inference.py:122:            semantic_target = SemanticEncoder_module(semantic_target)
X-Codec-2.0/inference.py:125:            vq_emb = torch.cat([semantic_target, vq_emb], dim=1)
X-Codec-2.0/inference.py:135:        sf.write(target_wav_path, recon, sr)
X-Codec-2.0/get_tsv.py:10:        rel_path = os.path.relpath(file_path, start=root_dir)
X-Codec-2.0/get_tsv.py:14:        for start in range(0, waveform.numel(), batch_size):
X-Codec-2.0/get_tsv.py:15:            end = min(start + batch_size, waveform.numel())
X-Codec-2.0/get_tsv.py:16:            if torch.isnan(waveform.view(-1)[start:end]).any():
X-Codec-2.0/inference_save_code.py:129:        if key.startswith('CodecEnc.'):
X-Codec-2.0/inference_save_code.py:132:        elif key.startswith('generator.'):
X-Codec-2.0/inference_save_code.py:135:        elif key.startswith('fc_post_a.'):
X-Codec-2.0/inference_save_code.py:138:        elif key.startswith('SemanticEncoder_module.'):
X-Codec-2.0/inference_save_code.py:141:        elif key.startswith('fc_prior.'):
X-Codec-2.0/inference_save_code.py:211:            semantic_target = semantic_model(feats[:,0,:,:].to(device))
X-Codec-2.0/inference_save_code.py:212:            semantic_target = semantic_target.hidden_states[16]
X-Codec-2.0/inference_save_code.py:213:            semantic_target = semantic_target.transpose(1, 2)
X-Codec-2.0/inference_save_code.py:214:            semantic_target = SemanticEncoder_module(semantic_target)
X-Codec-2.0/inference_save_code.py:216:            vq_emb = torch.cat([semantic_target, vq_emb], dim=1)
X-Codec-2.0/config/dataset/default.yaml:2:  _target_: data_module.FSDataset_add_STFT
X-Codec-2.0/vq/codec_decoder_vocos.py:13:# from rotary_embedding_torch import RotaryEmbedding
X-Codec-2.0/vq/codec_decoder_vocos.py:14:from torchtune.modules import RotaryPositionalEmbeddings
X-Codec-2.0/vq/codec_decoder_vocos.py:329:        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
X-Codec-2.0/vq/codec_decoder_vocos.py:333:            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
X-Codec-2.0/vq/codec_encoder.py:11:from torchtune.modules import RotaryPositionalEmbeddings
X-Codec-2.0/vq/codec_encoder.py:131:        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
X-Codec-2.0/vq/codec_encoder.py:134:            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
X-Codec-2.0/vq/codec_encoder.py:208:        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
X-Codec-2.0/vq/codec_encoder.py:212:            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
X-Codec-2.0/vq/bs_roformer5.py:8:# from rotary_embedding_torch import RotaryEmbedding
X-Codec-2.0/vq/bs_roformer5.py:10:from torchtune.modules import RotaryPositionalEmbeddings
X-Codec-2.0/vq/bs_roformer5.py:45:    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings):
X-Codec-2.0/vq/bs_roformer5.py:52:        self.rotary_embed = rotary_embed
X-Codec-2.0/vq/bs_roformer5.py:77:        q = self.rotary_embed(q)
X-Codec-2.0/vq/bs_roformer5.py:78:        k = self.rotary_embed(k)
X-Codec-2.0/vq/bs_roformer5.py:92:    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings):
X-Codec-2.0/vq/bs_roformer5.py:100:        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed)
X-Codec-2.0/vq/bs_roformer5.py:114:    rotary_embed_128 = RotaryPositionalEmbeddings(dim=128)
X-Codec-2.0/vq/bs_roformer5.py:118:        rotary_embed=rotary_embed_128


**Searching for dataset definitions**

I'm looking into the data_module to identify any dataset definitions or relevant data structures for the task.


rg -n "Dataset|data_module|manifest|jsonl|tar" -S X-Codec-2.0


Output:
Exit code: 0
Wall time: 0 seconds
Output:
X-Codec-2.0/requirements.txt:99:rotary-embedding-torch==0.8.4
X-Codec-2.0/train.py:14:from data_module import DataModule
X-Codec-2.0/train.py:46:        print("No checkpoint found, starting training from scratch.")
X-Codec-2.0/lightning_module.py:154:            semantic_target = self.semantic_model(feats[:,0,:,:])
X-Codec-2.0/lightning_module.py:156:            semantic_target = semantic_target.hidden_states[16]
X-Codec-2.0/lightning_module.py:157:            semantic_target = semantic_target.detach()
X-Codec-2.0/lightning_module.py:159:        semantic_target = semantic_target.transpose(1, 2)
X-Codec-2.0/lightning_module.py:160:        semantic_target_processed = self.SemanticEncoder_module(semantic_target)
X-Codec-2.0/lightning_module.py:162:        vq_emb = torch.cat([semantic_target_processed, vq_emb], dim=1)
X-Codec-2.0/lightning_module.py:192:            'semantic_recon_loss': F.mse_loss(semantic_recon, semantic_target),
X-Codec-2.0/get_tsv.py:10:        rel_path = os.path.relpath(file_path, start=root_dir)
X-Codec-2.0/get_tsv.py:14:        for start in range(0, waveform.numel(), batch_size):
X-Codec-2.0/get_tsv.py:15:            end = min(start + batch_size, waveform.numel())
X-Codec-2.0/get_tsv.py:16:            if torch.isnan(waveform.view(-1)[start:end]).any():
X-Codec-2.0/data_module.py:13:from torch.utils.data import Dataset, DataLoader
X-Codec-2.0/data_module.py:32:        ds = FSDataset(phase, self.cfg)
X-Codec-2.0/data_module.py:33:        # ds = FSDataset_add_STFT(phase, self.cfg)
X-Codec-2.0/data_module.py:53:class FSDataset(Dataset):
X-Codec-2.0/data_module.py:54:    """Dataset batching wav, mel
X-Codec-2.0/data_module.py:135:    data_module = DataModule(cfg)
X-Codec-2.0/data_module.py:138:    train_loader = data_module.val_dataloader()
X-Codec-2.0/inference_save_code.py:18:from torch.utils.data import Dataset, DataLoader
X-Codec-2.0/inference_save_code.py:54:class WaveDataset(Dataset):
X-Codec-2.0/inference_save_code.py:129:        if key.startswith('CodecEnc.'):
X-Codec-2.0/inference_save_code.py:132:        elif key.startswith('generator.'):
X-Codec-2.0/inference_save_code.py:135:        elif key.startswith('fc_post_a.'):
X-Codec-2.0/inference_save_code.py:138:        elif key.startswith('SemanticEncoder_module.'):
X-Codec-2.0/inference_save_code.py:141:        elif key.startswith('fc_prior.'):
X-Codec-2.0/inference_save_code.py:189:    dataset = WaveDataset(file_list=current_file_list, sampling_rate=sr, root_dir=args.input_dir)
X-Codec-2.0/inference_save_code.py:211:            semantic_target = semantic_model(feats[:,0,:,:].to(device))
X-Codec-2.0/inference_save_code.py:212:            semantic_target = semantic_target.hidden_states[16]
X-Codec-2.0/inference_save_code.py:213:            semantic_target = semantic_target.transpose(1, 2)
X-Codec-2.0/inference_save_code.py:214:            semantic_target = SemanticEncoder_module(semantic_target)
X-Codec-2.0/inference_save_code.py:216:            vq_emb = torch.cat([semantic_target, vq_emb], dim=1)
X-Codec-2.0/config/dataset/default.yaml:2:  _target_: data_module.FSDataset_add_STFT
X-Codec-2.0/vq/codec_decoder_vocos.py:13:# from rotary_embedding_torch import RotaryEmbedding
X-Codec-2.0/vq/codec_decoder_vocos.py:14:from torchtune.modules import RotaryPositionalEmbeddings
X-Codec-2.0/vq/codec_decoder_vocos.py:329:        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
X-Codec-2.0/vq/codec_decoder_vocos.py:333:            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
X-Codec-2.0/vq/codec_encoder.py:11:from torchtune.modules import RotaryPositionalEmbeddings
X-Codec-2.0/vq/codec_encoder.py:131:        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
X-Codec-2.0/vq/codec_encoder.py:134:            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
X-Codec-2.0/vq/codec_encoder.py:208:        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
X-Codec-2.0/vq/codec_encoder.py:212:            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
X-Codec-2.0/inference.py:41:        if key.startswith('CodecEnc.'):
X-Codec-2.0/inference.py:45:        elif key.startswith('generator.'):
X-Codec-2.0/inference.py:49:        elif key.startswith('fc_post_a.'):
X-Codec-2.0/inference.py:53:        elif key.startswith('SemanticEncoder_module.'):
X-Codec-2.0/inference.py:57:        elif key.startswith('fc_prior.'):
X-Codec-2.0/inference.py:98:        target_wav_path = join(wav_dir, basename(wav_path))
X-Codec-2.0/inference.py:117:            semantic_target = semantic_model(feat[:,  :,:])
X-Codec-2.0/inference.py:119:            semantic_target = semantic_target.hidden_states[16]
X-Codec-2.0/inference.py:121:            semantic_target = semantic_target.transpose(1, 2)
X-Codec-2.0/inference.py:122:            semantic_target = SemanticEncoder_module(semantic_target)
X-Codec-2.0/inference.py:125:            vq_emb = torch.cat([semantic_target, vq_emb], dim=1)
X-Codec-2.0/inference.py:135:        sf.write(target_wav_path, recon, sr)
X-Codec-2.0/vq/bs_roformer5.py:8:# from rotary_embedding_torch import RotaryEmbedding
X-Codec-2.0/vq/bs_roformer5.py:10:from torchtune.modules import RotaryPositionalEmbeddings
X-Codec-2.0/vq/bs_roformer5.py:45:    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings):
X-Codec-2.0/vq/bs_roformer5.py:52:        self.rotary_embed = rotary_embed
X-Codec-2.0/vq/bs_roformer5.py:77:        q = self.rotary_embed(q)
X-Codec-2.0/vq/bs_roformer5.py:78:        k = self.rotary_embed(k)
X-Codec-2.0/vq/bs_roformer5.py:92:    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings):
X-Codec-2.0/vq/bs_roformer5.py:100:        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed)
X-Codec-2.0/vq/bs_roformer5.py:114:    rotary_embed_128 = RotaryPositionalEmbeddings(dim=128)
X-Codec-2.0/vq/bs_roformer5.py:118:        rotary_embed=rotary_embed_128


Summary: auto


**Opening data module for dataset format**


sed -n '1,220p' X-Codec-2.0/data_module.py