File size: 3,072 Bytes
94336a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from dataclasses import dataclass, field
from typing import List
@dataclass
class TrainConfig:
# --- Paths ---
# Directory where setup.py downloaded the files
model_dir: str = "./pretrained_models"
# Path to your metadata CSV (Format: ID|RawText|NormText)
csv_path: str = "./MyTTSDataset/metadata.csv"
#metadata_path: str = "./metadata.json"
# Directory containing WAV files
wav_dir: str = "./MyTTSDataset/wavs"
#wav_dir: str = "./FileBasedDataset"
preprocessed_dir = "./MyTTSDataset/preprocess"
#preprocessed_dir = "./FileBasedDataset/preprocess"
# Output directory for the finetuned model
output_dir: str = "./chatterbox_output"
is_inference = False
inference_prompt_path: str = "./speaker_reference/reference.wav"
# Reusing the Hindi "[hi]" tag for Nepali (see prepare_dataset.py's docstring for why)
# -- must match what preprocess_ljspeech.py/inference.py feed to the tokenizer.
inference_test_text: str = "[hi]त्यो त राम्रो भयो। अनि, मसँग कुरा गर्नुभएकोमा धन्यवाद।"
ljspeech = True # Set True if the dataset format is ljspeech, and False if it's file-based.
json_format = False # Set True if the dataset format is json, and False if it's file-based or ljspeech.
preprocess = True # If you've already done preprocessing once, set it to false.
is_turbo: bool = False # Set True if you're training Turbo, False if you're training Normal.
is_lora: bool = True # True: Efficient LoRA training (Recommended for < 10h data)
# False: Full Fine-Tune (High VRAM, for massive datasets)
lora_r: int = 128
lora_alpha: int = 256
turbo_lora_target_modules: List[str] = field(default_factory=lambda: ["c_attn", "c_proj", "c_fc", "spkr_enc"])
lora_target_modules: List[str] = field(default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "spkr_enc"])
lora_modules_to_save: List[str] = field(default_factory=lambda: ["text_emb", "text_head"])
# --- Vocabulary ---
# The size of the NEW vocabulary (from tokenizer.json)
# Ensure this matches the JSON file generated by your tokenizer script.
# For Turbo mode: Use the exact number provided by setup.py (e.g., 52260)
new_vocab_size: int = 52260 if is_turbo else 2454
# --- Hyperparameters ---
batch_size: int = 32 # Adjust based on VRAM (2, 4, 8)
grad_accum: int = 1 # Effective Batch Size = Batch * Accum
learning_rate: float = 1e-4 if is_lora else 1e-5 # T3 is sensitive, keep low
num_epochs: int = 10 if is_lora else 30
save_steps: int = 500
save_total_limit: int = 5
dataloader_num_workers: int = 8
# --- Constraints ---
start_text_token = 255
stop_text_token = 0
max_text_len: int = 256
max_speech_len: int = 850 # Truncates very long audio
prompt_duration: float = 3.0 # Duration for the reference prompt (seconds)
|