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

import re
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import numpy as np
import pytorch_lightning as pl
import random
import librosa
from os.path import basename, exists, join
from torch.utils.data import Dataset, DataLoader
import hydra
import utils
import torchaudio
from transformers import AutoFeatureExtractor
from torchaudio.transforms import Resample
from tqdm import tqdm
from torchaudio.transforms import Resample
class DataModule(pl.LightningDataModule):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg

        ocwd = hydra.utils.get_original_cwd()
        self.ocwd = ocwd

    def get_loader(self, phase):
        phase_cfg = self.cfg.dataset.get(phase)
        batch_size = phase_cfg.batch_size
        ds = FSDataset(phase, self.cfg)
        # ds = FSDataset_add_STFT(phase, self.cfg)
        dl = DataLoader(ds,
                        batch_size=batch_size,
                        shuffle=phase_cfg.shuffle,
                        num_workers=28,
                        collate_fn=ds.collate_fn,
                        pin_memory=True,
                        persistent_workers=True)

        return dl

    def train_dataloader(self):
        return self.get_loader('train')

    def val_dataloader(self):
        return self.get_loader('val')

    def test_dataloader(self):
        pass

class FSDataset(Dataset):
    """Dataset batching wav, mel
    and other acoustic features

    Args:
        phase: train, val, test
        cfg: hydra config
    """
    def __init__(self, phase, cfg):
        self.phase = phase
        self.cfg = cfg
        self.phase_cfg = cfg.dataset.get(phase)
        self.ocwd = hydra.utils.get_original_cwd()

        self.sr = cfg.preprocess.audio.sr

        # self.filelist = utils.read_filelist(join(self.ocwd, self.phase_cfg.filelist))
        self.filelist = self.get_filelist(self.phase_cfg.filelist)
        self.min_audio_length = cfg.dataset.min_audio_length
        self.feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
    def __len__(self):
        return len(self.filelist)

    def load_wav(self, path):
        wav, sr = librosa.load(path, sr=self.sr)
        return wav

    def get_filelist(self, fpath):
        with open(fpath, 'r') as f:
            # flist = [l.strip() for l in f if l.strip()]
            flist = [l.strip().split('\t')[0] for l in f if l.strip()]
        return flist

    def __getitem__(self, idx):
        # (  wavpath,fid) = self.filelist[idx]
        wavpath  = self.filelist[idx]
        wavpath_full = join(self.cfg.preprocess.datasets.LibriSpeech.root, wavpath)
        # wav = self.load_wav(wavpath)
        # wav = torch.from_numpy(wav)

        wav,sr=torchaudio.load(wavpath_full)

        if sr != 16000:
            wav = Resample(sr, 16000)(wav)
        wav = wav[0,:]
        length = wav.shape[0]
        # length = wav.shape[1]
        if length < self.min_audio_length:
            wav = F.pad(wav, (0, self.min_audio_length - length))
            length = wav.shape[0]
        i = random.randint(0, length-self.min_audio_length)
        wav = wav[i:i+self.min_audio_length]

        wav_pad = F.pad(wav, (160, 160))
        feat = self.feature_extractor(wav_pad, sampling_rate=16000, return_tensors="pt") .data['input_features']
        out = {

            'wav': wav,
            'feat': feat,
            # 'paths': wavpath_full
        }

        return out

    def collate_fn(self, bs):

        wavs = [b['wav'] for b in bs]
        wavs = torch.stack(wavs)
        feats = [b['feat'] for b in bs]
        feats = torch.stack(feats)
        out = {

            'wav': wavs,
            'feats': feats,
            # 'paths': [b['paths'] for b in bs]
        }
        return out

@hydra.main(config_path='config', config_name='default', version_base=None)
def main(cfg):

    data_module = DataModule(cfg)

    train_loader = data_module.val_dataloader()

    valid_filelist = []

    for batch_idx, batch in enumerate(tqdm(train_loader, desc="Processing batches", unit="batch")):

        wavs = batch['wav']

if __name__ == "__main__":
    main()


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.


**Confirming training expects local filelists**

The current training pipeline uses filelists of WAV paths relative to root and doesn't employ tar or webdataset formats. To support .tar archives, either validation of the tar or adapting the training dataloader to handle webdatasets will be necessary. Next steps include checking dataset config files and related scripts for how tar files might be integrated or supported.


sed -n '1,200p' X-Codec-2.0/config/dataset/default.yaml


Output:
Exit code: 0
Wall time: 0 seconds
Output:
dataset:
  _target_: data_module.FSDataset_add_STFT

train:
  filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_higher_quality.txt
  batch_size: 8
  shuffle: true

val:
  filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_test.txt
  batch_size:  8
  shuffle: false

test:
  filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_test.txt
  batch_size: 1
  shuffle: false
padding_idx: 0
frame_length: 200
min_audio_length: 96000


Summary: auto


**Searching for alternative dataset classes**

I'm checking if there are other dataset classes like FSDataset_add_STFT that handle dataset file lists differently, since the current one is not for wds.


rg -n "class FSDataset_add_STFT|add_STFT" -S X-Codec-2.0


Output:
Exit code: 0
Wall time: 0 seconds
Output:
X-Codec-2.0/data_module.py:33:        # ds = FSDataset_add_STFT(phase, self.cfg)
X-Codec-2.0/config/dataset/default.yaml:2:  _target_: data_module.FSDataset_add_STFT


**Noticing missing dataset class and input format mismatch**

The dataset config references a missing dataset class, indicating inconsistency; training might not use webdataset tar inputs despite pipeline outputting tar shards. Next, I’ll explore other code paths to see if alternative training code handles webdataset or whether I must adapt the training to read from extracted files or modify it for webdataset support.


rg -n "WebDataset|tar|jsonl|manifest" -S


