More fixes.

Added sanitize_dataset.py to clean CSV files.

.gitignore

@@ -33,12 +33,15 @@ test.py
 *.json
 *.jsonl
 code_digest.txt
+uv.lock
 
 # =========================
 # Data, Logs, & Outputs
 # =========================
 wandb/
 logs/
+dataset/
+datasets/
 *.wav
 *.flac
 *.mp3

codec_train.py

@@ -13,7 +13,7 @@ from codec_dataset import LJSpeechDataset
 from codec.codec_decoder.decoder import SimpleDecoder
 
 # Import the config loader
-from config_loader import load_config
+from utils.config_loader import load_config
 
 def pad_collate(batch):
     """

generate_dataset.py

@@ -1,25 +1,29 @@
 """
-Converts a dataset in LJSpeech format into audio tokens for Soprano, using pre-defined train/val lists.
+Converts a dataset in LJSpeech format into audio tokens for Soprano.
+This script creates two JSON files for train and test splits in the provided directory.
 
 Usage:
-python generate_dataset_from_lists.py
+python generate_dataset.py
 """
 import pathlib
 import json
 import os
+import random
 import torch
 from tqdm import tqdm
-from encoder.codec import Encoder
+from huggingface_hub import hf_hub_download
+from codec.encoder.codec import Encoder
 
-from config_loader import load_config
+from utils.config_loader import load_config
 from utils.audio_utils import AudioPipeline
 
 def load_metadata(input_dir):
     print("Reading metadata...")
     meta_map = {}
-    meta_path = input_dir / 'metadata_orig.csv'
+    meta_path = input_dir / 'metadata.csv'
+    
     if not meta_path.exists():
-        meta_path = input_dir / 'metadata.csv'
+        raise FileNotFoundError(f"Could not find {meta_path}. Did you run sanitize.py?")
     
     with open(meta_path, encoding='utf-8') as f:
         for line in f:
@@ -30,92 +34,87 @@ def load_metadata(input_dir):
             meta_map[filename] = transcript
     return meta_map
 
-def process_list(list_file, meta_map, encoder, target_sr):
-    dataset = []
-    print(f"Processing {list_file}...")
-    with open(list_file, 'r') as f:
-        lines = [l.strip() for l in f if l.strip()]
-    
-    for line in tqdm(lines):
-        path_obj = pathlib.Path(line)
-        filename = path_obj.stem # LJxxx
-        
-        if filename not in meta_map:
-            print(f"Warning: {filename} not found in metadata. Skipping.")
-            continue
-            
-        transcript = meta_map[filename]
-        wav_path = str(path_obj)
-        
-        # Load and Encode with OS-aware pipeline
-        try:
-            audio, _ = AudioPipeline.load_audio(wav_path, target_sr)
-        except Exception as e:
-            print(f"Error loading {wav_path}: {e}")
-            continue
-            
-        with torch.no_grad():
-            audio_tokens = encoder(audio) 
-
-        dataset.append([transcript, audio_tokens.squeeze(0).tolist(), wav_path])
-        
-    return dataset
-
 def main():
     config = load_config("config.yaml")
     cfg_paths = config["paths"]
     cfg_codec = config["codec"]
+    cfg_data = config["data_generation"]
     
     input_dir = pathlib.Path(cfg_paths["dataset_root"])
-    output_dir = pathlib.Path(cfg_paths["save_dir"]) / "dataset_lists"
-    os.makedirs(output_dir, exist_ok=True)
-    
     target_sr = cfg_codec["sample_rate"]
     device = config["global"]["device"] if torch.cuda.is_available() else 'cpu'
+    seed = config["global"]["seed"]
 
     # Load Encoder
     print("Loading Encoder...")
     encoder = Encoder()
     speech_autoencoder_path = cfg_paths["pretrained_codec_path"]
     
-    if not speech_autoencoder_path or not os.path.exists(speech_autoencoder_path):
-        raise FileNotFoundError(f"pretrained_codec_path not found: {speech_autoencoder_path}")
+    if speech_autoencoder_path and os.path.exists(speech_autoencoder_path):
+        print(f"Loading custom weights from {speech_autoencoder_path}...")
+        full_ckpt = torch.load(speech_autoencoder_path, map_location='cpu')
         
-    print(f"Loading weights from {speech_autoencoder_path}")
-    full_ckpt = torch.load(speech_autoencoder_path, map_location='cpu')
+        encoder_state_dict = {}
+        for k, v in full_ckpt.items():
+            if k.startswith("encoder."):
+                new_k = k.replace("encoder.", "", 1)
+                encoder_state_dict[new_k] = v
                 
-    encoder_state_dict = {}
-    for k, v in full_ckpt.items():
-        if k.startswith("encoder."):
-            new_k = k.replace("encoder.", "", 1)
-            encoder_state_dict[new_k] = v
+        encoder.load_state_dict(encoder_state_dict)
+    else:
+        print("No custom codec path found in config. Downloading default Soprano-Encoder from Hugging Face...")
+        encoder_path = hf_hub_download(repo_id='ekwek/Soprano-Encoder', filename='encoder.pth')
+        encoder.load_state_dict(torch.load(encoder_path, map_location='cpu'))
         
-    encoder.load_state_dict(encoder_state_dict)
     encoder.to(device)
     encoder.eval()
     print("Encoder Loaded.")
 
     meta_map = load_metadata(input_dir)
 
-    # Process Train List
-    train_list_path = input_dir / 'train_list.txt'
-    if train_list_path.exists():
-        train_data = process_list(train_list_path, meta_map, encoder, target_sr)
-        with open(output_dir / 'train.json', 'w') as f:
-            json.dump(train_data, f, indent=2)
-        print(f"Saved {len(train_data)} train samples to {output_dir}/train.json")
-    else:
-        print(f"Error: {train_list_path} not found.")
+    print("Encoding audio...")
+    dataset = []
     
-    # Process Val List
-    val_list_path = input_dir / 'val_list.txt'
-    if val_list_path.exists():
-        val_data = process_list(val_list_path, meta_map, encoder, target_sr)
-        with open(output_dir / 'val.json', 'w') as f:
-            json.dump(val_data, f, indent=2)
-        print(f"Saved {len(val_data)} val samples to {output_dir}/val.json")
-    else:
-        print(f"Error: {val_list_path} not found.")
+    # Process all files found in the metadata
+    for filename, transcript in tqdm(meta_map.items()):
+        wav_path = input_dir / 'wavs' / f'{filename}.wav'
+        
+        if not wav_path.exists():
+            print(f"Warning: {wav_path} not found. Skipping.")
+            continue
+            
+        # Load and Encode with OS-aware pipeline
+        try:
+            audio, _ = AudioPipeline.load_audio(str(wav_path), target_sr)
+        except Exception as e:
+            print(f"Error loading {wav_path}: {e}")
+            continue
+            
+        audio = audio.to(device)
+        
+        with torch.no_grad():
+            audio_tokens = encoder(audio) 
+
+        dataset.append([transcript, audio_tokens.squeeze(0).tolist(), str(wav_path.resolve())])
+
+    print("Generating train/test splits...")
+    random.seed(seed)
+    random.shuffle(dataset)
+    num_val = min(int(cfg_data["val_prop"] * len(dataset)) + 1, cfg_data["val_max"])
+    
+    train_dataset = dataset[num_val:]
+    val_dataset = dataset[:num_val]
+    
+    print(f'# train samples: {len(train_dataset)}')
+    print(f'# val samples: {len(val_dataset)}')
+
+    print("Saving datasets...")
+    with open(input_dir / 'train.json', 'w', encoding='utf-8') as f:
+        json.dump(train_dataset, f, indent=2)
+    with open(input_dir / 'val.json', 'w', encoding='utf-8') as f:
+        json.dump(val_dataset, f, indent=2)
+        
+    print("Datasets saved successfully.")
 
 if __name__ == '__main__':
     main()

generate_dataset_from_lists.py

@@ -9,17 +9,19 @@ import json
 import os
 import torch
 from tqdm import tqdm
-from encoder.codec import Encoder
+from huggingface_hub import hf_hub_download
+from codec.encoder.codec import Encoder
 
-from config_loader import load_config
+from utils.config_loader import load_config
 from utils.audio_utils import AudioPipeline
 
 def load_metadata(input_dir):
     print("Reading metadata...")
     meta_map = {}
-    meta_path = input_dir / 'metadata_orig.csv'
+    meta_path = input_dir / 'metadata.csv'
+    
     if not meta_path.exists():
-        meta_path = input_dir / 'metadata.csv'
+        raise FileNotFoundError(f"Could not find {meta_path}. Did you run sanitize.py?")
     
     with open(meta_path, encoding='utf-8') as f:
         for line in f:
@@ -82,19 +84,22 @@ def main():
     encoder = Encoder()
     speech_autoencoder_path = cfg_paths["pretrained_codec_path"]
     
-    if not speech_autoencoder_path or not os.path.exists(speech_autoencoder_path):
-        raise FileNotFoundError(f"pretrained_codec_path not found: {speech_autoencoder_path}")
+    if speech_autoencoder_path and os.path.exists(speech_autoencoder_path):
+        print(f"Loading custom weights from {speech_autoencoder_path}")
+        full_ckpt = torch.load(speech_autoencoder_path, map_location='cpu')
         
-    print(f"Loading weights from {speech_autoencoder_path}")
-    full_ckpt = torch.load(speech_autoencoder_path, map_location='cpu')
+        encoder_state_dict = {}
+        for k, v in full_ckpt.items():
+            if k.startswith("encoder."):
+                new_k = k.replace("encoder.", "", 1)
+                encoder_state_dict[new_k] = v
                 
-    encoder_state_dict = {}
-    for k, v in full_ckpt.items():
-        if k.startswith("encoder."):
-            new_k = k.replace("encoder.", "", 1)
-            encoder_state_dict[new_k] = v
+        encoder.load_state_dict(encoder_state_dict)
+    else:
+        print("No custom codec path found in config. Downloading default Soprano-Encoder from Hugging Face...")
+        encoder_path = hf_hub_download(repo_id='ekwek/Soprano-Encoder', filename='encoder.pth')
+        encoder.load_state_dict(torch.load(encoder_path, map_location='cpu'))
         
-    encoder.load_state_dict(encoder_state_dict)
     encoder.to(device)
     encoder.eval()
     print("Encoder Loaded.")
@@ -109,7 +114,7 @@ def main():
             json.dump(train_data, f, indent=2)
         print(f"Saved {len(train_data)} train samples to {output_dir}/train.json")
     else:
-        print(f"Error: {train_list_path} not found.")
+        print(f"Error: {train_list_path} not found. Skipping train list generation.")
 
     # Process Val List
     val_list_path = input_dir / 'val_list.txt'
@@ -119,7 +124,7 @@ def main():
             json.dump(val_data, f, indent=2)
         print(f"Saved {len(val_data)} val samples to {output_dir}/val.json")
     else:
-        print(f"Error: {val_list_path} not found.")
+        print(f"Error: {val_list_path} not found. Skipping val list generation.")
 
 if __name__ == '__main__':
     main()

sanitize_dataset.py

@@ -0,0 +1,77 @@
+import re
+import shutil
+from pathlib import Path
+from utils.config_loader import load_config
+
+def clean_text(text):
+    """Sanitizes text for TTS training."""
+    # Replace weird curly quotes with standard straight quotes
+    text = text.replace('“', '"').replace('”', '"')
+    text = text.replace('‘', "'").replace('’', "'")
+    
+    # Replace em-dashes with standard dashes
+    text = text.replace('—', '-')
+    
+    # Remove leading/trailing whitespace
+    text = text.strip()
+    
+    # Collapse multiple spaces/tabs into a single space
+    text = re.sub(r'\s+', ' ', text)
+    
+    return text
+
+def main():
+    config = load_config("config.yaml")
+    dataset_dir = Path(config["paths"]["dataset_root"])
+    
+    input_csv = dataset_dir / "metadata.csv"
+    backup_csv = dataset_dir / "metadata.csv.bak"
+    
+    if not input_csv.exists():
+        print(f"Error: Could not find {input_csv}. Please check your config.yaml.")
+        return
+
+    print(f"Reading and sanitizing {input_csv}...")
+    
+    clean_lines = []
+    skipped = 0
+
+    # 1. Read and clean the data in memory first
+    with open(input_csv, "r", encoding="utf-8") as f:
+        for line_num, line in enumerate(f):
+            line = line.strip()
+            if not line:
+                continue
+            
+            # Split by pipe
+            parts = line.split("|")
+            
+            if len(parts) < 2:
+                print(f"Skipping line {line_num + 1} (not enough columns): {line}")
+                skipped += 1
+                continue
+                
+            filename = parts[0].strip()
+            
+            # Grab the transcript (we take parts[1] so we ignore the duplicate 3rd column if it exists)
+            raw_transcript = parts[1]
+            
+            # Clean the text
+            transcript = clean_text(raw_transcript)
+            
+            # Reformat to strict 2-column: filename|transcript
+            clean_lines.append(f"{filename}|{transcript}")
+
+    # 2. Create the backup
+    print(f"Creating backup at {backup_csv}...")
+    shutil.copy2(input_csv, backup_csv)
+
+    # 3. Overwrite the original file with the clean data
+    print(f"Overwriting {input_csv} with clean data...")
+    with open(input_csv, "w", encoding="utf-8") as f:
+        f.write("\n".join(clean_lines) + "\n")
+        
+    print(f"Done! Successfully processed {len(clean_lines)} lines. Skipped {skipped} invalid lines.")
+
+if __name__ == "__main__":
+    main()

simple_inference.py

@@ -7,7 +7,7 @@ from safetensors.torch import load_file
 
 # Ensure decoder module is importable
 from decoder.decoder import SopranoDecoder
-from config_loader import load_config
+from utils.config_loader import load_config
 
 def load_models(llm_path, decoder_path, device='cuda'):
     if not llm_path or not os.path.exists(llm_path):

train_decoder.py

@@ -22,7 +22,7 @@ from decoder.decoder import SopranoDecoder
 from decoder.discriminator import Discriminator
 from decoder.losses import MelSpectrogramWrapper, feature_matching_loss, discriminator_loss, generator_loss, MultiResolutionSTFTLoss
 
-from config_loader import load_config
+from utils.config_loader import load_config
 
 
 def worker_seed_init(_):

train_llm.py

@@ -17,7 +17,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 from safetensors.torch import load_file
 
 from dataset import AudioDataset
-from config_loader import load_config
+from utils.config_loader import load_config
 
 
 def worker_seed_init(_):