update llm training

speech/config.yaml

@@ -187,23 +187,15 @@ feat_extractor: !name:matcha.utils.audio.mel_spectrogram
     fmin: 0
     fmax: 8000
     center: False
-compute_fbank: !name:cosyvoice.dataset.processor.compute_fbank
-    feat_extractor: !ref <feat_extractor>
-    token_mel_ratio: !ref <token_mel_ratio>
-compute_f0: !name:cosyvoice.dataset.processor.compute_f0
-    sample_rate: !ref <sample_rate>
-    hop_size: 480
-parse_embedding: !name:cosyvoice.dataset.processor.parse_embedding
-    normalize: True
+
 shuffle: !name:cosyvoice.dataset.processor.shuffle
     shuffle_size: 1000
 sort: !name:cosyvoice.dataset.processor.sort
     sort_size: 500  # sort_size should be less than shuffle_size
 batch: !name:cosyvoice.dataset.processor.batch
     batch_type: 'dynamic'
-    max_frames_in_batch: 5000
+    max_frames_in_batch: 25000
 padding: !name:cosyvoice.dataset.processor.padding
-    use_spk_embedding: False # change to True during sft
     use_speaker_encoder: !ref <use_speaker_encoder>
 
 
@@ -213,9 +205,7 @@ data_pipeline: [
     !ref <tokenize>,
     !ref <filter>,
     !ref <resample>,
-    !ref <compute_fbank>,
     !ref <extract_reference_mel>,  # Add this for speaker encoder
-    !ref <parse_embedding>,
     !ref <shuffle>,
     !ref <sort>,
     !ref <batch>,

speech/cosyvoice/dataset/processor.py

@@ -24,17 +24,15 @@ import pyworld as pw
 import glob
 import os
 import json
-
+import traceback
 AUDIO_FORMAT_SETS = {'flac', 'mp3', 'm4a', 'ogg', 'opus', 'wav', 'wma'}
 
 
-def individual_file_opener(data, mode='train', tts_data={}):
-    """Load data from individual files instead of parquet
+def individual_file_opener(data, mode='train', tts_data={}, token_latent_ratio=3):
+    """Load data from individual files listed in files.txt
     
     Args:
-        data: Iterable[{src}] where src is either:
-            - Path to a directory containing audio files
-            - Path to a JSON index file
+        data: Iterable[{src}] where src is path to files.txt containing audio paths
         mode: 'train' or 'test'
         tts_data: Dict for TTS mode
         
@@ -45,51 +43,93 @@ def individual_file_opener(data, mode='train', tts_data={}):
         assert 'src' in sample
         src = sample['src']
         
-        # Determine if src is a directory or index file
-        if src.endswith('.json'):
-            # Load from index file
+        # Load file list from files.txt
+        file_list = []
+        
+        # Check if src is a files.txt file
+        if src.endswith('.txt'):
+            with open(src, 'r') as f:
+                wav_files = [line.strip() for line in f if line.strip()]
+            
+            for wav_path in wav_files:
+                # Skip empty lines or comments
+                if not wav_path or wav_path.startswith('#'):
+                    continue
+                    
+                # Verify wav file exists
+                if not os.path.exists(wav_path):
+                    logging.warning(f'Audio file not found: {wav_path}, skipping')
+                    continue
+                
+                # Check if all required files exist
+                txt_path = wav_path.replace('.wav', '.txt')
+                token_path = wav_path.replace('.wav', '_fsq.pt')
+                latent_path = wav_path.replace('.wav', '_latent.pt')
+                
+                if not os.path.exists(txt_path):
+                    logging.warning(f'Text file not found for {wav_path}, skipping')
+                    continue
+                
+                if not os.path.exists(token_path):
+                    logging.warning(f'Token file not found for {wav_path}, skipping')
+                    continue
+                    
+                if not os.path.exists(latent_path):
+                    logging.warning(f'Latent file not found for {wav_path}, skipping')
+                    continue
+                
+                # Extract metadata
+                utt = os.path.basename(wav_path).replace('.wav', '')
+                # Try to extract speaker from filename (assuming format: spk_*.wav)
+                spk = utt.split('_')[0] if '_' in utt else 'default'
+                
+                file_info = {
+                    'utt': utt,
+                    'spk': spk,
+                    'wav': wav_path,
+                    'text_path': txt_path,
+                    'token_path': token_path,
+                    'latent_path': latent_path,
+                }
+                logging.info(f'file_info {file_info}')
+                file_list.append(file_info)
+        
+        elif src.endswith('.json'):
+            # Keep backward compatibility with JSON index files
             with open(src, 'r') as f:
                 index_data = json.load(f)
                 file_list = index_data.get('data', [])
+        
         else:
-            # Scan directory for wav files
+            # Assume it's a directory for backward compatibility
             wav_files = glob.glob(os.path.join(src, '*/*/*wav'))
             if not wav_files:
-                # Try different patterns
                 wav_files = glob.glob(os.path.join(src, '**/*.wav'), recursive=True)
             
-            file_list = []
             for wav_path in wav_files:
-                # Check if all required files exist
-                txt_path = wav_path.replace('.wav', '.normalized.txt')
-                embedding_path = wav_path.replace('.wav', '_embedding.pt')
-                token_path = wav_path.replace('.wav', '_tokens.pt')
+                txt_path = wav_path.replace('.wav', '.txt')
+                token_path = wav_path.replace('.wav', '_fsq.pt')
+                latent_path = wav_path.replace('.wav', '_latent.pt')
                 
                 if not os.path.exists(txt_path):
                     logging.warning(f'Text file not found for {wav_path}, skipping')
                     continue
                 
-                # Extract metadata
                 utt = os.path.basename(wav_path).replace('.wav', '')
                 spk = utt.split('_')[0]
                 
-                # Find speaker embedding
-                spk_embed_dir = os.path.join(os.path.dirname(src), 'spk_embeddings')
-                if not os.path.exists(spk_embed_dir):
-                    spk_embed_dir = os.path.join(src, 'spk_embeddings')
-                spk_embedding_path = os.path.join(spk_embed_dir, f'{spk}_embedding.pt')
-                
                 file_info = {
                     'utt': utt,
                     'spk': spk,
                     'wav': wav_path,
                     'text_path': txt_path,
-                    'embedding_path': embedding_path,
                     'token_path': token_path,
-                    'spk_embedding_path': spk_embedding_path
+                    'latent_path': latent_path,
                 }
                 file_list.append(file_info)
         
+        logging.info(f'Found {len(file_list)} valid audio files from {src}')
+        
         # Process each file
         for file_info in file_list:
             try:
@@ -98,38 +138,24 @@ def individual_file_opener(data, mode='train', tts_data={}):
                     audio_data = f.read()
                 
                 # Read text
-                with open(file_info['text_path'], 'r') as f:
+                with open(file_info['text_path'], 'r', encoding='utf-8') as f:
                     text = ''.join(l.strip() for l in f.readlines())
                 
-                # Load embeddings if they exist
-                if os.path.exists(file_info['embedding_path']):
-                    utt_embedding = torch.load(file_info['embedding_path'], weights_only=False)
-                    if isinstance(utt_embedding, torch.Tensor):
-                        utt_embedding = utt_embedding.tolist()
-                else:
-                    logging.warning(f"Utterance embedding not found: {file_info['embedding_path']}")
-                    # Create a dummy embedding
-                    utt_embedding = [0.0] * 192  # Assuming 192-dim embeddings
-                
-                # Load tokens if they exist
-                if os.path.exists(file_info['token_path']):
-                    speech_token = torch.load(file_info['token_path'], weights_only=False)
-                    if isinstance(speech_token, torch.Tensor):
-                        speech_token = speech_token.tolist()
-                else:
-                    logging.warning(f"Token file not found: {file_info['token_path']}")
-                    speech_token = []
-                
-                # Load speaker embedding
-                if os.path.exists(file_info['spk_embedding_path']):
-                    spk_embedding = torch.load(file_info['spk_embedding_path'], weights_only=False)
-                    if isinstance(spk_embedding, torch.Tensor):
-                        spk_embedding = spk_embedding.tolist()
-                else:
-                    logging.warning(f"Speaker embedding not found: {file_info['spk_embedding_path']}")
-                    # Use utterance embedding as fallback
-                    spk_embedding = utt_embedding
-                
+                # Load speech token
+                speech_token = torch.load(file_info['token_path'], map_location='cpu', weights_only=False)
+                if isinstance(speech_token, torch.Tensor):
+                    speech_token = speech_token.tolist()
+
+                # Load speech latent
+                speech_latent = torch.load(file_info['latent_path'], map_location='cpu', weights_only=False)
+                speech_latent = speech_latent['z'].transpose(0, 1)
+
+                if token_latent_ratio != 0:
+                    # trim to align speech_token and speech_feat
+                    token_len = int(min(speech_latent.shape[0] / token_latent_ratio, len(speech_token)))
+                    speech_latent = speech_latent[:token_latent_ratio * token_len]
+                    speech_token = speech_token[:token_len]
+
                 # Build sample dict
                 sample_dict = {
                     'utt': file_info['utt'],
@@ -137,10 +163,9 @@ def individual_file_opener(data, mode='train', tts_data={}):
                     'audio_data': audio_data,
                     'text': text,
                     'text_token': [],  # Will be filled by tokenize processor
-                    'utt_embedding': utt_embedding,
-                    'spk_embedding': spk_embedding,
                     'speech_token': speech_token,
                     'wav': file_info['wav'],  # Keep original path for reference
+                    'speech_latent': speech_latent,
                 }
                 
                 # Copy over any additional fields from the original sample
@@ -237,8 +262,10 @@ def filter(data,
             continue
         if num_frames != 0:
             if len(sample['text_token']) / num_frames < min_output_input_ratio:
+                print('continue text_token')
                 continue
             if len(sample['text_token']) / num_frames > max_output_input_ratio:
+                print('continue text_token')
                 continue
         yield sample
 
@@ -261,6 +288,7 @@ def resample(data, resample_rate=22050, min_sample_rate=16000, mode='train'):
         waveform = sample['speech']
         if sample_rate != resample_rate:
             if sample_rate < min_sample_rate:
+                print('continue sample_rate')
                 continue
             sample['sample_rate'] = resample_rate
             sample['speech'] = torchaudio.transforms.Resample(
@@ -292,43 +320,6 @@ def truncate(data, truncate_length=24576, mode='train'):
         yield sample
 
 
-def compute_fbank(data,
-                  feat_extractor,
-                  token_mel_ratio=0,
-                  mode='train'):
-    """ Extract fbank
-
-        Args:
-            data: Iterable[{key, wav, label, sample_rate}]
-
-        Returns:
-            Iterable[{key, feat, label}]
-    """
-    for sample in data:
-        assert 'sample_rate' in sample
-        assert 'speech' in sample
-        assert 'utt' in sample
-        assert 'text_token' in sample
-        waveform = sample['speech']
-        feat = feat_extractor(waveform).squeeze(dim=0).transpose(0, 1)
-        if token_mel_ratio != 0:
-            
-            if isinstance(sample["speech_token"], list):
-                speech_token_tensor = torch.tensor(sample["speech_token"])
-            else:
-                speech_token_tensor = sample["speech_token"]
-            
-            # trim to align speech_token and speech_feat
-            token_len = int(min(feat.shape[0] / token_mel_ratio, speech_token_tensor.shape[0]))
-            feat = feat[:token_mel_ratio * token_len]
-            
-            # Update speech_token - keep as tensor for consistency
-            sample["speech_token"] = speech_token_tensor[:token_len]
-
-        sample['speech_feat'] = feat
-        yield sample
-
-
 def extract_reference_mel_from_speech(
     data,
     feat_extractor,
@@ -361,6 +352,7 @@ def extract_reference_mel_from_speech(
             sample['reference_mels'] = []
             sample['reference_mel_lengths'] = []
             sample['num_references'] = 0
+            print('continue num_references')
             yield sample
             continue
         
@@ -403,48 +395,6 @@ def extract_reference_mel_from_speech(
         
         yield sample
 
-def compute_f0(data, sample_rate, hop_size, mode='train'):
-    """ Extract f0
-
-        Args:
-            data: Iterable[{key, wav, label, sample_rate}]
-
-        Returns:
-            Iterable[{key, feat, label}]
-    """
-    frame_period = hop_size * 1000 / sample_rate
-    for sample in data:
-        assert 'sample_rate' in sample
-        assert 'speech' in sample
-        assert 'utt' in sample
-        assert 'text_token' in sample
-        waveform = sample['speech']
-        _f0, t = pw.harvest(waveform.squeeze(dim=0).numpy().astype('double'), sample_rate, frame_period=frame_period)
-        if sum(_f0 != 0) < 5:  # this happens when the algorithm fails
-            _f0, t = pw.dio(waveform.squeeze(dim=0).numpy().astype('double'), sample_rate, frame_period=frame_period)  # if harvest fails, try dio
-        f0 = pw.stonemask(waveform.squeeze(dim=0).numpy().astype('double'), _f0, t, sample_rate)
-        f0 = F.interpolate(torch.from_numpy(f0).view(1, 1, -1), size=sample['speech_feat'].shape[0], mode='linear').view(-1)
-        sample['pitch_feat'] = f0
-        yield sample
-
-
-def parse_embedding(data, normalize, mode='train'):
-    """ Parse utt_embedding/spk_embedding
-
-        Args:
-            data: Iterable[{key, wav, label, sample_rate}]
-
-        Returns:
-            Iterable[{key, feat, label}]
-    """
-    for sample in data:
-        sample['utt_embedding'] = torch.tensor(sample['utt_embedding'], dtype=torch.float32)
-        sample['spk_embedding'] = torch.tensor(sample['spk_embedding'], dtype=torch.float32)
-        if normalize:
-            sample['utt_embedding'] = F.normalize(sample['utt_embedding'], dim=0)
-            sample['spk_embedding'] = F.normalize(sample['spk_embedding'], dim=0)
-        yield sample
-
 
 def tokenize(data, get_tokenizer, allowed_special, mode='train'):
     """ Decode text to chars or BPE
@@ -505,12 +455,12 @@ def sort(data, sort_size=500, mode='train'):
     for sample in data:
         buf.append(sample)
         if len(buf) >= sort_size:
-            buf.sort(key=lambda x: x['speech_feat'].size(0))
+            buf.sort(key=lambda x: x['speech_latent'].size(0))
             for x in buf:
                 yield x
             buf = []
     # The sample left over
-    buf.sort(key=lambda x: x['speech_feat'].size(0))
+    buf.sort(key=lambda x: x['speech_latent'].size(0))
     for x in buf:
         yield x
 
@@ -549,9 +499,9 @@ def dynamic_batch(data, max_frames_in_batch=12000, mode='train'):
     buf = []
     longest_frames = 0
     for sample in data:
-        assert 'speech_feat' in sample
-        assert isinstance(sample['speech_feat'], torch.Tensor)
-        new_sample_frames = sample['speech_feat'].size(0)
+        assert 'speech_latent' in sample
+        assert isinstance(sample['speech_latent'], torch.Tensor)
+        new_sample_frames = sample['speech_latent'].size(0)
         longest_frames = max(longest_frames, new_sample_frames)
         frames_after_padding = longest_frames * (len(buf) + 1)
         if frames_after_padding > max_frames_in_batch:
@@ -574,7 +524,7 @@ def batch(data, batch_type='static', batch_size=16, max_frames_in_batch=12000, m
     else:
         logging.fatal('Unsupported batch type {}'.format(batch_type))
 
-def padding(data, use_spk_embedding, mode='train', gan=False, dpo=False, use_speaker_encoder=False):
+def padding(data, mode='train', gan=False, dpo=False, use_speaker_encoder=False):
     """ Padding the data into training data
 
         Args:
@@ -586,9 +536,9 @@ def padding(data, use_spk_embedding, mode='train', gan=False, dpo=False, use_spe
     """
     for sample in data:
         assert isinstance(sample, list)
-        speech_feat_len = torch.tensor([x['speech_feat'].size(0) for x in sample],  # Changed from size(1) to size(0)
+        speech_latent_len = torch.tensor([x['speech_latent'].size(0) for x in sample],  # Changed from size(1) to size(0)
                                        dtype=torch.int32)
-        order = torch.argsort(speech_feat_len, descending=True)
+        order = torch.argsort(speech_latent_len, descending=True)
 
         utts = [sample[i]['utt'] for i in order]
         speech = [sample[i]['speech'].squeeze(dim=0) for i in order]
@@ -607,17 +557,16 @@ def padding(data, use_spk_embedding, mode='train', gan=False, dpo=False, use_spe
                                     batch_first=True,
                                     padding_value=0)
         
-        speech_feat = [sample[i]['speech_feat'] for i in order]
-        speech_feat_len = torch.tensor([i.size(0) for i in speech_feat], dtype=torch.int32)
-        speech_feat = pad_sequence(speech_feat,
+        speech_latent = [sample[i]['speech_latent'] for i in order]
+
+        speech_latent = pad_sequence(speech_latent,
                                    batch_first=True,
                                    padding_value=0)
+
         text = [sample[i]['text'] for i in order]
         text_token = [torch.tensor(sample[i]['text_token']) for i in order]
         text_token_len = torch.tensor([i.size(0) for i in text_token], dtype=torch.int32)
         text_token = pad_sequence(text_token, batch_first=True, padding_value=0)
-        utt_embedding = torch.stack([sample[i]['utt_embedding'] for i in order], dim=0)
-        spk_embedding = torch.stack([sample[i]['spk_embedding'] for i in order], dim=0)
         
         batch = {
             "utts": utts,
@@ -625,13 +574,11 @@ def padding(data, use_spk_embedding, mode='train', gan=False, dpo=False, use_spe
             "speech_len": speech_len,
             "speech_token": speech_token,
             "speech_token_len": speech_token_len,
-            "speech_feat": speech_feat,
-            "speech_feat_len": speech_feat_len,
+            "speech_latent": speech_latent,
+            "speech_latent_len": speech_latent,
             "text": text,
             "text_token": text_token,
             "text_token_len": text_token_len,
-            "utt_embedding": utt_embedding,
-            "spk_embedding": spk_embedding,
         }
         
         # Handle reference mels for speaker encoder

speech/cosyvoice/utils/executor.py

@@ -78,6 +78,10 @@ class Executor:
             info_dict["epoch"] = self.epoch
             info_dict["batch_idx"] = batch_idx
 
+            for key, value in batch_dict.items():
+                if isinstance(value, torch.Tensor):
+                    print(f'{key} {value.shape}')
+
 
             if use_ddp and (batch_idx + 1) % info_dict["accum_grad"] != 0:
                 context = model.no_sync
@@ -86,6 +90,7 @@ class Executor:
 
 
             with context():
+                logger.info(f'{self.rank} batch_forward')
                 info_dict = batch_forward(
                     model,
                     batch_dict,
@@ -94,12 +99,13 @@ class Executor:
                     ref_model=self.ref_model,
                     dpo_loss=self.dpo_loss,
                 )
-
+                logger.info(f'{self.rank} batch_backward')
                 info_dict = batch_backward(model, scaler, info_dict)
-
+            logger.info(f'{self.rank} update_parameter_and_lr')
             info_dict = update_parameter_and_lr(
                 model, optimizer, scheduler, scaler, info_dict, model_type=model_type
             )
+            logger.info(f'{self.rank} log_per_step')
             log_per_step(experiment, info_dict)
 
             if (

speech/cosyvoice/utils/train_utils.py

@@ -312,7 +312,7 @@ def batch_forward(model, batch, scaler, info_dict, ref_model=None, dpo_loss=None
 
     with autocast:
         info_dict['loss_dict'] = model(batch, device)
-        # print('infor_dict loss_dict : ', info_dict['loss_dict'])
+        print('infor_dict loss_dict : ', info_dict['loss_dict'])
         if ref_model is not None and dpo_loss is not None:
             chosen_logps = info_dict['loss_dict']["chosen_logps"]
             rejected_logps = info_dict['loss_dict']["rejected_logps"]

speech/tools/generate_json_index.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python3
+"""
+Generate JSON index file for dataset
+This creates a JSON file with all valid audio files and their metadata
+"""
+
+import os
+import json
+import glob
+import argparse
+from pathlib import Path
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import hashlib
+from datetime import datetime
+
+def validate_file_set(wav_path):
+    """Check if all required files exist for a given wav file"""
+    txt_path = wav_path.replace('.wav', '.txt')
+    token_path = wav_path.replace('.wav', '_fsq.pt')
+    latent_path = wav_path.replace('.wav', '_latent.pt')
+    
+    # Check all files exist
+    if not all(os.path.exists(p) for p in [wav_path, txt_path, token_path, latent_path]):
+        return None
+    
+    # Get file sizes for validation
+    try:
+        wav_size = os.path.getsize(wav_path)
+        txt_size = os.path.getsize(txt_path)
+        token_size = os.path.getsize(token_path)
+        latent_size = os.path.getsize(latent_path)
+        
+        # Skip if any file is empty
+        if any(size == 0 for size in [wav_size, txt_size, token_size, latent_size]):
+            return None
+        
+        # Extract metadata
+        utt = os.path.basename(wav_path).replace('.wav', '')
+        spk = utt.split('_')[0] if '_' in utt else 'default'
+        
+        # Get file modification time
+        mtime = os.path.getmtime(wav_path)
+        
+        return {
+            'utt': utt,
+            'spk': spk,
+            'wav': wav_path,
+            'text_path': txt_path,
+            'token_path': token_path,
+            'latent_path': latent_path,
+            'wav_size': wav_size,
+            'txt_size': txt_size,
+            'token_size': token_size,
+            'latent_size': latent_size,
+            'mtime': mtime,
+        }
+    except Exception as e:
+        print(f"Error processing {wav_path}: {e}")
+        return None
+
+def process_directory(directory, pattern='**/*.wav'):
+    """Process a directory and find all valid audio files"""
+    print(f"Scanning directory: {directory}")
+    wav_files = glob.glob(os.path.join(directory, pattern), recursive=True)
+    print(f"Found {len(wav_files)} wav files")
+    
+    valid_files = []
+    
+    # Process files in parallel
+    with ThreadPoolExecutor(max_workers=16) as executor:
+        futures = [executor.submit(validate_file_set, wav_path) for wav_path in wav_files]
+        
+        for future in tqdm(as_completed(futures), total=len(futures), desc="Validating files"):
+            result = future.result()
+            if result is not None:
+                valid_files.append(result)
+    
+    return valid_files
+
+def process_files_txt(files_txt):
+    """Process files from a text file list"""
+    print(f"Reading file list from: {files_txt}")
+    
+    with open(files_txt, 'r') as f:
+        wav_files = [line.strip() for line in f if line.strip() and not line.startswith('#')]
+    
+    print(f"Found {len(wav_files)} files in list")
+    
+    valid_files = []
+    
+    # Process files in parallel
+    with ThreadPoolExecutor(max_workers=16) as executor:
+        futures = [executor.submit(validate_file_set, wav_path) for wav_path in wav_files]
+        
+        for future in tqdm(as_completed(futures), total=len(futures), desc="Validating files"):
+            result = future.result()
+            if result is not None:
+                valid_files.append(result)
+    
+    return valid_files
+
+def generate_statistics(file_list):
+    """Generate statistics about the dataset"""
+    stats = {
+        'total_files': len(file_list),
+        'total_size_gb': sum(f['wav_size'] + f['txt_size'] + f['token_size'] + f['latent_size'] 
+                           for f in file_list) / (1024**3),
+        'speakers': {},
+        'file_sizes': {
+            'wav_total_gb': sum(f['wav_size'] for f in file_list) / (1024**3),
+            'txt_total_mb': sum(f['txt_size'] for f in file_list) / (1024**2),
+            'token_total_gb': sum(f['token_size'] for f in file_list) / (1024**3),
+            'latent_total_gb': sum(f['latent_size'] for f in file_list) / (1024**3),
+        }
+    }
+    
+    # Count files per speaker
+    for f in file_list:
+        spk = f['spk']
+        if spk not in stats['speakers']:
+            stats['speakers'][spk] = 0
+        stats['speakers'][spk] += 1
+    
+    stats['num_speakers'] = len(stats['speakers'])
+    
+    return stats
+
+def generate_json_index(input_paths, output_file, split_ratio=None):
+    """
+    Generate JSON index file from input paths
+    
+    Args:
+        input_paths: List of paths (directories or files.txt)
+        output_file: Output JSON file path
+        split_ratio: Optional tuple (train_ratio, val_ratio, test_ratio)
+    """
+    all_files = []
+    
+    # Process each input path
+    for path in input_paths:
+        if os.path.isdir(path):
+            files = process_directory(path)
+        elif path.endswith('.txt'):
+            files = process_files_txt(path)
+        else:
+            print(f"Warning: Skipping unknown input type: {path}")
+            continue
+        
+        all_files.extend(files)
+    
+    # Remove duplicates based on utterance ID
+    unique_files = {}
+    for f in all_files:
+        utt = f['utt']
+        if utt not in unique_files:
+            unique_files[utt] = f
+        else:
+            # Keep the one with newer modification time
+            if f['mtime'] > unique_files[utt]['mtime']:
+                unique_files[utt] = f
+    
+    file_list = list(unique_files.values())
+    
+    # Sort by utterance ID for consistency
+    file_list.sort(key=lambda x: x['utt'])
+    
+    print(f"\nTotal unique files: {len(file_list)}")
+    
+    # Generate statistics
+    stats = generate_statistics(file_list)
+    
+    # Create index
+    index = {
+        'version': '1.0',
+        'created': datetime.now().isoformat(),
+        'statistics': stats,
+        'data': file_list
+    }
+    
+    # Optional: Create train/val/test splits
+    if split_ratio:
+        import random
+        random.seed(42)  # For reproducibility
+        
+        train_ratio, val_ratio, test_ratio = split_ratio
+        assert abs(train_ratio + val_ratio + test_ratio - 1.0) < 0.001, "Split ratios must sum to 1"
+        
+        # Shuffle for random split
+        shuffled = file_list.copy()
+        random.shuffle(shuffled)
+        
+        n = len(shuffled)
+        train_end = int(n * train_ratio)
+        val_end = int(n * (train_ratio + val_ratio))
+        
+        splits = {
+            'train': shuffled[:train_end],
+            'val': shuffled[train_end:val_end],
+            'test': shuffled[val_end:]
+        }
+        
+        # Save separate files for each split
+        base_name = output_file.replace('.json', '')
+        
+        for split_name, split_data in splits.items():
+            split_index = {
+                'version': '1.0',
+                'created': datetime.now().isoformat(),
+                'split': split_name,
+                'statistics': generate_statistics(split_data),
+                'data': split_data
+            }
+            
+            split_file = f"{base_name}_{split_name}.json"
+            with open(split_file, 'w') as f:
+                json.dump(split_index, f, indent=2)
+            
+            print(f"Saved {split_name} split: {split_file} ({len(split_data)} files)")
+    
+    # Save main index
+    with open(output_file, 'w') as f:
+        json.dump(index, f, indent=2)
+    
+    print(f"\nSaved index to: {output_file}")
+    print(f"Total files: {stats['total_files']}")
+    print(f"Total size: {stats['total_size_gb']:.2f} GB")
+    print(f"Number of speakers: {stats['num_speakers']}")
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate JSON index for dataset")
+    parser.add_argument('--input', nargs='+', required=True,
+                      help='Input paths (directories or files.txt)')
+    parser.add_argument('--output', default='dataset_index.json',
+                      help='Output JSON file (default: dataset_index.json)')
+    parser.add_argument('--split', nargs=3, type=float, metavar=('TRAIN', 'VAL', 'TEST'),
+                      help='Create train/val/test splits (e.g., --split 0.8 0.1 0.1)')
+    
+    args = parser.parse_args()
+    
+    # Validate split ratios if provided
+    split_ratio = None
+    if args.split:
+        split_ratio = tuple(args.split)
+        if abs(sum(split_ratio) - 1.0) > 0.001:
+            parser.error("Split ratios must sum to 1.0")
+    
+    generate_json_index(args.input, args.output, split_ratio)
+
+if __name__ == "__main__":
+    main()
+
+# Example usage:
+# python generate_json_index.py --input /data/dataset/emilia /data/dataset/vivoice --output dataset_index.json
+# python generate_json_index.py --input train_files.txt --output train_index.json
+# python generate_json_index.py --input /data/dataset/emilia --output dataset_index.json --split 0.8 0.1 0.1