Output:
Exit code: 0
Wall time: 0 seconds
Total output lines: 654
Output:
scripts/fetch_languages.py:19:        if not stripped or stripped.startswith("#") or "=" not in stripped:
scripts/fetch_languages.py:21:        if stripped.startswith("export "):
objective.md:9:WebDataset shards: shard-000000.tar, shard-000001.tar, …
objective.md:10:Each sample inside the tar:
objective.md:12:sample_id.json (per-sample metadata) OR a global manifest.jsonl indexed by sample_id
objective.md:13:Why WebDataset:
objective.md:18:0.5–2 GB per tar
objective.md:25:      shard-000000.tar
objective.md:26:      shard-000001.tar
objective.md:29:      shard-000000.tar
objective.md:31:  manifests/
objective.md:32:    train.jsonl
objective.md:33:    val.jsonl
objective.md:38:2) What to put in your manifest (metadata before training)
objective.md:39:Create a train.jsonl where each line is one sample:
objective.md:42:shard (which tar)
objective.md:43:offset (optional if you do indexed tar reads; not required for pure streaming)
objective.md:105:• Pick training SR: 16k or 24k (start here even if you’ll do 48k later)
objective.md:109:target SR
objective.md:113:target 3–12 seconds
objective.md:116:Write a single manifest_raw.jsonl
objective.md:128:pack into tar shards of 0.5–2GB
objective.md:129:write final train.jsonl with shard paths
objective.md:168: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:
objective.md:177:Use WebDataset tar shards (streaming-friendly)
objective.md:178:r2://bucket/xcodec16/train/shard-000000.tar
objective.md:179:r2://bucket/xcodec16/val/shard-000000.tar
objective.md:181:512MB – 1GB per tar is a sweet spot on S3/R2
objective.md:187:uid.json (optional; can also use a global manifest)
objective.md:201:Segmentation target (recommended):
objective.md:205:Make a train.jsonl with one line per sample:
objective.md:247:speech_segments: list of [start, end] times from VAD
objective.md:292:Store 8–15s VAD segments (not 6s fixed) in 512MB–1GB WebDataset tar shards on R2
objective.md:298:Building train.jsonl
objective.md:299:Writing WebDataset shards directly to R2
X-Codec-2.0/requirements.txt:99:rotary-embedding-torch==0.8.4
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),
scripts/normalize_languages.py:17:        if not stripped or stripped.startswith("#") or "=" not in stripped:
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)
scripts/bulk_download.py:1:"""Bulk download all tars for specified languages from R2"""
scripts/bulk_download.py:19:        if not stripped or stripped.startswith("#") or "=" not in stripped:
scripts/bulk_download.py:21:        if stripped.startswith("export "):
scripts/bulk_download.py:45:    Returns: {language: [(youtube_id, r2_tar_key), ...]}
scripts/bulk_download.py:52:                    SELECT youtube_id, r2_tar_key
scripts/bulk_download.py:55:                      AND r2_tar_key IS NOT NULL
scripts/bulk_download.py:65:    """Try to find the tar file with prefix fallback"""
scripts/bulk_download.py:74:        if not key.startswith(cleaned_prefix):
scripts/bulk_download.py:101:    """Download a single video tar
scripts/bulk_download.py:110:        local_path = output_dir / f"{youtube_id}.tar"
scripts/bulk_download.py:121:    parser = argparse.ArgumentParser(description="Bulk download tars for languages")
scripts/bulk_download.py:177:    print(f"Starting downloads with {args.workers} workers...\n")
scripts/fill_missing_languages.py:22:        if not stripped or stripped.startswith("#") or "=" not in stripped:
scripts/fill_missing_languages.py:164:        start = time.time()
scripts/fill_missing_languages.py:173:            for chunk_start in range(0, total, chunk_size):
scripts/fill_missing_languages.py:174:                chunk = video_ids[chunk_start : chunk_start + chunk_size]
scripts/fill_missing_languages.py:196:                        elapsed = time.time() - start
scripts/fill_missing_languages.py:207:    elapsed = time.time() - start
strategy.md:6:- **Format**: `videoID.tar` files containing `metadata.json` + `segments/*.flac` (original quality from YouTube WebM, compressed as FLAC)
strategy.md:55:For each segment in each `videoID.tar`, compute and store:
strategy.md:79:- `start`, `end`, `original_start`, `original_end`, `duration`
strategy.md:84:1. Load the segment FLAC from the tar (already at original quality from YouTube WebM).
strategy.md:86:3. Assign buckets and resample targets based on **actual bandwidth**, not container SR:
strategy.md:91:4. Record `resample_target_sr` for each emitted sample.
strategy.md:125:- **Silence-aware chunking**: Only cut when a low-energy boundary exists near the target window; otherwise keep the full segment.
strategy.md:133:├── manifests/
strategy.md:134:│   ├── telugu_train.jsonl
strategy.md:135:│   ├── telugu_val.jsonl
strategy.md:138:│   ├── shard-000000.tar
strategy.md:139:│   ├── shard-000000.jsonl
strategy.md:142:│   ├── shard-000000.tar
strategy.md:143:│   ├── shard-000000.jsonl
strategy.md:146:│   ├── shard-000000.tar
strategy.md:147:│   ├── shard-000000.jsonl
strategy.md:150:│   ├── shard-000000.tar
strategy.md:151:│   ├── shard-000000.jsonl
strategy.md:166:shard-000000.tar
strategy.md:167:shard-000000.jsonl
strategy.md:170:`shard-000000.jsonl` is one line per sample with:
strategy.md:178:  "resample_target_sr": 16000,
strategy.md:280:   - `train_codes.npy`, `train_codes_index.npy`, `train_samples.jsonl`
strategy.md:330:- Progressive training: start with 32k decoder, then fine-tune to 44k
strategy.md:351:| 4 | Create sharding pipeline | WebDataset tars + JSONL |
strategy.md:352:| 4 | Generate train/val manifests | JSONL files |
strategy.md:362:| `scripts/create_shards.py` | Pack segments into WebDataset tars + JSONL |
strategy.md:363:| `scripts/generate_manifests.py` | Create train/val JSONL files |
strategy.md:421:- `start`, `end`, `duration`, `original_start`, `original_end`
strategy.md:431:- `start`, `end`, `original_start`, `original_end` (time boundaries)
strategy.md:437:- `resample_target_sr` (16000, 24000, 32000, or 44100)
scripts/supabase_summary.py:58:        if not stripped or stripped.startswith("#") or "=" not in stripped:
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)
scripts/prepare_shards.py:9:import tarfile
scripts/prepare_shards.py:61:    start = time.perf_counter()
scripts/prepare_shards.py:65:        elapsed = time.perf_counter() - start
scripts/prepare_shards.py:134:    r2_tar_key: str
scripts/prepare_shards.py:167:        if not stripped or stripped.startswith("#") or "=" not in stripped:
scripts/prepare_shards.py:169:        if stripped.startswith("export "):
scripts/prepare_shards.py:200:        if not key.startswith(cleaned_prefix):
scripts/prepare_shards.py:222:    if cleaned.startswith(default_bucket + "/"):
scripts/prepare_shards.py:238:                select youtube_id, r2_tar_key
scripts/prepare_shards.py:241:                  and r2_tar_key is not null
scripts/prepare_shards.py:253:            for youtube_id, r2_tar_key in cur.fetchall():
scripts/prepare_shards.py:254:                if youtube_id and r2_tar_key:
scripts/prepare_shards.py:255:                    records.append(VideoRecord(youtube_id, r2_tar_key))
scripts/prepare_shards.py:416:    search_start: float,
scripts/prepare_shards.py:424:    mask = (times >= search_start) & (times <= search_end)
scripts/prepare_shards.py:479:        target_len = rng.uniform(cfg.min_sec, cfg.max_sec)
scripts/prepare_shards.py:480:        if remaining - target_len < cfg.min_sec:
scripts/prepare_shards.py:481:            target_len = remaining - cfg.min_sec
scripts/prepare_shards.py:482:        target_end = offset + target_len
scripts/prepare_shards.py:483:        search_low = max(offset + cfg.min_sec, target_end - silence_cfg.search_sec)
scripts/prepare_shards.py:484:        search_high = min(offset + cfg.max_sec, target_end + silence_cfg.search_sec, latest_cut)
scripts/prepare_shards.py:515:    """Parse segment filename to extract speaker and start time.
scripts/prepare_shards.py:517:    Filename format: segments/SPEAKER_XX_NNNN_start-end.flac
scripts/prepare_shards.py:520:    Returns (speaker, start_time) for matching against metadata.
scripts/prepare_shards.py:543:        start_str = time_range.split("-")[0]
scripts/prepare_shards.py:544:        start_time = float(start_str)
scripts/prepare_shards.py:548:    return speaker, start_time
scripts/prepare_shards.py:592:        start: Optional[float],
scripts/prepare_shards.py:594:        original_start: Optional[float],
scripts/prepare_shards.py:596:        target_srs: list[int],
scripts/prepare_shards.py:621:            "start": start,
scripts/prepare_shards.py:623:            "original_start": original_start,
scripts/prepare_shards.py:646:        for target_sr in target_srs:
scripts/prepare_shards.py:647:            resampled = resample_audio(audio_slice, sr, target_sr)
scripts/prepare_shards.py:649:            audio_bytes = encode_flac(resampled, target_sr)
scripts/prepare_shards.py:651:            meta["sample_rate"] = target_sr
scripts/prepare_shards.py:652:            meta["resample_target_sr"] = target_sr
scripts/prepare_shards.py:653:            outputs.append(SegmentOutput(target_sr, audio_bytes, meta))
scripts/prepare_shards.py:655:        if target_srs == [16000]:
scripts/prepare_shards.py:663:    segment_start = segment_meta.get("start")
scripts/prepare_shards.py:665:    original_start = segment_meta.get("original_start")
scripts/prepare_shards.py:670:        start_idx = int(round(offset_sec * sr))
scripts/prepare_shards.py:672:        audio_slice = audio[start_idx:end_idx]
scripts/prepare_shards.py:678:        chunk_start = float(segment_start) + offset_sec if segment_start is not None else None
scripts/prepare_shards.py:679:        chunk_end = chunk_start + slice_duration if chunk_start is not None else None
scripts/prepare_shards.py:680:        chunk_original_start = (
scripts/prepare_shards.py:681:            float(original_start) + offset_sec if original_start is not None else None
scripts/prepare_shards.py:684:            chunk_original_start + slice_duration
scripts/prepare_shards.py:685:            if chunk_original_start is not None
scripts/prepare_shards.py:701:            chunk_start,
scripts/prepare_shards.py:703:            chunk_original_start,
scripts/prepare_shards.py:711:    targets_high: list[int] = []
scripts/prepare_shards.py:713:        targets_high.append(24000)
scripts/prepare_shards.py:715:        targets_high.append(32000)
scripts/prepare_shards.py:717:        targets_high.append(44100)
scripts/prepare_shards.py:719:    if targets_high:
scripts/prepare_shards.py:722:                start_idx = int(round(offset_sec * sr))
scripts/prepare_shards.py:724:                audio_slice = audio[start_idx:end_idx]
scripts/prepare_shards.py:730:                slice_targets: list[int] = []
scripts/prepare_shards.py:732:                    slice_targets.append(24000)
scripts/prepare_shards.py:734:                    slice_targets.append(32000)
scripts/prepare_shards.py:736:                    slice_targets.append(44100)
scripts/prepare_shards.py:737:                if not slice_targets:
scripts/prepare_shards.py:740:                chunk_start = (
scripts/prepare_shards.py:741:                    float(segment_start) + offset_sec if segment_start is not None else None
scripts/prepare_shards.py:744:                    chunk_start + slice_duration if chunk_start is not None else None
scripts/prepare_shards.py:746:                chunk_original_start = (
scripts/prepare_shards.py:747:                    float(original_start) + offset_sec if original_start is not None else None
scripts/prepare_shards.py:750:                    chunk_original_start + slice_duration
scripts/prepare_shards.py:751:                    if chunk_original_start is not None
scripts/prepare_shards.py:762:                    chunk_start,
scripts/prepare_shards.py:764:                    chunk_original_start,
scripts/prepare_shards.py:766:                    slice_targets,
scripts/prepare_shards.py:775:                segment_start,
scripts/prepare_shards.py:777:                original_start,
scripts/prepare_shards.py:779:                targets_high,
scripts/prepare_shards.py:786:    def __init__(self, output_dir: Path, prefix: str, target_size_bytes: int) -> None:
scripts/prepare_shards.py:789:        self.target_size_bytes = target_size_bytes
scripts/prepare_shards.py:791:        self.current_tar: Optional[tarfile.TarFile] = None
scripts/prepare_shards.py:792:        self.current_jsonl = None
scripts/prepare_shards.py:799:        if self.current_tar:
scripts/prepare_shards.py:802:        tar_path = self.output_dir / f"{shard_name}.tar"
scripts/prepare_shards.py:803:        jsonl_path = self.output_dir / f"{shard_name}.jsonl"
scripts/prepare_shards.py:804:        self.current_tar = tarfile.open(tar_path, "w")
scripts/prepare_shards.py:805:        self.current_jsonl = jsonl_path.open("w", encoding="utf-8")
scripts/prepare_shards.py:812:        if self.current_tar is None:
scripts/prepare_shards.py:816:            and self.current_size + len(audio_bytes) > self.target_size_bytes
scripts/prepare_shards.py:820:        assert self.current_tar is not None
scripts/prepare_shards.py:821:        assert self.current_jsonl is not None
scripts/prepare_shards.py:824:        info = tarfile.TarInfo(name=arcname)
scripts/prepare_shards.py:827:        self.current_tar.addfile(info, io.BytesIO(audio_bytes))
scripts/prepare_shards.py:828:        self.current_jsonl.write(json.dumps(metadata) + "\n")
scripts/prepare_shards.py:835:        if self.current_tar:
scripts/prepare_shards.py:836:            self.current_tar.close()
scripts/prepare_shards.py:837:        if self.current_jsonl:
scripts/prepare_shards.py:838:            self.current_jsonl.close()
scripts/prepare_shards.py:839:        self.current_tar = None
scripts/prepare_shards.py:840:        self.current_jsonl = None
scripts/prepare_shards.py:901:    tar_key = f"{remote_prefix}/{shard_name}.tar"
scripts/prepare_shards.py:902:    jsonl_key = f"{remote_prefix}/{shard_name}.jsonl"
scripts/prepare_shards.py:903:    tar_path = tmp_dir / f"{shard_name}.tar"
scripts/prepare_shards.py:904:    jsonl_path = tmp_dir / f"{shard_name}.jsonl"
scripts/prepare_shards.py:906:    s3.download_file(bucket, tar_key, str(tar_path))
scripts/prepare_shards.py:907:    s3.download_file(bucket, jsonl_key, str(jsonl_path))
scripts/prepare_shards.py:909:    with tarfile.open(tar_path, "r") as tf:
scripts/prepare_shards.py:910:        tar_names = set(tf.getnames())
scripts/prepare_shards.py:911:        with jsonl_path.open("r", encoding="utf-8") as fh:
scripts/prepare_shards.py:919:            if audio_name not in tar_names:
scripts/prepare_shards.py:936:    parser = argparse.ArgumentParser(description="Prepare sharded datasets from cleaned tars.")
scripts/prepare_shards.py:941:    parser.add_argument("--target-shard-size-mb", type=int, default=5120)
scripts/prepare_shards.py:953:        help="Use only local tar files; do not download from R2.",
scripts/prepare_shards.py:956:        "--local-tars-dir",
scripts/prepare_shards.py:959:        help="Directory containing pre-downloaded tar files for this language.",
scripts/prepare_shards.py:1101:            ["manifests", "stats", "reports", "checks"],
scripts/prepare_shards.py:1124:        resolved = resolve_r2_key(s3, cleaned_bucket, record.r2_tar_key, prefix)
scripts/prepare_shards.py:1126:            raise FileNotFoundError(f"Missing tar in R2: {record.r2_tar_key}")
scripts/prepare_shards.py:1127:        local_path = download_dir / f"{record.youtube_id}.tar"
scripts/prepare_shards.py:1134:        local_dir = args.local_tars_dir or download_dir
scripts/prepare_shards.py:1136:            raise SystemExit(f"Local tar directory not found: {local_dir}")
scripts/prepare_shards.py:1139:            local_path = local_dir / f"{record.youtube_id}.tar"
scripts/prepare_shards.py:1145:            print(f"Missing local tars: {missing} (skipping)")
scripts/prepare_shards.py:1167:    manifest_dir = output_dir / "manifests"
scripts/prepare_shards.py:1169:    manifest_dir.mkdir(parents=True, exist_ok=True)
scripts/prepare_shards.py:1174:        16000: ShardWriter(output_dir / f"{language}_16", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
scripts/prepare_shards.py:1175:        24000: ShardWriter(output_dir / f"{language}_24", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
scripts/prepare_shards.py:1176:        32000: ShardWriter(output_dir / f"{language}_32", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
scripts/prepare_shards.py:1177:        44100: ShardWriter(output_dir / f"{language}_44", shard_prefix, args.target_shard_size_mb * 1024 * 1024),
scripts/prepare_shards.py:1180:    manifest_suffix = f"_{partition_tag}" if partition_tag else ""
scripts/prepare_shards.py:1181:    train_manifest = (manifest_dir / f"{language}_trai…13181 tokens truncated…packages/e1/41/9b873a8c055582859b239be17902a85339bec6a30ad162f98c9b0288a2cc/soundfile-0.13.1.tar.gz", hash = "sha256:b2c68dab1e30297317080a5b43df57e302584c49e2942defdde0acccc53f0e5b", size = 46156 }
tts/uv.lock:4035:sdist = { url = "https://files.pythonhosted.org/packages/02/c0/4429bf9b3be10e749149e286aa5c53775399ec62891c6b970456c6dca325/soxr-0.5.0.post1.tar.gz", hash = "sha256:7092b9f3e8a416044e1fa138c8172520757179763b85dc53aa9504f4813cff73", size = 170853 }
tts/uv.lock:4062:sdist = { url = "https://files.pythonhosted.org/packages/c4/53/6050dc3dde1671eb3db592c13b55a8005e5040131f7509cef0215212cb84/standard_aifc-3.13.0.tar.gz", hash = "sha256:64e249c7cb4b3daf2fdba4e95721f811bde8bdfc43ad9f936589b7bb2fae2e43", size = 15240 }
tts/uv.lock:4071:sdist = { url = "https://files.pythonhosted.org/packages/43/06/ce1bb165c1f111c7d23a1ad17204d67224baa69725bb6857a264db61beaf/standard_chunk-3.13.0.tar.gz", hash = "sha256:4ac345d37d7e686d2755e01836b8d98eda0d1a3ee90375e597ae43aaf064d654", size = 4672 }
tts/uv.lock:4083:sdist = { url = "https://files.pythonhosted.org/packages/66/e3/ce8d38cb2d70e05ffeddc28bb09bad77cfef979eb0a299c9117f7ed4e6a9/standard_sunau-3.13.0.tar.gz", hash = "sha256:b319a1ac95a09a2378a8442f403c66f4fd4b36616d6df6ae82b8e536ee790908", size = 9368 }
tts/uv.lock:4089:name = "starlette"
tts/uv.lock:4096:sdist = { url = "https://files.pythonhosted.org/packages/0a/69/662169fdb92fb96ec3eaee218cf540a629d629c86d7993d9651226a6789b/starlette-0.47.1.tar.gz", hash = "sha256:aef012dd2b6be325ffa16698f9dc533614fb1cebd593a906b90dc1025529a79b", size = 2583072 }
tts/uv.lock:4098:    { url = "https://files.pythonhosted.org/packages/82/95/38ef0cd7fa11eaba6a99b3c4f5ac948d8bc6ff199aabd327a29cc000840c/starlette-0.47.1-py3-none-any.whl", hash = "sha256:5e11c9f5c7c3f24959edbf2dffdc01bba860228acf657129467d8a7468591527", size = 72747 },
tts/uv.lock:4114:sdist = { url = "https://files.pythonhosted.org/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040 }
tts/uv.lock:4136:sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921 }
tts/uv.lock:4145:sdist = { url = "https://files.pythonhosted.org/packages/a3/4d/6a19536c50b849338fcbe9290d562b52cbdcf30d8963d3588a68a4107df1/tenacity-8.5.0.tar.gz", hash = "sha256:8bc6c0c8a09b31e6cad13c47afbed1a567518250a9a171418582ed8d9c20ca78", size = 47309 }
tts/uv.lock:4154:sdist = { url = "https://files.pythonhosted.org/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274 }
tts/uv.lock:4167:sdist = { url = "https://files.pythonhosted.org/packages/ea/cf/756fedf6981e82897f2d570dd25fa597eb3f4459068ae0572d7e888cfd6f/tiktoken-0.9.0.tar.gz", hash = "sha256:d02a5ca6a938e0490e1ff957bc48c8b078c88cb83977be1625b1fd8aac792c5d", size = 35991 }
tts/uv.lock:4209:sdist = { url = "https://files.pythonhosted.org/packages/80/4c/60088c867d4750a085cff05a1b4f65e0e6b38fba973a72108d23decb4a4c/timm-1.0.17.tar.gz", hash = "sha256:90ccf07894d6ae38259b7aa7c94d683002f5e360fd43471c84f9da6c3afb8a0d", size = 2345699 }
tts/uv.lock:4221:sdist = { url = "https://files.pythonhosted.org/packages/ab/2d/b0fce2b8201635f60e8c95990080f58461cc9ca3d5026de2e900f38a7f21/tokenizers-0.21.2.tar.gz", hash = "sha256:fdc7cffde3e2113ba0e6cc7318c40e3438a4d74bbc62bf04bcc63bdfb082ac77", size = 351545 }
tts/uv.lock:4243:sdist = { url = "https://files.pythonhosted.org/packages/18/87/302344fed471e44a87289cf4967697d07e532f2421fdaf868a303cbae4ff/tomli-2.2.1.tar.gz", hash = "sha256:cd45e1dc79c835ce60f7404ec8119f2eb06d38b1deba146f07ced3bbc44505ff", size = 17175 }
tts/uv.lock:4282:sdist = { url = "https://files.pythonhosted.org/packages/19/75/241269d1da26b624c0d5e110e8149093c759b7a286138f4efd61a60e75fe/tomli_w-1.2.0.tar.gz", hash = "sha256:2dd14fac5a47c27be9cd4c976af5a12d87fb1f0b4512f81d69cce3b35ae25021", size = 7184 }
tts/uv.lock:4291:sdist = { url = "https://files.pythonhosted.org/packages/cc/18/0bbf3884e9eaa38819ebe46a7bd25dcd56b67434402b66a58c4b8e552575/tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1", size = 185207 }
tts/uv.lock:4488:sdist = { url = "https://files.pythonhosted.org/packages/a4/67/e4c79da3f15777b9bc2b655d47dac553fc31e40360500fef5e66d6877ce8/torchlibrosa-0.1.0.tar.gz", hash = "sha256:62a8beedf9c9b4141a06234df3f10229f7ba86e67678ccee02489ec4ef044028", size = 11719 }
tts/uv.lock:4505:sdist = { url = "https://files.pythonhosted.org/packages/96/80/bcf680d3d7de7646707884ed7ce0859c83be33380239971a97255147184f/torchmetrics-1.7.4.tar.gz", hash = "sha256:506a1a5c7c304cd77ba323ca4b009e46b814fd2be9dcf0f4ccc2e5c0f5b4b0c1", size = 567802 }
tts/uv.lock:4604:sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
tts/uv.lock:4625:sdist = { url = "https://files.pythonhosted.org/packages/4c/67/80f51466ec447028fd84469b208eb742533ce06cc8fad2e3181380199e5c/transformers-4.53.2.tar.gz", hash = "sha256:6c3ed95edfb1cba71c4245758f1b4878c93bf8cde77d076307dacb2cbbd72be2", size = 9201233 }
tts/uv.lock:4651:sdist = { url = "https://files.pythonhosted.org/packages/05/8c/cd2aae1c77b9180cf589e44def6f40c87e6b31603cdbea84c3d6c6825277/trl-0.16.0.tar.gz", hash = "sha256:0e76c6f697bd93c27fd1221ff18646413cf601ba543afef32c2fd87ecaabce4e", size = 351764 }
tts/uv.lock:4660:sdist = { url = "https://files.pythonhosted.org/packages/38/04/1cd43f72c241fedcf0d9a18d0783953ee301eac9e5d9db1df0f0f089d9af/trove_classifiers-2025.5.9.12.tar.gz", hash = "sha256:7ca7c8a7a76e2cd314468c677c69d12cc2357711fcab4a60f87994c1589e5cb5", size = 16940 }
tts/uv.lock:4769:sdist = { url = "https://files.pythonhosted.org/packages/c7/68/71c1a15b5f65f40e91b65da23b8224dad41349894535a97f63a52e462196/typeguard-4.4.4.tar.gz", hash = "sha256:3a7fd2dffb705d4d0efaed4306a704c89b9dee850b688f060a8b1615a79e5f74", size = 75203 }
tts/uv.lock:4784:sdist = { url = "https://files.pythonhosted.org/packages/c5/8c/7d682431efca5fd290017663ea4588bf6f2c6aad085c7f108c5dbc316e70/typer-0.16.0.tar.gz", hash = "sha256:af377ffaee1dbe37ae9440cb4e8f11686ea5ce4e9bae01b84ae7c63b87f1dd3b", size = 102625 }
tts/uv.lock:4793:sdist = { url = "https://files.pythonhosted.org/packages/98/5a/da40306b885cc8c09109dc2e1abd358d5684b1425678151cdaed4731c822/typing_extensions-4.14.1.tar.gz", hash = "sha256:38b39f4aeeab64884ce9f74c94263ef78f3c22467c8724005483154c26648d36", size = 107673 }
tts/uv.lock:4805:sdist = { url = "https://files.pythonhosted.org/packages/f8/b1/0c11f5058406b3af7609f121aaa6b609744687f1d158b3c3a5bf4cc94238/typing_inspection-0.4.1.tar.gz", hash = "sha256:6ae134cc0203c33377d43188d4064e9b357dba58cff3185f22924610e70a9d28", size = 75726 }
tts/uv.lock:4814:sdist = { url = "https://files.pythonhosted.org/packages/95/32/1a225d6164441be760d75c2c42e2780dc0873fe382da3e98a2e1e48361e5/tzdata-2025.2.tar.gz", hash = "sha256:b60a638fcc0daffadf82fe0f57e53d06bdec2f36c4df66280ae79bce6bd6f2b9", size = 196380 }
tts/uv.lock:4823:sdist = { url = "https://files.pythonhosted.org/packages/94/7d/a8a765761bbc0c836e397a2e48d498305a865b70a8600fd7a942e85dcf63/Unidecode-1.4.0.tar.gz", hash = "sha256:ce35985008338b676573023acc382d62c264f307c8f7963733405add37ea2b23", size = 200149 }
tts/uv.lock:4832:sdist = { url = "https://files.pythonhosted.org/packages/15/22/9ee70a2574a4f4599c47dd506532914ce044817c7752a79b6a51286319bc/urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760", size = 393185 }
tts/uv.lock:4844:sdist = { url = "https://files.pythonhosted.org/packages/d5/b7/30753098208505d7ff9be5b3a32112fb8a4cb3ddfccbbb7ba9973f2e29ff/userpath-1.9.2.tar.gz", hash = "sha256:6c52288dab069257cc831846d15d48133522455d4677ee69a9781f11dbefd815", size = 11140 }
tts/uv.lock:4853:sdist = { url = "https://files.pythonhosted.org/packages/43/c5/d1defc44eea72552ee7667212359bded10ac8f7e39b62b042a3d4a9d4040/uv-0.7.22.tar.gz", hash = "sha256:f5cf159907d594e33433f14737d1ee843dc8799edfcf57b5b8c0f282d1117051", size = 3387947 }
tts/uv.lock:4883:sdist = { url = "https://files.pythonhosted.org/packages/5e/42/e0e305207bb88c6b8d3061399c6a961ffe5fbb7e2aa63c9234df7259e9cd/uvicorn-0.35.0.tar.gz", hash = "sha256:bc662f087f7cf2ce11a1d7fd70b90c9f98ef2e2831556dd078d131b96cc94a01", size = 78473 }
tts/uv.lock:4899:sdist = { url = "https://files.pythonhosted.org/packages/57/58/2a8c3b2f6b13ee34aaa940e552109becf0148f3cccbdea22b7d3b2fcbdc3/vector_quantize_pytorch-1.17.8.tar.gz", hash = "sha256:3b9b9514dcdce7af122b7e8abae1a01b4050c3446b78b4b8a07da74f0008d3ea", size = 592912 }
tts/uv.lock:4913:sdist = { url = "https://files.pythonhosted.org/packages/56/2c/444f465fb2c65f40c3a104fd0c495184c4f2336d65baf398e3c75d72ea94/virtualenv-20.31.2.tar.gz", hash = "sha256:e10c0a9d02835e592521be48b332b6caee6887f332c111aa79a09b9e79efc2af", size = 6076316 }
tts/uv.lock:4934:sdist = { url = "https://files.pythonhosted.org/packages/73/09/c84264a219e20efd615e4d5d150cc7d359d57d51328d3fa94ee02d70ed9c/wandb-0.21.0.tar.gz", hash = "sha256:473e01ef200b59d780416062991effa7349a34e51425d4be5ff482af2dc39e02", size = 40085784 }
tts/uv.lock:4952:sdist = { url = "https://files.pythonhosted.org/packages/6c/63/53559446a878410fc5a5974feb13d31d78d752eb18aeba59c7fef1af7598/wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5", size = 101301 }
tts/uv.lock:4961:sdist = { url = "https://files.pythonhosted.org/packages/21/e6/26d09fab466b7ca9c7737474c52be4f76a40301b08362eb2dbc19dcc16c1/websockets-15.0.1.tar.gz", hash = "sha256:82544de02076bafba038ce055ee6412d68da13ab47f0c60cab827346de828dee", size = 177016 }
tts/uv.lock:5026:sdist = { url = "https://files.pythonhosted.org/packages/8a/98/2d9906746cdc6a6ef809ae6338005b3f21bb568bea3165cfc6a243fdc25c/wheel-0.45.1.tar.gz", hash = "sha256:661e1abd9198507b1409a20c02106d9670b2576e916d58f520316666abca6729", size = 107545 }
tts/uv.lock:5035:sdist = { url = "https://files.pythonhosted.org/packages/c3/fc/e91cc220803d7bc4db93fb02facd8461c37364151b8494762cc88b0fbcef/wrapt-1.17.2.tar.gz", hash = "sha256:41388e9d4d1522446fe79d3213196bd9e3b301a336965b9e27ca2788ebd122f3", size = 55531 }
tts/uv.lock:5074:sdist = { url = "https://files.pythonhosted.org/packages/00/5e/d6e5258d69df8b4ed8c83b6664f2b47d30d2dec551a29ad72a6c69eafd31/xxhash-3.5.0.tar.gz", hash = "sha256:84f2caddf951c9cbf8dc2e22a89d4ccf5d86391ac6418fe81e3c67d0cf60b45f", size = 84241 }
tts/uv.lock:5152:sdist = { url = "https://files.pythonhosted.org/packages/3c/fb/efaa23fa4e45537b827620f04cf8f3cd658b76642205162e072703a5b963/yarl-1.20.1.tar.gz", hash = "sha256:d017a4997ee50c91fd5466cef416231bb82177b93b029906cefc542ce14c35ac", size = 186428 }
tts/uv.lock:5246:sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547 }
tts/uv.lock:5258:sdist = { url = "https://files.pythonhosted.org/packages/ed/f6/2ac0287b442160a89d726b17a9184a4c615bb5237db763791a7fd16d9df1/zstandard-0.23.0.tar.gz", hash = "sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09", size = 681701 }
tts/tts/training/codec/codec_datasets.py:26:        samples_path = os.path.join(dataset_dir, f"{split}_samples.jsonl")
tts/tts/training/codec/codec_datasets.py:81:            self.samples[idx].wav_path, target_sample_rate=self.sample_rate
tts/tts/training/codec/train_codec.py:104:    logging.info("Starting data setup...")
tts/tts/training/codec/train_codec.py:213:    logging.info("Starting training...")
tts/tts/training/codec/gan_training_loop.py:36:    statistics.start_micro_batch_training()
tts/tts/training/codec/gan_training_loop.py:51:                    "Restarting the iterator."
tts/tts/training/rlhf/rlhf_main.py:73:    logging.info("Starting data setup. Max seq len: %d.", max_seq_len)
tts/tts/training/rlhf/rlhf_main.py:108:    logging.info("Starting training...")
tts/tts/data/datasets/rlhf.py:27:        self.number_of_codes = sum(end - start for start, end in self.indexes)
tts/tts/data/datasets/rlhf.py:35:        self.speech_start_id = tokenizer.convert_tokens_to_ids(
tts/tts/training/rlhf/rewards.py:116:        start_time = time.time()
tts/tts/training/rlhf/rewards.py:157:                f"time: {end_time - start_time:.2f} seconds."
tts/tts/training/rlhf/rewards.py:180:        start_time = time.time()
tts/tts/training/rlhf/rewards.py:198:                f"time: {end_time - start_time:.2f} seconds."
tts/tts/training/rlhf/rewards.py:236:        start_time = time.time()
tts/tts/training/rlhf/rewards.py:268:                f"{similarity_scores}, time: {end_time - start_time:.2f} seconds."
tts/tts/core/codec/decoder_modules.py:261:        rotary_embed: torchtune_module.RotaryPositionalEmbeddings,
tts/tts/core/codec/decoder_modules.py:270:        self.rotary_embed = rotary_embed
tts/tts/core/codec/decoder_modules.py:280:        q = self.rotary_embed(q)
tts/tts/core/codec/decoder_modules.py:281:        k = self.rotary_embed(k)
tts/tts/core/codec/decoder_modules.py:300:        rotary_embed: torchtune_module.RotaryPositionalEmbeddings,
tts/tts/core/codec/decoder_modules.py:308:        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed)
tts/tts/core/codec/decoder_modules.py:363:        time_rotary_embed = torchtune_module.RotaryPositionalEmbeddings(dim=pos_meb_dim)
tts/tts/core/codec/decoder_modules.py:367:                dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed
tts/tts/data/datasets/finetuning.py:31:        self.number_of_codes = sum(end - start for start, end in self.indexes)
tts/tts/data/datasets/finetuning.py:40:        self.speech_start_id = tokenizer.convert_tokens_to_ids(
tts/tts/data/datasets/finetuning.py:46:            "Pad/speech start/speed end tokens: [%d, %d, %d]",
tts/tts/data/datasets/finetuning.py:51:            self.speech_start_id,
tts/tts/data/datasets/finetuning.py:79:                (input_ids == self.speech_start_id).nonzero(as_tuple=True)[0].item()
tts/tts/data/datasets/finetuning.py:82:            logging.warning("Speech start not found in the input ids.")
tts/tts/data/datasets/finetuning.py:130:        if not raw_text.startswith("<human>:"):
tts/tts/data/datasets/finetuning.py:131:            raise ValueError("Sample does not start with <human>:")
tts/tts/data/datasets/finetuning.py:169:        response_start_idx = min(right_most_end_header_id + 1, self._max_seq_len - 1)
tts/tts/data/datasets/finetuning.py:175:        labels[:response_start_idx] = constants.LOSS_IGNORE_TOKEN_ID
tts/tts/core/constants.py:6:SPEECH_START_TOKEN = "<|speech_start|>"
tts/tts/core/constants.py:8:TEXT_PROMPT_START_TOKEN = "<|text_prompt_start|>"
tts/tts/core/constants.py:10:VOICE_DESCRIPTION_START_TOKEN = "<|voice_description_start|>"
tts/tts/core/constants.py:12:SOUND_EFFECT_START_TOKEN = "<|sound_effect_start|>"
tts/tts/training/environment.py:208:        logging.info("Start syncing distributed SLURM processes via fabric...")
tts/tts/core/modeling.py:86:    start_time = time.time()
tts/tts/core/modeling.py:91:        time.time() - start_time,
tts/tts/core/modeling.py:105:    start_time = time.time()
tts/tts/core/modeling.py:113:    logging.info("Model loaded in %.2f seconds.", time.time() - start_time)
tts/tts/core/codec/decoder.py:35:                f"factors {self.upsample_factors} do not match the target "
tts/tts/core/codec/decoder.py:102:                if key.startswith("generator."):
tts/tts/core/codec/decoder.py:105:                elif key.startswith("fc_post_a."):
tts/tts/core/codec/decoder.py:117:                if k.startswith("generator.")
tts/tts/data/data_utils.py:24:    start_idx = worker_id * chunk_size
tts/tts/data/data_utils.py:25:    end_idx = total if (worker_id == num_workers - 1) else (start_idx + chunk_size)
tts/tts/data/data_utils.py:26:    chunk = work_items[start_idx:end_idx]
tts/tts/data/data_utils.py:30:        start_idx,
tts/tts/data/data_utils.py:48:    if not dataset_path.endswith(".jsonl"):
tts/tts/data/data_utils.py:49:        raise ValueError("Dataset path must end with .jsonl")
tts/tts/data/data_utils.py:53:    start_time = time.time()
tts/tts/data/data_utils.py:54:    logging.info("Starting to load dataset %s", dataset_path)
tts/tts/data/data_utils.py:91:        time.time() - start_time,
tts/tts/data/data_utils.py:103:    samples_path = os.path.join(dataset_dir, f"{split}_samples.jsonl")
tts/tts/data/data_utils.py:171:    wav_path: str, target_sample_rate: int | None = None
tts/tts/data/data_utils.py:177:        target_sample_rate: Optional target sample rate to resample to. If None,
tts/tts/data/data_utils.py:186:    if target_sample_rate and sr != target_sample_rate:
tts/tts/data/data_utils.py:187:        wav = torchaudio.functional.resample(wav, sr, target_sample_rate)
tts/tts/data/data_utils.py:188:        sr = target_sample_rate
tts/tts/training/checkpointing.py:60:                if f.startswith("checkpoint_") and f.endswith(".pt")
tts/tts/core/lora.py:33:    if config.target_modules:
tts/tts/core/lora.py:34:        target_modules = list(config.target_modules)
tts/tts/core/lora.py:36:        # When target_modules aren't provided, automatically find all linear components.
tts/tts/core/lora.py:37:        target_modules = _find_linear_modules(model)
tts/tts/core/lora.py:38:        logging.info("No |target_modules| provided, using all linear layers instead.")
tts/tts/core/lora.py:43:        target_modules=target_modules,
tts/tts/inference/quality_validation.py:168:            prompt_wav_path, target_sample_rate=constants.CODEC_SAMPLE_RATE
tts/tts/inference/quality_validation.py:199:            "Starting to synthesize test phrases for step %d. They will be saved in %s",
tts/tts/inference/quality_validation.py:270:            "Starting prompt continuation validation for step %d. "
tts/tts/inference/quality_validation.py:279:                prompt_wav_path, target_sample_rate=constants.CODEC_SAMPLE_RATE
tts/tts/training/main.py:112:    logging.info("Starting data setup. Max seq len: %d.", max_seq_len)
tts/tts/training/main.py:224:    logging.info("Starting training...")
tts/tts/core/optimization.py:30:        self._start_lr = learning_rate / 10.0
tts/tts/core/optimization.py:42:            return self._start_lr
tts/tts/core/optimization.py:52:        return self._start_lr + coeff * (self._peak_lr - self._start_lr)
tts/tts/inference/inferencing.py:57:        if token_str.startswith("<|s_") and token_str.endswith("|>"):
tts/tts/inference/inferencing.py:242:    speech_start_id = tokenizer.convert_tokens_to_ids(constants.SPEECH_START_TOKEN)
tts/tts/inference/inferencing.py:253:    input_ids = [speech_start_id] + input_ids
tts/tts/inference/inferencing.py:264:    # Remove the speech start and end tokens.
tts/tts/core/codec/encoder.py:95:                if key.startswith("CodecEnc."):
tts/tts/core/codec/encoder.py:98:                elif key.startswith("generator.quantizer."):
tts/tts/core/codec/encoder.py:101:                elif key.startswith("SemanticEncoder_module."):
tts/tts/core/codec/encoder.py:104:                elif key.startswith("fc_prior."):
tts/tts/training/training_loop.py:115:    statistics.start_micro_batch_training()
tts/tts/training/training_loop.py:129:                    "Restarting the iterator."
tts/tts/training/training_loop.py:226:                logging.info("Eval step %d: %s", statistics.step, "Starting evaluation")
tts/tts/training/training_loop.py:309:            logging.info("Step [%d]: starting checkpointing...", statistics.step)


**Inspecting TTS training datasets**

I'm checking the codec_datasets.py in tts/training to understand the dataset format used for TTS training, as there seems to be no dataset for 'tar' but possibly jsonl-based data.


sed -n '1,200p' tts/tts/training/codec/codec_datasets.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
import json
import os
import random
from collections.abc import Callable
from typing import Any

import numpy as np
import torch
from absl import logging

from tts.data import data_sample, data_utils, tts_datasets

class CodecTrainingDataset(torch.utils.data.Dataset):
    """Implements a TTS fine-tuning dataset."""

    def __init__(
        self,
        dataset_dir: str,
        split: str,
        audio_window_size: int,
        sample_rate: int,
        minimum_data_sample_rate: int,
    ):
        dataset_name = os.path.basename(dataset_dir) + "_" + split
        samples_path = os.path.join(dataset_dir, f"{split}_samples.jsonl")
        self.samples = []
        self.sample_rate = sample_rate
        self.audio_window_size = audio_window_size
        self.minimum_data_sample_rate = minimum_data_sample_rate
        self.code_window_size = int(self.audio_window_size / self.sample_rate * 50)
        self.hop_length = self.sample_rate // 50
        self.half_hop_length = self.hop_length // 2

        self.filtered_samples_index = []
        self.original_length = 0
        with open(samples_path) as f:
            for i, line in enumerate(f):
                sp = data_sample.Sample.from_json(json.loads(line), dataset_dir)
                self.original_length += 1
                if sp.sample_rate >= self.minimum_data_sample_rate:
                    self.samples.append(sp)
                    self.filtered_samples_index.append(i)

        codes_path = os.path.join(dataset_dir, f"{split}_codes.npy")
        codes_index_path = os.path.join(dataset_dir, f"{split}_codes_index.npy")
        self.codes = np.memmap(codes_path, dtype=np.int32, mode="r")
        self.number_of_codes = self.codes.shape[0]
        self.codes_index = np.load(codes_index_path)

        self.length = len(self.filtered_samples_index)
        if self.original_length != self.codes_index.shape[0]:
            raise ValueError("The number of samples and codes indices must match!")

        logging.info(
            "Loadded dataset [%s] with [%d] samples in total, [%d] samples "
            "removed due to low sample rate.",
            dataset_name,
            self.length,
            self.original_length - self.length,
        )

    def __len__(self) -> int:
        return self.length

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        # convert filtered idx to the original index
        c_idx = self.filtered_samples_index[idx]

        left_boundary = self.codes_index[c_idx]
        right_boundary = (
            self.codes_index[c_idx + 1]
            if c_idx < self.original_length - 1
            else self.number_of_codes
        )
        audio_codes = self.codes[left_boundary:right_boundary]
        audio_codes = torch.from_numpy(audio_codes).unsqueeze(0)

        # load and pad audio
        wav, _ = data_utils.load_wav(
            self.samples[idx].wav_path, target_sample_rate=self.sample_rate
        )

        wav = torch.nn.functional.pad(
            wav, (0, self.hop_length - (wav.shape[1] % self.hop_length))
        )

        # TODO: figure out whether extra padding is needed
        # wav = torch.nn.functional.pad(
        #     wav, (self.half_hop_length, self.half_hop_length))

        # repeat wav and audio_codes if it's too short
        if wav.shape[0] < self.audio_window_size:
            n_repeat = self.audio_window_size // wav.shape[1] + 1
            wav = wav.repeat(1, n_repeat)

            audio_codes = audio_codes.repeat(1, n_repeat)
            wav = wav[:, : self.audio_window_size]

            audio_codes = audio_codes[:, : self.code_window_size]

        st = random.randint(0, audio_codes.shape[1] - self.code_window_size)
        audio_codes = audio_codes[:, st : st + self.code_window_size]
        wav = wav[
            :, (st * self.hop_length) : st * self.hop_length + self.audio_window_size
        ]

        return {
            "audio_codes": audio_codes,
            "wav": wav,
            "audio_processed_sec": wav.shape[1] / self.sample_rate,
        }

def collate_fn(features: list[Any]) -> dict[str, Any]:
    """collate_fn for codec datasets."""
    audio_codes = torch.cat([f["audio_codes"] for f in features], dim=0)
    wav = torch.cat([f["wav"] for f in features], dim=0)
    audio_processed_sec = torch.tensor([f["audio_processed_sec"] for f in features])

    return {
        "audio_codes": audio_codes,
        "wav": wav,
        "audio_processed_sec": audio_processed_sec,
    }

def merge_datasets(
    weighted_datasets: dict[str, float],
    audio_window_size: int,
    split: str,
    sample_rate: int,
    minimum_data_sample_rate: int,
) -> torch.utils.data.Dataset:
    """Merges multiple datasets into a single dataset."""
    logging.info("-*-" * 10 + " %s " + "-*-" * 10, split)

    # Normalize the dataset weights for non-training splits.
    normalized_dataset_weights = dict(weighted_datasets.items())
    if split != "train":
        max_dataset_weight = max(normalized_dataset_weights.values())
        normalized_dataset_weights = {
            dataset_path: dataset_weight / max_dataset_weight
            for dataset_path, dataset_weight in normalized_dataset_weights.items()
        }

    datasets = []
    for dataset_path in weighted_datasets:
        dataset_name = os.path.basename(dataset_path)
        dataset = CodecTrainingDataset(
            dataset_path,
            split,
            audio_window_size,
            sample_rate,
            minimum_data_sample_rate,
        )
        datasets.append(
            tts_datasets.WeightedDataset(
                name=dataset_name,
                dataset=dataset,
                epochs=normalized_dataset_weights[dataset_path],
            )
        )
        logging.info(
            f"Dataset [{dataset_name}] has [{len(dataset)}] samples and "
            f"epochs [{normalized_dataset_weights[dataset_path]:.1f}]."
        )

    merged_dataset = tts_datasets.CombinedDataset(datasets)
    if len(datasets) > 1:
        logging.info(
            "Merged dataset has %d sources and %d samples total.",
            len(merged_dataset.sources),
            len(merged_dataset),
        )
    logging.info("-*-" * (18 + len(split)))
    return merged_dataset

def get_dataloader(
    dataset: torch.utils.data.Dataset,
    batch_size: int,
    collate_fn: Callable,
    shuffle: bool,
    num_workers: int,
) -> torch.utils.data.DataLoader:
    """Creates a dataloader for the dataset."""
    # pin_memory must be False to avoid error.
    # num_workers must be > 0 for normal audio loading speed.
    return torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        pin_memory=False,
        collate_fn=collate_fn,
        num_workers=num_workers,
    )


**Identifying data sample structure**

The TTS dataset stores audio paths and codes in JSONL format and uses a Sample class in data_sample; I'll inspect data_sample next to understand its structure.


sed -n '1,200p' tts/tts/data/data_sample.py


Output:
Exit code: 0
Wall time: 0 seconds
Output:
import dataclasses
import uuid
from typing import Any

_DEFAULT_VALUES = {
    "speaker_id": "",
    "emotion": "",
    "language": "unknown",
    "dnsmos_mos_ovr": 0.0,
    "style": "",
}
_REQUIRED_FIELDS = frozenset(("wav_path",))

@dataclasses.dataclass
class Sample:
    """A TTS sample. Can and should be adjusted for the specific use-cases.

    Args:
        id: The id of the sample.
        wav_path: The path to the wav file.
        speaker_id: The speaker id.

        language: The language of the sample.
        emotion: The emotion of the sample.
        transcript: The transcript of the sample.
        voice_description: The voice description of the sample.
        sound_effect: The sound effect of the sample.

        duration: The duration of the sample in seconds.
        sample_rate: Original sample rate of the wav file.
        dataset_name: The name of the dataset the sample is from.
        dnsmos_mos_ovr: The dnsmos mos ovr of the sample.
        style: The style of the sample (e.g., whispering, shouting, laughing, etc.).
        original_data: The original sample data used to support flexible formatting.
    """

    # Metadata.
    id: str
    wav_path: str
    speaker_id: str

    # Annotations.
    language: str
    emotion: str
    transcript: str
    voice_description: str
    sound_effect: str

    # Properties.
    duration: float
    sample_rate: int
    dataset_name: str
    dnsmos_mos_ovr: float
    style: str
    original_data: dict[str, Any] = dataclasses.field(default_factory=dict)

    def __post_init__(self):
        """Validates sample fields."""
        if not self.transcript and not self.voice_description and not self.sound_effect:
            raise ValueError(
                "At least one of transcript, voice_description, "
                "or sound_effect must be set."
            )

    def to_json(self) -> dict[str, str]:
        return {k: v for k, v in dataclasses.asdict(self).items() if v}

    @classmethod
    def from_json(cls, json_data: dict[str, Any], dataset_name: str) -> "Sample":
        if not dataset_name:
            raise ValueError("dataset_name is required")
        for field in _REQUIRED_FIELDS:
            if json_data.get(field) is None:
                raise ValueError(f"{field} is required for sample: {json_data}")

        return Sample(
            id=json_data.get("id", str(uuid.uuid4())),
            wav_path=json_data["wav_path"],
            speaker_id=json_data.get("speaker_id", _DEFAULT_VALUES["speaker_id"]),
            emotion=json_data.get("emotion", _DEFAULT_VALUES["emotion"]).lower(),
            transcript=json_data.get("transcript", ""),
            voice_description=json_data.get("voice_description", ""),
            sound_effect=json_data.get("sound_effect", ""),
            language=json_data.get("language", _DEFAULT_VALUES["language"]),
            duration=json_data.get("duration", -1.0),
            sample_rate=json_data.get("sample_rate", -1),
            dataset_name=dataset_name,
            dnsmos_mos_ovr=json_data.get(
                "dnsmos_mos_ovr", _DEFAULT_VALUES["dnsmos_mos_ovr"]
            ),
            style=json_data.get("style", _DEFAULT_VALUES["style"]).lower(),
            original_data=json_data.get("original_data", {}),
        )


**Assessing dataset compatibility with X-Codec training**

I found that X-Codec-2.0 expects filelists with local audio files, not .tar archives, so the current TTS .tar plus JSONL manifest format won't work directly. To support streaming from tar files, we'll need to either modify the training code or add a preprocessing step creating filelists from tar contents or implement a WebDataset dataloader.


**Planning WebDataset support without external dependencies**

I'm considering adding a dataset class in X-Codec-2.0 that reads from tar archives and JSONL metadata using built-in tarfile for random access to audio samples, avoiding new dependencies. This approach may be slower but keeps minimal setup; alternatively, we can provide a script to generate filelists from tar contents or adapt the pipeline to emit extracted flat files, balancing user requests and training code constraints.


Output:
Exit code: 2
Wall time: 0 seconds
Output:
sed: can't read CURRENT_STATUS.md: No such file or directory


**Listing open tabs content**


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


**Planning validation of downloaded audio segments**

I'll use the prepare_shards script with a limited download from R2 to generate audio shards and metadata, then manually verify each segment's length against metadata and possibly write a Python validation script to ensure all segments align correctly.


**Reviewing pipeline implementation against plan and preparing validation**

I've mapped key plan details to the prepare_shards pipeline, noting it covers download, metrics, resampling, and shard creation inline without separate steps. I'll validate segment counts, durations, and metadata consistency on one R2 video, check output shards and manifests, and verify audio sample rates and lengths. Also, I'll assess data compatibility with X-Codec training requirements, especially filelist expectations and segment duration handling, to identify any needed adjustments.


sed -n '1,200p' X-Codec-2.0/config/default.yaml


Output:
Exit code: 0
Wall time: 0 seconds
Output:
defaults:
  - preprocess: default
  - dataset: default
  - model: default
  - train: default
log_dir: /aifs4su/data/zheny/bigcodec_final/logs_1020/v10_31
debug: false
ckpt: null # infer
input_dir: null # infer
output_dir: null # infer
hydra:
  output_subdir: hydra
  job:
    chdir: True


sed -n '1,200p' X-Codec-2.0/config/preprocess/default.yaml


Output:
Exit code: 0
Wall time: 0 seconds
Output:
datasets:
    LibriSpeech:
        root: /aifs4su/data/zheny/data/data_8_21_2/
        trainsets: ['dev', 'test','train']
        testsets: ['test-clean']
view:
    train_filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_higher_quality.txt
    test_filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_test.txt
audio:
  sr: 16000
  normalization: False
  max_wav_value: 32768.0
stft:
  n_fft: 1280
  window_length: 1280
  hop_length: 320


**Planning dataset and path compatibility**

I'm weighing whether to modify the dataset for absolute filelist paths and WebDataset support or add a script to generate local filelists from manifests; a minor dataset change to handle absolute paths seems necessary, and I’m considering adding WebDataset reading with the `webdataset` library to support tar files.


**Planning tar-based dataset class**

I'm considering adding a new dataset in X-Codec-2.0 that reads from the manifest JSONL with tar shard and member info, managing tar file access per DataLoader worker, ideally leveraging or adding the `webdataset` library; adapting training configs to use this will ensure compatibility with our tar data.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
import os
from os.path import join, isdir
from pathlib import Path
from multiprocessing import Pool, cpu_count
from functools import partial
from tqdm import tqdm

def find_files_in_subdir(subdir, extension):
    """
    在给定的子目录中查找具有特定扩展名的文件。

    Args:
        subdir (str): 子目录路径。
        extension (str): 文件扩展名（例如 ".txt"）。

    Returns:
        list: 符合条件的文件列表，每个元素是 [文件名（无扩展名），文件路径]。
    """
    out = []
    try:
        for root, dirs, filenames in os.walk(subdir):
            for f in filenames:
                if f.endswith(extension):
                    out.append([str(Path(f).stem), os.path.join(root, f)])
    except Exception as e:
        print(f"Error processing {subdir}: {e}")
    return out

def get_second_level_subdirs(path_dir):
    """
    获取指定目录下的所有二级子目录。
    如果一级子目录没有二级子目录，则将一级子目录本身作为处理单元。

    Args:
        path_dir (str): 目标目录路径。

    Returns:
        list: 所有二级子目录路径列表。
    """
    second_level_subdirs = []
    try:
        first_level = [join(path_dir, d) for d in os.listdir(path_dir) if isdir(join(path_dir, d))]
        for first_subdir in first_level:
            second_subdirs = [join(first_subdir, sd) for sd in os.listdir(first_subdir) if isdir(join(first_subdir, sd))]
            if second_subdirs:
                second_level_subdirs.extend(second_subdirs)
            else:
                # 如果一级子目录没有二级子目录，则将一级子目录本身作为处理单元
                second_level_subdirs.append(first_subdir)
    except Exception as e:
        print(f"Error accessing directories in {path_dir}: {e}")
    return second_level_subdirs

def find_all_files(path_dir, extension):
    """
    使用多进程查找指定目录及其二级子目录中所有具有特定扩展名的文件。

    Args:
        path_dir (str): 目标目录路径。
        extension (str): 文件扩展名（例如 ".txt"）。

    Returns:
        list: 所有符合条件的文件列表。
    """
    out = []
    subdirs = get_second_level_subdirs(path_dir)

    if not subdirs:
        subdirs = [path_dir]

    func = partial(find_files_in_subdir, extension=extension)

    with Pool(processes=100) as pool:
        for result in tqdm(pool.imap(func, subdirs), total=len(subdirs), desc="Processing subdirectories"):
            out.extend(result)

    return out

def read_filelist(path, delimiter='|'):
    """
    读取文件列表，每行使用指定的分隔符分割。

    Args:
        path (str): 文件路径。
        delimiter (str): 分隔符，默认为 '|'.

    Returns:
        list: 分割后的文件列表，每个元素是一个列表。
    """
    with open(path, encoding='utf8') as f:
        lines = [line.strip().split(delimiter) for line in f if line.strip()]
    return lines

def write_filelist(filelists, path, delimiter='|'):
    """
    将文件列表写入指定路径，每行使用指定的分隔符连接。

    Args:
        filelists (list): 文件列表，每个元素是一个列表。
        path (str): 输出文件路径。
        delimiter (str): 分隔符，默认为 '|'.
    """
    with open(path, 'w', encoding='utf8') as f:
        for line in filelists:
            f.write(delimiter.join(line) + '\n')

# 示例用法
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="多进程查找文件并显示进度条（按二级目录）")
    parser.add_argument("directory", help="目标目录路径")
    parser.add_argument("extension", help="目标文件扩展名，例如 .txt")
    parser.add_argument("output", help="输出文件列表的路径")
    args = parser.parse_args()

    directory = args.directory
    extension = args.extension
    output_path = args.output

    all_files = find_all_files(directory, extension)
    write_filelist(all_files, output_path)

    print(f"查找完成，共找到 {len(all_files)} 个文件。结果已写入 {output_path}。")


Summary: auto


**Planning tar-backed dataset**

I'll implement a new dataset class to read audio from tar files via the JSONL manifest, focusing on simple per-sample tar openings for now and adding config options to switch between tar-based and shard-based datasets. Next, I'll review the existing tsv file generator to check if it supports tar or needs adjustments.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
from multiprocessing import Pool, cpu_count
import os
import torchaudio
import torch
from tqdm import tqdm

def process_file(args):
    file_path, root_dir = args  # 解包传入的参数
    try:
        rel_path = os.path.relpath(file_path, start=root_dir)
        waveform, sample_rate = torchaudio.load(file_path)
        nsample = waveform.shape[1]
        batch_size = 10000  # 选择适合的批处理大小
        for start in range(0, waveform.numel(), batch_size):
            end = min(start + batch_size, waveform.numel())
            if torch.isnan(waveform.view(-1)[start:end]).any():
                print(rel_path)
                return None  # 如果包含nan，则不处理这个文件

        if nsample == 0:
            return None
        return f"{rel_path}\t{nsample}\n"
    except Exception as e:
        print(f"Error processing {file_path}: {e}")
        return None

def list_audio_files(root_dir, output_file, exclude_dirs=None):
    if exclude_dirs is None:
        exclude_dirs = []

    audio_files = []
    for root, dirs, files in os.walk(root_dir):
        # 排除指定的子文件夹
        dirs[:] = [d for d in dirs if os.path.join(root, d) not in exclude_dirs]

        for filename in files:
            if filename.endswith(('.wav', '.flac', '.mp3')):
                file_path = os.path.join(root, filename)
                audio_files.append((file_path, root_dir))  # 将root_dir与文件路径一起打包

    # 按文件名排序
    audio_files.sort(key=lambda x: x[0])

    # 使用多进程处理文件
    pool = Pool(processes=int(cpu_count() / 2))  # 使用一半的CPU核心
    results = list(tqdm(pool.imap(process_file, audio_files), total=len(audio_files), desc="Processing audio files"))

    # 写入结果到文件
    with open(output_file, 'w') as file:
        file.write(f"{root_dir}\n")
        for result in results:
            if result:  # 只有当result不为None时才写入文件
                file.write(result)

# 示例使用
root_directory = '/aifs4su/data/zheny/data/data_8_21_2'
output_tsv = '/aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_higher_quality.txt'
exclude_folders = ['/aifs4su/data/zheny/data/data_8_21_2/test-clean']

list_audio_files(root_directory, output_tsv, exclude_dirs=exclude_folders)


**Planning complex multi-step task**

I'm creating a detailed plan to inspect pipeline docs and scripts, map out the strategy, identify outputs and gaps, then integrate and validate using actual pipeline runs for reliability.


**Detailing stepwise validation and integration**

The plan now includes selecting a test video, validating tar and metadata, running prepare_shards with output checks, and reviewing X-Codec data loading for .tar compatibility with fixes and config updates.


Review strategy vs current pipeline scripts to map requirements and identify initial gaps.

Pick a Supabase video, download its cleaned tar from R2, and validate metadata/segment audio alignment.

Run the sharding pipeline on that video and verify shard/manifest/audio outputs match expectations.

Check X-Codec-2.0 data loading compatibility with tar shards; implement minimal fixes or adapters as needed.

Explanation: Plan to verify the pipeline end-to-end on one video, then align training code with the produced tar shards.


Output: Plan updated


**Preparing Python snippet to query sample video**

I'm drafting a Python script to parse the .env for the database URL and query a sample video with a non-null and completed status, filtering by a random language from the configured list and printing the video ID, language, and key without exposing secrets.


python - <<'PY'
from pathlib import Path
import psycopg

env = {}
for line in Path('.env').read_text().splitlines():
    s = line.strip()
    if not s or s.startswith('#') or '=' not in s:
        continue
    if s.startswith('export '):
        s = s[len('export '):]
    k, v = s.split('=', 1)
    env[k.strip()] = v.strip()

db_url = env.get('DATABASE_URL')
if not db_url:
    raise SystemExit('DATABASE_URL missing')

query = """
    select youtube_id, lower(language) as language, r2_tar_key
    from videos
    where r2_tar_key is not null
      and status = 'COMPLETED'
      and language is not null
    order by random()
    limit 1
"""

with psycopg.connect(db_url) as conn:
    with conn.cursor() as cur:
        cur.execute(query)
        row = cur.fetchone()
        if not row:
            raise SystemExit('No completed rows found')
        youtube_id, language, r2_tar_key = row
        print(youtube_id)
        print(language)
        print(r2_tar_key)
PY


Output:
Exit code: 127
Wall time: 0 seconds
Output:
/bin/bash: line 1: python: command not found


**Adjusting to use python3**

I'm switching to use `python3` instead of `python` to ensure the correct interpreter is found.


python3 - <<'PY'
from pathlib import Path
import psycopg

env = {}
for line in Path('.env').read_text().splitlines():
    s = line.strip()
    if not s or s.startswith('#') or '=' not in s:
        continue
    if s.startswith('export '):
        s = s[len('export '):]
    k, v = s.split('=', 1)
    env[k.strip()] = v.strip()

db_url = env.get('DATABASE_URL')
if not db_url:
    raise SystemExit('DATABASE_URL missing')

query = """
    select youtube_id, lower(language) as language, r2_tar_key
    from videos
    where r2_tar_key is not null
      and status = 'COMPLETED'
      and language is not null
    order by random()
    limit 1
"""

with psycopg.connect(db_url) as conn:
    with conn.cursor() as cur:
        cur.execute(query)
        row = cur.fetchone()
        if not row:
            raise SystemExit('No completed rows found')
        youtube_id, language, r2_tar_key = row
        print(youtube_id)
        print(language)
        print(r2_tar_key)
PY


Output:
Exit code: 124
Wall time: 10 seconds
Output:
command timed out after 10007 milliseconds


**Extending command timeout**

I'm increasing the command timeout from the default 10 seconds to 20 seconds to avoid premature termination due to connection delays.


Output:
Exit code: 0
Wall time: 11.1 seconds
Output:
xcN9cg-sR3o
telugu
cleaned/trail/xcN9cg-sR3o.tar


**Extracting bucket info from environment**

I'm planning to parse the .env file in Python to safely read bucket names and prefixes without exposing credentials, ensuring these values align with keys used for R2 tar file access and shard preparation.


python3 - <<'PY'
from pathlib import Path
keys = ['R2_CLEANED_BUCKET','R2_ClEANED_BUCKET','R2_BUCKET_SHARDS','R2_PREFIX']
vals = {}
for line in Path('.env').read_text().splitlines():
    s = line.strip()
    if not s or s.startswith('#') or '=' not in s:
        continue
    if s.startswith('export '):
        s = s[len('export '):]
    k, v = s.split('=', 1)
    k = k.strip()
    if k in keys:
        vals[k] = v.strip()
for k in keys:
    if k in vals:
        print(f"{k}={vals[k]}")
PY


Output:
Exit code: 0
Wall time: 0 seconds
Output:
R2_CLEANED_BUCKET=1-cleaned-data
R2_PREFIX=cleaned/trail


**Planning precise tar validation**

The current prepare_shards script picks a random video limiting by language, which won't guarantee testing our chosen video. Considering modifying it to accept a specific video ID or using local tar inputs with skip-downloads to ensure controlled validation of the exact tar.


**Designing targeted video filtering patch**

Deciding to add optional `--video-id` args to prepare_shards for explicit record filtering, enabling precise selection without random DB queries. This improves control for validation but will be done carefully to avoid side effects and preserve existing behavior.


python3 scripts/prepare_shards.py --language telugu --limit 1 --output-dir /tmp/pipeline_test --tmp-dir /tmp/xcodec_work --download-workers 4 --segment-workers 4


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

download:   0%|          | 0/1 [00:00<?, ?it/s]
download: 100%|██████████| 1/1 [00:00<00:00,  2.30it/s]
download: 100%|██████████| 1/1 [00:00<00:00,  2.30it/s]

process:   0%|          | 0/1 [00:00<?, ?it/s]
process: 100%|██████████| 1/1 [00:00<00:00, 63.56it/s]
Done. Stats written to /tmp/pipeline_test/stats/telugu_stats.json


**Confirming download location and filename**

The downloaded video is stored under /tmp/xcodec_work/downloads with the filename based on the YouTube video ID; next step is listing files there to verify.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
U2c2AMyTYRo.tar


**Planning video metadata validation**

I'll validate the video using its ID and query Supabase to confirm related keys and language, skipping the optional DB r2_key check.


python3 - <<'PY'
from pathlib import Path
import psycopg

env = {}
for line in Path('.env').read_text().splitlines():
    s = line.strip()
    if not s or s.startswith('#') or '=' not in s:
        continue
    if s.startswith('export '):
        s = s[len('export '):]
    k, v = s.split('=', 1)
    env[k.strip()] = v.strip()

db_url = env.get('DATABASE_URL')
video_id = 'U2c2AMyTYRo'
with psycopg.connect(db_url) as conn:
    with conn.cursor() as cur:
        cur.execute(
            """
            select youtube_id, lower(language) as language, r2_tar_key, status
            from videos
            where youtube_id = %s
            """,
            (video_id,),
        )
        row = cur.fetchone()
        print(row)
PY


Output:
Exit code: 0
Wall time: 0.7 seconds
Output:
('U2c2AMyTYRo', 'telugu', 'cleaned/trail/U2c2AMyTYRo.tar', 'COMPLETED')


Summary: auto


**Planning tar metadata validation**

I'm outlining a Python script to open the tar, confirm the metadata.json path, parse segment status and counts, and cross-verify each flac file's presence and duration against metadata. First, I'll inspect metadata structure by listing tar members and sampling metadata keys.


python3 - <<'PY'
import io
import json
import tarfile
from pathlib import Path
import numpy as np
import soundfile as sf

video_id = 'U2c2AMyTYRo'
tar_path = Path('/tmp/xcodec_work/downloads') / f'{video_id}.tar'

with tarfile.open(tar_path, 'r') as tf:
    meta_member = tf.getmember(f"{video_id}/metadata.json")
    meta = json.loads(tf.extractfile(meta_member).read())

    segments_meta = meta.get('segments', [])
    usable_segments = [s for s in segments_meta if s.get('status') == 'usable']

    members = [m for m in tf.getmembers() if m.name.endswith('.flac')]

    print('segments_total', len(segments_meta))
    print('segments_usable', len(usable_segments))
    print('segment_files', len(members))

    # Build lookup: (speaker, start rounded to 2) -> meta
    seg_map = {(s.get('speaker'), round(float(s.get('start', 0.0)), 2)): s for s in usable_segments}

    mismatched = 0
    duration_mismatch = 0
    sr_mismatch = 0
    duration_diffs = []
    for m in members:
        base = m.name.split('/')[-1]
        if not base.endswith('.flac'):
            continue
        base = base[:-5]
        parts = base.split('_')
        if len(parts) < 3:
            mismatched += 1
            continue
        speaker = '_'.join(parts[:-2])
        time_range = parts[-1]
        if '-' not in time_range:
            mismatched += 1
            continue
        start_time = float(time_range.split('-')[0])
        seg_meta = seg_map.get((speaker, round(start_time, 2)))
        if not seg_meta:
            # try fuzzy match within 0.1s
            candidates = [s for s in usable_segments if s.get('speaker') == speaker and abs(s.get('start', 0.0) - start_time) < 0.1]
            if not candidates:
                mismatched += 1
                continue
            seg_meta = candidates[0]

        audio_bytes = tf.extractfile(m).read()
        audio, sr = sf.read(io.BytesIO(audio_bytes))
        duration = audio.shape[0] / sr
        start = seg_meta.get('start')
        end = seg_meta.get('end')
        dur_meta = None
        if start is not None and end is not None:
            dur_meta = float(end) - float(start)
        elif seg_meta.get('duration') is not None:
            dur_meta = float(seg_meta['duration'])
        if dur_meta is not None:
            diff = abs(duration - dur_meta)
            duration_diffs.append(diff)
            if diff > 0.05:  # 50ms tolerance
                duration_mismatch += 1
        if sr != seg_meta.get('sample_rate', sr):
            # sample_rate may not exist; track mismatch if present
            if seg_meta.get('sample_rate') is not None:
                sr_mismatch += 1

    print('segment_meta_missing', mismatched)
    if duration_diffs:
        arr = np.array(duration_diffs)
        print('duration_diff_max', float(arr.max()))
        print('duration_diff_p95', float(np.percentile(arr, 95)))
        print('duration_diff_mean', float(arr.mean()))
    print('duration_mismatch_gt50ms', duration_mismatch)
    print('sr_mismatch', sr_mismatch)
PY


Output:
Exit code: 0
Wall time: 0.5 seconds
Output:
segments_total 77
segments_usable 22
segment_files 22
segment_meta_missing 0
duration_diff_max 1.0416666696988486e-05
duration_diff_p95 1.0416666693568999e-05
duration_diff_mean 6.155303033888499e-06
duration_mismatch_gt50ms 0
sr_mismatch 0


**Validating segment integrity and stats**

All 22 tar segments match metadata with matching durations and counts; next, I'll verify audio sample rates and per-segment stats from the pipeline outputs and analyze the stats JSON for consistency.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
{
  "videos": 1,
  "segments_total": 22,
  "segments_kept": {
    "16000": 21,
    "24000": 16,
    "32000": 16,
    "44100": 16
  },
  "dropped_too_short": 12,
  "dropped_too_long": 0,
  "dropped_clipped": 0,
  "segments_chunked": 5,
  "chunks_emitted": 21,
  "native_sr_buckets": {
    "32-44k": 21
  }
}


**Clarifying segment counts and chunking**

The stats indicate 22 original segments with 12 dropped due to short length, but chunking long segments creates multiple outputs, leading to 21 kept chunks overall. Next, I'll inspect a manifest file to verify how these entries correspond.


**Reviewing manifest split logic**

The validation count for a single video is zero due to how val_count is computed, so all examples go to training; next, I'll examine a few lines from the manifest to confirm data allocation.


head -n 5 /tmp/pipeline_test/manifests/telugu_train.jsonl


Output:
Exit code: 0
Wall time: 0 seconds
Output:
{"id": "U2c2AMyTYRo__SPEAKER_00_0000", "video_id": "U2c2AMyTYRo", "language": "telugu", "segment_index": 0, "chunk_index": null, "start": 0.03096875, "end": 7.304093750000001, "original_start": 180.03096875, "original_end": 187.30409375, "duration_sec": 7.273125, "native_sr": 48000, "effective_nyquist_hz": 19080.4, "spectral_rolloff_hz": 804.8000000000001, "high_freq_energy_ratio": 0.0032576530049229736, "is_truly_wideband": true, "native_sr_bucket": "32-44k", "loudness_lufs": -22.7992540755854, "loudness_rms_db": -22.398625217960028, "clipping_fraction": 0.0, "snr_proxy": 24.68301158801352, "quality_score": 0.8302001741078947, "music_prob": 0.0647, "music_stats": {"music_mean": 0.0647, "music_max": 0.0835, "music_ratio": 0.0, "noise_mean": 0.0392, "noise_max": 0.0482, "noise_ratio": 0.0, "contamination_ratio": 0.0, "has_music": false, "has_noise": false, "has_contamination": false, "decision": "clean", "chunks_analyzed": 5, "chunks_with_music": 0, "chunks_with_noise": 0, "chunks_with_contamination": 0}, "speaker": "SPEAKER_00", "vad_speech_regions": [[0.0, 7.273125]], "duration_bucket": "6-12s", "sampling_weight": 1.0, "sample_rate": 16000, "resample_target_sr": 16000, "shard": "shard-000000.tar", "audio": "U2c2AMyTYRo__SPEAKER_00_0000.flac"}
{"id": "U2c2AMyTYRo__SPEAKER_00_0000", "video_id": "U2c2AMyTYRo", "language": "telugu", "segment_index": 0, "chunk_index": null, "start": 0.03096875, "end": 7.30409375, "original_start": 180.03096875, "original_end": 187.30409375, "duration_sec": 7.273125, "native_sr": 48000, "effective_nyquist_hz": 19080.4, "spectral_rolloff_hz": 804.8000000000001, "high_freq_energy_ratio": 0.0032576530049229736, "is_truly_wideband": true, "native_sr_bucket": "32-44k", "loudness_lufs": -22.7992540755854, "loudness_rms_db": -22.398625217960028, "clipping_fraction": 0.0, "snr_proxy": 24.68301158801352, "quality_score": 0.8302001741078947, "music_prob": 0.0647, "music_stats": {"music_mean": 0.0647, "music_max": 0.0835, "music_ratio": 0.0, "noise_mean": 0.0392, "noise_max": 0.0482, "noise_ratio": 0.0, "contamination_ratio": 0.0, "has_music": false, "has_noise": false, "has_contamination": false, "decision": "clean", "chunks_analyzed": 5, "chunks_with_music": 0, "chunks_with_noise": 0, "chunks_with_contamination": 0}, "speaker": "SPEAKER_00", "vad_speech_regions": [[0.0, 7.273125]], "duration_bucket": "6-12s", "sampling_weight": 1.0, "sample_rate": 24000, "resample_target_sr": 24000, "shard": "shard-000000.tar", "audio": "U2c2AMyTYRo__SPEAKER_00_0000.flac"}
{"id": "U2c2AMyTYRo__SPEAKER_00_0000", "video_id": "U2c2AMyTYRo", "language": "telugu", "segment_index": 0, "chunk_index": null, "start": 0.03096875, "end": 7.30409375, "original_start": 180.03096875, "original_end": 187.30409375, "duration_sec": 7.273125, "native_sr": 48000, "effective_nyquist_hz": 19080.4, "spectral_rolloff_hz": 804.8000000000001, "high_freq_energy_ratio": 0.0032576530049229736, "is_truly_wideband": true, "native_sr_bucket": "32-44k", "loudness_lufs": -22.7992540755854, "loudness_rms_db": -22.398625217960028, "clipping_fraction": 0.0, "snr_proxy": 24.68301158801352, "quality_score": 0.8302001741078947, "music_prob": 0.0647, "music_stats": {"music_mean": 0.0647, "music_max": 0.0835, "music_ratio": 0.0, "noise_mean": 0.0392, "noise_max": 0.0482, "noise_ratio": 0.0, "contamination_ratio": 0.0, "has_music": false, "has_noise": false, "has_contamination": false, "decision": "clean", "chunks_analyzed": 5, "chunks_with_music": 0, "chunks_with_noise": 0, "chunks_with_contamination": 0}, "speaker": "SPEAKER_00", "vad_speech_regions": [[0.0, 7.273125]], "duration_bucket": "6-12s", "sampling_weight": 1.0, "sample_rate": 32000, "resample_target_sr": 32000, "shard": "shard-000000.tar", "audio": "U2c2AMyTYRo__SPEAKER_00_0000.flac"}
{"id": "U2c2AMyTYRo__SPEAKER_00_0000", "video_id": "U2c2AMyTYRo", "language": "telugu", "segment_index": 0, "chunk_index": null, "start": 0.03096875, "end": 7.30409375, "original_start": 180.03096875, "original_end": 187.30409375, "duration_sec": 7.273125, "native_sr": 48000, "effective_nyquist_hz": 19080.4, "spectral_rolloff_hz": 804.8000000000001, "high_freq_energy_ratio": 0.0032576530049229736, "is_truly_wideband": true, "native_sr_bucket": "32-44k", "loudness_lufs": -22.7992540755854, "loudness_rms_db": -22.398625217960028, "clipping_fraction": 0.0, "snr_proxy": 24.68301158801352, "quality_score": 0.8302001741078947, "music_prob": 0.0647, "music_stats": {"music_mean": 0.0647, "music_max": 0.0835, "music_ratio": 0.0, "noise_mean": 0.0392, "noise_max": 0.0482, "noise_ratio": 0.0, "contamination_ratio": 0.0, "has_music": false, "has_noise": false, "has_contamination": false, "decision": "clean", "chunks_analyzed": 5, "chunks_with_music": 0, "chunks_with_noise": 0, "chunks_with_contamination": 0}, "speaker": "SPEAKER_00", "vad_speech_regions": [[0.0, 7.273125]], "duration_bucket": "6-12s", "sampling_weight": 1.0, "sample_rate": 44100, "resample_target_sr": 44100, "shard": "shard-000000.tar", "audio": "U2c2AMyTYRo__SPEAKER_00_0000.flac"}
{"id": "U2c2AMyTYRo__SPEAKER_00_0004", "video_id": "U2c2AMyTYRo", "language": "telugu", "segment_index": 4, "chunk_index": null, "start": 122.52659375, "end": 131.99998958333333, "original_start": 302.52659375, "original_end": 311.99998958333333, "duration_sec": 9.473395833333333, "native_sr": 48000, "effective_nyquist_hz": 18179.600000000002, "spectral_rolloff_hz": 759.8000000000001, "high_freq_energy_ratio": 0.010404405807753443, "is_truly_wideband": true, "native_sr_bucket": "32-44k", "loudness_lufs": -22.472722652253275, "loudness_rms_db": -22.06710979899238, "clipping_fraction": 0.0, "snr_proxy": 26.23447584905503, "quality_score": 0.7430550493609828, "music_prob": 0.1298, "music_stats": {"music_mean": 0.1298, "music_max": 0.4158, "music_ratio": 0.0909, "noise_mean": 0.041, "noise_max": 0.051, "noise_ratio": 0.0, "contamination_ratio": 0.0909, "has_music": true, "has_noise": false, "has_contamination": true, "decision": "needs_demucs", "chunks_analyzed": 11, "chunks_with_music": 1, "chunks_with_noise": 0, "chunks_with_contamination": 1}, "speaker": "SPEAKER_00", "vad_speech_regions": [[0.0, 9.473395833333333]], "duration_bucket": "6-12s", "sampling_weight": 0.5, "sample_rate": 16000, "resample_target_sr": 16000, "shard": "shard-000000.tar", "audio": "U2c2AMyTYRo__SPEAKER_00_0004.flac"}


**Noticing manifest duplicates across sample rates**

Manifests list identical audio IDs repeated for different sample rates without distinguishing shard paths, causing ambiguity if used standalone; likely the pipeline intends separate per-sample-rate datasets despite combined manifests including duplicates.


**Confirming per-language manifests cover all sample rates**

The manifests likely include all sample rates per language, so filtering by `sample_rate` is needed for training specific subsets like 16k; the next step is verifying shard contents in each bucket directory.


Output:
Exit code: 0
Wall time: 0 seconds
Output:
shard-000000.jsonl
shard-000000.tar


**Verifying sample counts and rates**

I'm planning to use Python to inspect sample counts and audio sample rates by reading the tar archive and matching entries with the JSONL metadata.


python3 - <<'PY'
import io
import json
import tarfile
from pathlib import Path
import soundfile as sf

bucket_dir = Path('/tmp/pipeline_test/telugu_16')
jsonl_path = bucket_dir / 'shard-000000.jsonl'
tar_path = bucket_dir / 'shard-000000.tar'

entries = [json.loads(line) for line in jsonl_path.read_text().splitlines() if line.strip()]
with tarfile.open(tar_path, 'r') as tf:
    names = set(tf.getnames())
    missing = [e['id'] for e in entries if f"{e['id']}.flac" not in names]
    print('entries', len(entries))
    print('missing', len(missing))
    # Check first 5 sample rates
    for e in entries[:5]:
        member = tf.getmember(f"{e['id']}.flac")
        audio_bytes = tf.extractfile(member).read()
        audio, sr = sf.read(io.BytesIO(audio_bytes))
        print(e['id'], 'sr', sr, 'duration', audio.shape[0]/sr)
PY


Output:
Exit code: 0
Wall time: 0.2 seconds
Output:
entries 21
missing 0
U2c2AMyTYRo__SPEAKER_00_0000 sr 16000 duration 7.273125
U2c2AMyTYRo__SPEAKER_00_0004 sr 16000 duration 9.4734375
U2c2AMyTYRo__SPEAKER_00_0007 sr 16000 duration 6.125625
U2c2AMyTYRo__SPEAKER_00_0008 sr 16000 duration 5.821875
U2c2AMyTYRo__SPEAKER_00_0010 sr 16000 duration 7.509375


**Confirming match counts**

The current bucket has 16 entries matching the expected count; next, I'll verify if other buckets also have 16 entries and assess the total count around 24,000.


python3 - <<'PY'
import io
import json
import tarfile
from pathlib import Path
import soundfile as sf

bucket_dir = Path('/tmp/pipeline_test/telugu_24')
jsonl_path = bucket_dir / 'shard-000000.jsonl'
tar_path = bucket_dir / 'shard-000000.tar'

entries = [json.loads(line) for line in jsonl_path.read_text().splitlines() if line.strip()]
with tarfile.open(tar_path, 'r') as tf:
    names = set(tf.getnames())
    missing = [e['id'] for e in entries if f"{e['id']}.flac" not in names]
    print('entries', len(entries))
    print('missing', len(missing))
    # Check first 3 sample rates
    for e in entries[:3]:
        member = tf.getmember(f"{e['id']}.flac")
        audio_bytes = tf.extractfile(member).read()
        audio, sr = sf.read(io.BytesIO(audio_bytes))
        print(e['id'], 'sr', sr, 'duration', audio.shape[0]/sr)
PY


Output:
Exit code: 0
Wall time: 0.1 seconds
Output:
entries 16
missing 0
U2c2AMyTYRo__SPEAKER_00_0000 sr 24000 duration 7.273125
U2c2AMyTYRo__SPEAKER_00_0004 sr 24000 duration 9.473416666666667
U2c2AMyTYRo__SPEAKER_00_0007 sr 24000 duration 6.125625


**Updating plan progress**

Step 1 review is partially complete with some gaps noted; Step 2 validated tar segments; Step 3 pipeline run and outputs done; Step 4 is still pending.


Explanation: Reviewed the pipeline and validated one R2 video + sharding outputs; next is aligning X-Codec-2.0 data loading with tar shards.


Output: Plan updated


**Planning dataset compatibility fix**

The current X-Codec-2.0 code and config lack support for the tar shard dataset and reference a nonexistent dataset class, so I'll propose adding a new dataset class tailored to the tar manifest and a separate config to use it safely without altering defaults.


**Designing tar shard dataset with manifest parsing**

I'm outlining a new dataset class that loads audio from tar shards using a manifest JSONL, caching tarfile handles per worker, filtering to 16k sample rate entries, and decoding audio with soundfile. I'll also address config inconsistencies by creating a separate dataset class and config while planning to fix the hardcoded dataset instantiation in the data module for better configurability.


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

import pytorch_lightning as pl
import hydra
import torch
import random
import time
from os.path import join, basename, exists
from pytorch_lightning import seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.strategies import DDPStrategy,FSDPStrategy
from torch.utils.data import DataLoader
from data_module import DataModule
from lightning_module import CodecLightningModule
from pytorch_lightning.loggers import WandbLogger
from omegaconf import OmegaConf

seed = 1024
seed_everything(seed)

@hydra.main(config_path='config', config_name='default', version_base=None)
def train(cfg):
    checkpoint_callback = ModelCheckpoint(dirpath=cfg.log_dir,
                            save_top_k=-1, save_last=True,
                            every_n_train_steps=20000, monitor='mel_loss', mode='min')

    lr_monitor = LearningRateMonitor(logging_interval='step')
    callbacks = [checkpoint_callback, lr_monitor]

    datamodule = DataModule(cfg)
    lightning_module = CodecLightningModule(cfg)
    log_dir_name = os.path.basename(os.path.normpath(cfg.log_dir))
    wandb_logger = WandbLogger(
        project='xcodec2',  # 替换为您的项目名称
        name=log_dir_name,              # 替换为您的运行名称
        config=OmegaConf.to_container(cfg, resolve=True)  # 将 Hydra 配置转换为字典并传递
    )

    ckpt_path = None
    last_ckpt = os.path.join(cfg.log_dir, 'last.ckpt')
    if os.path.exists(last_ckpt):
        ckpt_path = last_ckpt
        print(f"Resuming from checkpoint: {ckpt_path}")
    else:
        print("No checkpoint found, starting training from scratch.")

    trainer = pl.Trainer(
        **cfg.train.trainer,
        strategy=DDPStrategy(find_unused_parameters=True),
        callbacks=callbacks,
        logger=wandb_logger,
        profiler="simple",  # 启用 Profiler
        limit_train_batches=1.0 if not cfg.debug else 0.001
    )
    torch.backends.cudnn.benchmark = True
    # lightning_module.strict_loading = False
    # LightningModule.strict_loading = True
    trainer.fit(lightning_module, datamodule=datamodule,ckpt_path=ckpt_path )
    print(f'Training ends, best score: {checkpoint_callback.best_model_score}, ckpt path: {checkpoint_callback.best_model_path}')

if __name__ == '__main__':
    train()


Summary: auto


**Inspecting dataset config**


sed -n '1,200p' X-Codec-2.0/config/train/default.yaml


Output:
Exit code: 0
Wall time: 0 seconds
Output:
trainer:
  accelerator: 'gpu'
  num_nodes: 1
  devices: 8
  min_steps: 3000000
  max_steps: 3000000
  precision: 16-mixed
  limit_val_batches: 2560
  val_check_interval: 4000
  num_sanity_val_steps: 0
  accumulate_grad_batches:  1

# loss
lambdas:
  lambda_disc: 1.0
  lambda_feat_match_loss: 1.0
  lambda_mel_loss: 15.0
  lambda_adv: 1.0
  lambda_stft_loss: 1.0
  lambda_semantic_loss: 5
  lambda_perceptual_loss: 0
use_mel_loss: true
use_feat_match_loss: true
use_stft_loss: false

stft_loss_params:
  fft_sizes: [512,2048]
  hop_sizes: [128,512]
  win_lengths: [512,2048]
  window: hann_window

# optimizer
gen_optim_params:
  lr: 1.0
  betas: [0.8, 0.9]
disc_optim_params:
  lr: 1.0
  betas: [0.8, 0.9]
gen_grad_clip: 1.0
disc_grad_clip: 1.0

# scheduler
gen_schedule_params:
  warmup_step: 3000
  down_step: 400000
  min_lr: 2.0e-5 # this should be the final lr
  max_lr: 1.0e-4
disc_schedule_params:
  warmup_step: 3000
  down_step: 400000
  min_lr: 2.0e-5 # this should be the final lr
  max_lr: 1.0e-4


**Identifying dataset config key conflicts**

Discovered a naming conflict where both dataset and training configs use a top-level `train` key, likely causing `cfg.train` to reference only training parameters and overwriting dataset filelist info; this explains why DataModule’s use of `cfg.dataset.get(phase)` fails since dataset splits are top-level, not nested under `dataset`.


**Clarifying config structure and planning dataset refactor**

Confirmed dataset config keys like `train` and `val` are nested under `cfg.dataset`, while training config keys like `trainer` live under `cfg.train` separately, so no key conflict occurs; detected that DataModule currently ignores dataset's `_target_` and always uses a fixed dataset class, so planning to refactor DataModule to instantiate dataset classes dynamically via Hydra using dataset config’s `_target_` for flexibility and to add a new ShardDataset class with its own config to handle tarred shard datasets, including loading manifests and resampling audio.