inv code

speech/tools/create_data_list.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+"""Create data list files for training with individual files"""
+
+import argparse
+import os
+import json
+
+def create_data_lists(src_dir, output_dir):
+    """Create data list files pointing to directories or index files
+    
+    Args:
+        src_dir: Directory containing processed audio files
+        output_dir: Directory to save list files
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Option 1: Create a list pointing to the source directory
+    with open(os.path.join(output_dir, 'data.list'), 'w') as f:
+        f.write(src_dir + '\n')
+    
+    # Option 2: If you have an index file, point to it
+    index_file = os.path.join(src_dir, 'data_index.json')
+    if os.path.exists(index_file):
+        with open(os.path.join(output_dir, 'data_index.list'), 'w') as f:
+            f.write(index_file + '\n')
+    
+    print(f"Created data lists in {output_dir}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--src_dir', type=str, required=True,
+                       help='Source directory with processed files')
+    parser.add_argument('--output_dir', type=str, required=True,
+                       help='Output directory for list files')
+    args = parser.parse_args()
+    
+    create_data_lists(args.src_dir, args.output_dir)

speech/tools/extract_embedding.py

@@ -1,17 +1,4 @@
 #!/usr/bin/env python3
-# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu)
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 import argparse
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import onnxruntime
@@ -19,54 +6,102 @@ import torch
 import torchaudio
 import torchaudio.compliance.kaldi as kaldi
 from tqdm import tqdm
+import os
+import glob
+import logging
 
+logger = logging.getLogger()
 
-def single_job(utt):
-    audio, sample_rate = torchaudio.load(utt2wav[utt])
+
+def process_single_audio(wav_path):
+    # Extract utterance ID and speaker ID from filename
+    utt = os.path.basename(wav_path).replace('.wav', '')
+    spk = utt.split('_')[0]
+    
+    # Check if text file exists
+    txt_path = wav_path.replace('.wav', '.normalized.txt')
+    if not os.path.exists(txt_path):
+        logger.warning(f'{txt_path} does not exist, skipping {wav_path}')
+        return None
+    
+    # Process audio
+    audio, sample_rate = torchaudio.load(wav_path)
     if sample_rate != 16000:
         audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(audio)
+    
     feat = kaldi.fbank(audio,
                        num_mel_bins=80,
                        dither=0,
                        sample_frequency=16000)
     feat = feat - feat.mean(dim=0, keepdim=True)
-    embedding = ort_session.run(None, {ort_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten().tolist()
-    return utt, embedding
+    
+    # Generate embedding
+    embedding = ort_session.run(None, {ort_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten()
+    
+    # Save individual embedding file
+    embedding_path = wav_path.replace('.wav', '_embedding.pt')
+    torch.save(embedding, embedding_path)
+    
+    return {
+        'wav_path': wav_path,
+        'utt': utt,
+        'spk': spk,
+        'embedding': embedding,
+        'embedding_path': embedding_path
+    }
 
 
 def main(args):
-    all_task = [executor.submit(single_job, utt) for utt in utt2wav.keys()]
-    utt2embedding, spk2embedding = {}, {}
-    for future in tqdm(as_completed(all_task)):
-        utt, embedding = future.result()
-        utt2embedding[utt] = embedding
-        spk = utt2spk[utt]
-        if spk not in spk2embedding:
-            spk2embedding[spk] = []
-        spk2embedding[spk].append(embedding)
-    for k, v in spk2embedding.items():
-        spk2embedding[k] = torch.tensor(v).mean(dim=0).tolist()
-    torch.save(utt2embedding, "{}/utt2embedding.pt".format(args.dir))
-    torch.save(spk2embedding, "{}/spk2embedding.pt".format(args.dir))
+    # Find all wav files
+    wav_files = list(glob.glob('{}/*/*/*wav'.format(args.src_dir)))
+    print(f"Found {len(wav_files)} wav files")
+    
+    # Process all audio files
+    all_tasks = [executor.submit(process_single_audio, wav_path) for wav_path in wav_files]
+    
+    # Collect results
+    spk2embeddings = {}
+    successful_files = []
+    
+    for future in tqdm(as_completed(all_tasks), total=len(all_tasks)):
+        result = future.result()
+        if result is None:
+            continue
+            
+        successful_files.append(result)
+        
+        # Collect embeddings by speaker
+        spk = result['spk']
+        if spk not in spk2embeddings:
+            spk2embeddings[spk] = []
+        spk2embeddings[spk].append(result['embedding'])
+    
+    # Calculate and save speaker embeddings
+    spk_embed_dir = os.path.join(args.src_dir, "spk_embeddings")
+    os.makedirs(spk_embed_dir, exist_ok=True)
+    
+    for spk, embeddings in spk2embeddings.items():
+        spk_embedding = torch.stack([torch.tensor(e) for e in embeddings]).mean(dim=0)
+        spk_embedding_path = os.path.join(spk_embed_dir, f"{spk}_embedding.pt")
+        torch.save(spk_embedding, spk_embedding_path)
+        print(f"Saved speaker embedding for {spk} with {len(embeddings)} utterances")
+    
+    # Save a summary file for reference
+    summary_path = os.path.join(args.src_dir, "embedding_summary.txt")
+    with open(summary_path, 'w') as f:
+        f.write(f"Processed {len(successful_files)} files successfully\n")
+        f.write(f"Found {len(spk2embeddings)} speakers\n")
+        for result in successful_files:
+            f.write(f"{result['utt']} {result['wav_path']} {result['embedding_path']}\n")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--dir", type=str)
-    parser.add_argument("--onnx_path", type=str)
+    parser.add_argument("--src_dir", type=str, help="Source directory containing audio files")
+    parser.add_argument("--onnx_path", type=str, help="Path to campplus.onnx model")
     parser.add_argument("--num_thread", type=int, default=8)
     args = parser.parse_args()
 
-    utt2wav, utt2spk = {}, {}
-    with open('{}/wav.scp'.format(args.dir)) as f:
-        for l in f:
-            l = l.replace('\n', '').split()
-            utt2wav[l[0]] = l[1]
-    with open('{}/utt2spk'.format(args.dir)) as f:
-        for l in f:
-            l = l.replace('\n', '').split()
-            utt2spk[l[0]] = l[1]
-
     option = onnxruntime.SessionOptions()
     option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
     option.intra_op_num_threads = 1
@@ -74,4 +109,4 @@ if __name__ == "__main__":
     ort_session = onnxruntime.InferenceSession(args.onnx_path, sess_options=option, providers=providers)
     executor = ThreadPoolExecutor(max_workers=args.num_thread)
 
-    main(args)
+    main(args)

speech/tools/extract_speech_token.py

@@ -1,17 +1,4 @@
 #!/usr/bin/env python3
-# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu)
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 import argparse
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import logging
@@ -21,47 +8,87 @@ import onnxruntime
 import numpy as np
 import torchaudio
 import whisper
+import glob
+import os
 
+logger = logging.getLogger()
 
-def single_job(utt):
-    audio, sample_rate = torchaudio.load(utt2wav[utt], backend='soundfile')
+
+def process_single_audio(wav_path):
+    # Check if text file exists
+    txt_path = wav_path.replace('.wav', '.normalized.txt')
+    if not os.path.exists(txt_path):
+        logger.warning(f'{txt_path} does not exist, skipping {wav_path}')
+        return None
+    
+    # Extract utterance ID
+    utt = os.path.basename(wav_path).replace('.wav', '')
+    
+    # Process audio
+    audio, sample_rate = torchaudio.load(wav_path, backend='soundfile')
     if sample_rate != 16000:
         audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(audio)
+    
     # Convert audio to mono
     if audio.shape[0] > 1:
         audio = audio.mean(dim=0, keepdim=True)
+    
     if audio.shape[1] / 16000 > 30:
-        logging.warning('do not support extract speech token for audio longer than 30s')
+        logging.warning(f'Audio longer than 30s, skipping tokenization for {wav_path}')
         speech_token = []
     else:
         feat = whisper.log_mel_spectrogram(audio, n_mels=128)
-        speech_token = ort_session.run(None, {ort_session.get_inputs()[0].name: feat.detach().cpu().numpy(),
-                                              ort_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32)})[0].flatten().tolist()
-    return utt, speech_token
+        speech_token = ort_session.run(None, {
+            ort_session.get_inputs()[0].name: feat.detach().cpu().numpy(),
+            ort_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32)
+        })[0].flatten().tolist()
+    
+    # Save individual token file
+    token_path = wav_path.replace('.wav', '_tokens.pt')
+    torch.save(speech_token, token_path)
+    
+    return {
+        'wav_path': wav_path,
+        'utt': utt,
+        'token_path': token_path,
+        'num_tokens': len(speech_token)
+    }
 
 
 def main(args):
-    all_task = [executor.submit(single_job, utt) for utt in utt2wav.keys()]
-    utt2speech_token = {}
-    for future in tqdm(as_completed(all_task)):
-        utt, speech_token = future.result()
-        utt2speech_token[utt] = speech_token
-    torch.save(utt2speech_token, '{}/utt2speech_token.pt'.format(args.dir))
+    # Find all wav files
+    wav_files = list(glob.glob('{}/*/*/*wav'.format(args.src_dir)))
+    print(f"Found {len(wav_files)} wav files")
+    
+    # Process all audio files
+    all_tasks = [executor.submit(process_single_audio, wav_path) for wav_path in wav_files]
+    
+    # Collect results
+    successful_files = []
+    
+    for future in tqdm(as_completed(all_tasks), total=len(all_tasks)):
+        result = future.result()
+        if result is None:
+            continue
+        successful_files.append(result)
+    
+    # Save a summary file for reference
+    summary_path = os.path.join(args.src_dir, "token_summary.txt")
+    with open(summary_path, 'w') as f:
+        f.write(f"Processed {len(successful_files)} files successfully\n")
+        total_tokens = sum(r['num_tokens'] for r in successful_files)
+        f.write(f"Total tokens generated: {total_tokens}\n")
+        for result in successful_files:
+            f.write(f"{result['utt']} {result['wav_path']} {result['token_path']} {result['num_tokens']}\n")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--dir", type=str)
-    parser.add_argument("--onnx_path", type=str)
+    parser.add_argument("--src_dir", type=str, help="Source directory containing audio files")
+    parser.add_argument("--onnx_path", type=str, help="Path to speech_tokenizer_v2.onnx model")
     parser.add_argument("--num_thread", type=int, default=8)
     args = parser.parse_args()
 
-    utt2wav = {}
-    with open('{}/wav.scp'.format(args.dir)) as f:
-        for l in f:
-            l = l.replace('\n', '').split()
-            utt2wav[l[0]] = l[1]
-
     option = onnxruntime.SessionOptions()
     option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
     option.intra_op_num_threads = 1
@@ -69,4 +96,4 @@ if __name__ == "__main__":
     ort_session = onnxruntime.InferenceSession(args.onnx_path, sess_options=option, providers=providers)
     executor = ThreadPoolExecutor(max_workers=args.num_thread)
 
-    main(args)
+    main(args)

speech/tools/inv_file_processor.py

@@ -0,0 +1,109 @@
+# Add this to your processor.py file or create a new file
+
+import logging
+import json
+import torch
+import glob
+import os
+from pathlib import Path
+
+def individual_file_opener(data, mode='train', tts_data={}):
+    """Load data from individual files instead of parquet
+    
+    Args:
+        data: Iterable[{src}] where src is either:
+            - Path to index JSON file
+            - Directory path containing wav files
+        mode: 'train' or 'test'
+        tts_data: Dict for TTS mode
+        
+    Yields:
+        Dict with all required fields for training
+    """
+    for sample in data:
+        src = sample['src']
+        
+        # Check if src is a JSON index file or a directory
+        if src.endswith('.json'):
+            # Load from index file
+            with open(src, 'r') as f:
+                index_data = json.load(f)
+                file_list = index_data.get('data', [])
+        else:
+            # Find all wav files in directory
+            wav_files = glob.glob(os.path.join(src, '*/*/*wav'))
+            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')
+                
+                if not all(os.path.exists(p) for p in [txt_path, embedding_path, token_path]):
+                    logging.warning(f'Missing files for {wav_path}, skipping')
+                    continue
+                
+                # Extract metadata
+                utt = os.path.basename(wav_path).replace('.wav', '')
+                spk = utt.split('_')[0]
+                
+                file_list.append({
+                    'utt': utt,
+                    'spk': spk,
+                    'wav': wav_path,
+                    'text_path': txt_path,
+                    'embedding_path': embedding_path,
+                    'token_path': token_path,
+                    'spk_embedding_path': os.path.join(os.path.dirname(src), f"spk_embeddings/{spk}_embedding.pt")
+                })
+        
+        # Process each file
+        for file_info in file_list:
+            try:
+                # Read audio data
+                with open(file_info['wav'], 'rb') as f:
+                    audio_data = f.read()
+                
+                # Read text
+                with open(file_info['text_path'], 'r') as f:
+                    text = ''.join(l.strip() for l in f.readlines())
+                
+                # Load embeddings
+                utt_embedding = torch.load(file_info['embedding_path']).tolist()
+                speech_token = torch.load(file_info['token_path'])
+                
+                # Load speaker embedding
+                if os.path.exists(file_info['spk_embedding_path']):
+                    spk_embedding = torch.load(file_info['spk_embedding_path']).tolist()
+                else:
+                    logging.warning(f"Speaker embedding not found: {file_info['spk_embedding_path']}")
+                    spk_embedding = utt_embedding  # Fallback to utterance embedding
+                
+                # Build sample dict
+                sample_dict = {
+                    'utt': file_info['utt'],
+                    'spk': file_info['spk'],
+                    '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
+                }
+                
+                # Merge with original sample data
+                sample_dict.update(sample)
+                
+                if mode == 'train':
+                    yield sample_dict
+                else:
+                    # For TTS mode
+                    if file_info['utt'] in tts_data:
+                        for index, tts_text in enumerate(tts_data[file_info['utt']]):
+                            yield {**sample_dict, 'tts_index': index, 'tts_text': tts_text}
+                    else:
+                        yield sample_dict
+                        
+            except Exception as ex:
+                logging.warning(f'Failed to process {file_info["wav"]}: {ex}')