File size: 4,597 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import os
import random
import torch
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence
from src.utils import setup_logger
logger = setup_logger(__name__)
class ChatterboxDataset(Dataset):
def __init__(self, config):
self.cfg = config
self.preprocessed_dir = config.preprocessed_dir
if not os.path.exists(self.preprocessed_dir):
raise FileNotFoundError(f"Preprocessing folder not found: {self.preprocessed_dir}.")
self.files = [f for f in os.listdir(self.preprocessed_dir) if f.endswith(".pt")]
if len(self.files) == 0:
raise RuntimeError(f"There are no .pt files in the folder: {self.preprocessed_dir}")
logger.info(f"Dataset loaded. Total sample: {len(self.files)}")
self.sot_token = config.start_text_token
self.eot_token = config.stop_text_token
def __len__(self):
return len(self.files)
def __getitem__(self, idx):
try:
filename = self.files[idx]
pt_path = os.path.join(self.preprocessed_dir, filename)
data = torch.load(pt_path)
text_tokens = data["text_tokens"]
if text_tokens.size(0) > self.cfg.max_text_len - 2:
text_tokens = text_tokens[:self.cfg.max_text_len - 2]
sot = torch.tensor([self.sot_token], dtype=torch.long)
eot = torch.tensor([self.eot_token], dtype=torch.long)
text_tokens = torch.cat([sot, text_tokens, eot])
speech_tokens = data["speech_tokens"]
if speech_tokens.size(0) > self.cfg.max_speech_len:
speech_tokens = speech_tokens[:self.cfg.max_speech_len]
speaker_emb = data["speaker_emb"]
prompt_tokens = data["prompt_tokens"]
if random.random() < 0.20:
speaker_emb = torch.zeros_like(speaker_emb)
prompt_tokens = torch.zeros(1, dtype=torch.long)
return {
"text_tokens": text_tokens,
"speech_tokens": speech_tokens,
"speaker_emb": speaker_emb,
"prompt_tokens": prompt_tokens
}
except Exception as e:
logger.error(f"Error loading {filename}: {e}")
return None
def data_collator_standart(batch):
batch = [item for item in batch if item is not None]
if not batch:
return {}
# Padding
text_tokens = pad_sequence([x["text_tokens"] for x in batch], batch_first=True, padding_value=0)
speech_tokens = pad_sequence([x["speech_tokens"] for x in batch], batch_first=True, padding_value=0)
prompt_tokens = pad_sequence([x["prompt_tokens"] for x in batch], batch_first=True, padding_value=0)
speaker_embs = torch.stack([x["speaker_emb"] for x in batch])
# Lengths
text_lens = torch.tensor([len(x["text_tokens"]) for x in batch], dtype=torch.long)
speech_lens = torch.tensor([len(x["speech_tokens"]) for x in batch], dtype=torch.long)
return {
"text_tokens": text_tokens,
"text_token_lens": text_lens,
"speech_tokens": speech_tokens,
"speech_token_lens": speech_lens,
"speaker_emb": speaker_embs,
"prompt_tokens": prompt_tokens
}
def data_collator_turbo(batch):
batch = [item for item in batch if item is not None]
if not batch:
return {}
# 1. Text Tokens Padding
text_tokens = pad_sequence([x["text_tokens"] for x in batch], batch_first=True, padding_value=0)
text_lens = torch.tensor([len(x["text_tokens"]) for x in batch], dtype=torch.long)
# 2. Speech Tokens Padding
speech_tokens = pad_sequence([x["speech_tokens"] for x in batch], batch_first=True, padding_value=0)
speech_lens = torch.tensor([len(x["speech_tokens"]) for x in batch], dtype=torch.long)
# 3. Prompt Tokens Padding
prompt_tokens = pad_sequence([x["prompt_tokens"] for x in batch], batch_first=True, padding_value=0)
prompt_lens = torch.tensor([x["prompt_tokens"].shape[0] for x in batch], dtype=torch.long)
# 4. Speaker Embedding
speaker_embs = torch.stack([x["speaker_emb"] for x in batch])
return {
"text_tokens": text_tokens,
"text_token_lens": text_lens,
"speech_tokens": speech_tokens,
"speech_token_lens": speech_lens,
"speaker_emb": speaker_embs,
"prompt_tokens": prompt_tokens,
"prompt_lens": prompt_lens
} |