import logging import os import random from hashlib import sha256 import pandas as pd from lightning import LightningDataModule from sklearn.model_selection import train_test_split from torch.utils.data import DataLoader from .datasets import TextAudioDataset from steerable_retrieval.utils.instantiators import instantiate def instantiate_datasets(datasets_cfg): datasets = [] dataset_names = [] for dataset_name in datasets_cfg: dataset = instantiate(datasets_cfg[dataset_name]) datasets.append(dataset) dataset_names.append(dataset_name) return datasets, dataset_names def get_song_describer_annotations(data_path = None, csv_path = None, val_split = 0.1): df = pd.read_csv(csv_path) df = df[['path','caption','is_valid_subset','caption_id']].rename(columns = {'path':'file_path'}) df['file_path'] = os.path.join(data_path) + '/' + df['file_path'] #replace .mp3 with .2min.mp3 df['file_path'] = df['file_path'].apply(lambda x: x.replace('.mp3','.2min.mp3')) records = df.to_dict(orient = 'records') for record in records: record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']} if val_split == 0.0: print('No validation split') for record in records: record['split'] = 'train' return records train_indices, val_indices = train_test_split(range(len(records)), test_size = val_split, random_state = 42) for idx in train_indices: records[idx]['split'] = 'train' for idx in val_indices: records[idx]['split'] = 'val' return records def get_musiccaps_annotations(data_path = None, csv_path = None, val_split = 0.1, test_split = 0.1): df = pd.read_csv(csv_path) df['file_path'] = data_path + '/' + df['ytid'] + '.wav' records = df.to_dict(orient = 'records') for record in records: record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']} if val_split == 0.0: print('No validation split') for record in records: record['split'] = 'train' return records # split into train, val, test train_indices, test_indices = train_test_split(range(len(records)), test_size = test_split + val_split, random_state = 42) val_indices, test_indices = train_test_split(test_indices, test_size = test_split/(test_split + val_split), random_state = 42) for idx in train_indices: records[idx]['split'] = 'train' for idx in val_indices: records[idx]['split'] = 'val' for idx in test_indices: records[idx]['split'] = 'test' return records def get_musiccaps_truncated_annotations(data_path = None, csv_path = None, val_split = 0.1, test_split = 0.1): df = pd.read_csv(csv_path) df['file_path'] = data_path + '/' + df['ytid'] + '.wav' import random records = df.to_dict(orient = 'records') print('Truncating captions') for record in records: # select a random sentence from the caption sentences = record['caption'].split('.') random_sentence = random.choice(sentences) record['caption'] = {sha256(random_sentence.encode('utf-8')).hexdigest(): random_sentence} if val_split == 0.0: print('No validation split') for record in records: record['split'] = 'train' return records # split into train, val, test train_indices, test_indices = train_test_split(range(len(records)), test_size = test_split + val_split, random_state = 42) val_indices, test_indices = train_test_split(test_indices, test_size = test_split/(test_split + val_split), random_state = 42) for idx in train_indices: records[idx]['split'] = 'train' for idx in val_indices: records[idx]['split'] = 'val' for idx in test_indices: records[idx]['split'] = 'test' return records def get_maxcaps_annotations(data_path = None, csv_path = None): """Read JSONL file line-by-line to avoid memory issues with large files.""" df = pd.read_csv(csv_path) df['file_path'] = data_path + '/' + df['file_path'] records = df.to_dict(orient = 'records') for record in records: record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']} return records def get_yt8m_annotations(data_path = None, csv_path = None): df = pd.read_csv(csv_path) df['file_path'] = data_path + '/' + df['file_path'] records = df.to_dict(orient = 'records') for record in records: record['caption'] = {sha256(record['caption'].encode('utf-8')).hexdigest(): record['caption']} return records def get_music4all_annotations(manifest_csv=None, audio_status_ok=True): """Annotations for music4all from the MuQ-MuLan embedding manifest. ``file_path`` is set directly to the pre-extracted ``.npy`` audio-embedding path so the preextracted-feature loader reads it as-is. Splits (train/val/test) and captions are taken from the manifest. Only rows with a successfully embedded audio are kept. """ df = pd.read_csv(manifest_csv) if audio_status_ok and 'audio_embedding_status' in df.columns: df = df[df['audio_embedding_status'] == 'ok'] df = df[df['audio_embedding_path'].notna()] records = [] for row in df.itertuples(index=False): caption = str(getattr(row, 'caption', '') or '') split = getattr(row, 'split', 'train') records.append({ 'file_path': row.audio_embedding_path, 'caption': {sha256(caption.encode('utf-8')).hexdigest(): caption}, 'split': split if split in ('train', 'val', 'test') else 'train', 'track_id': getattr(row, 'track_id', None), }) logging.info(f"Loaded {len(records)} music4all annotations from {manifest_csv}") return records def get_folder_annotations(data_path = None): # recursively get all files in the data_path directory that are audio files, and their paths audio_files = [] for root, dirs, files in os.walk(data_path): audio_files += [os.path.join(root, file) for file in files if file.endswith('.wav') or file.endswith('.mp3')] records = [{'file_path': file, 'caption': '', 'split': 'train'} for file in audio_files] logging.info(f"Found {len(records)} audio files in {data_path}") return records class TextAudioDataModule(LightningDataModule): def __init__(self, datasets, return_audio = True, return_text = True, concept = None, target_n_samples = 96000, target_sr = 48000, batch_size = 32, num_workers = 0, preextracted_features = False, truncate_preextracted = 50, root_dir = None, new_dir = None, **kwargs): super().__init__() self.annotations = [] dataset_names, datasets = list(datasets.keys()), list(datasets.values()) self.datasets = datasets self.dataset_names = dataset_names for dataset in self.datasets: dataset.split = dataset.split self.return_audio = return_audio self.return_text = return_text self.concept = concept self.target_n_samples = target_n_samples self.target_sr = target_sr self.batch_size = batch_size self.num_workers = num_workers self.preextracted_features = preextracted_features self.truncate_preextracted = truncate_preextracted self.truncate_preextracted = [self.truncate_preextracted for _ in range(len(self.datasets))] self.root_dirs = [dataset.root_dir for dataset in self.datasets] self.new_dirs = [dataset.new_dir for dataset in self.datasets] self.train_annotations = [] self.val_annotations = [] self.test_annotations = [] self.val_dataset_names = [] self.test_dataset_names = [] self.val_dataset_indices = [] # Track original dataset indices for val self.test_dataset_indices = [] # Track original dataset indices for test for i, dataset_ in enumerate(self.datasets): self.train_annotations.extend([annot for annot in dataset_.annotations if annot['split'] == 'train']) val_annots = [annot for annot in dataset_.annotations if annot['split'] == 'val'] test_annots = [annot for annot in dataset_.annotations if annot['split'] == 'test'] if val_annots: self.val_annotations.append(val_annots) self.val_dataset_names.append(self.dataset_names[i]) self.val_dataset_indices.append(i) if test_annots: self.test_annotations.append(test_annots) self.test_dataset_names.append(self.dataset_names[i]) self.test_dataset_indices.append(i) # filter for empty validation and test annotations (already filtered above, but keeping for safety) self.val_annotations = [annot for annot in self.val_annotations if annot] self.test_annotations = [annot for annot in self.test_annotations if annot] # Create dataloader_names mapping: maps dataloader_idx -> dataset_name (for val and test) # This will be used by callbacks to properly name metrics self.dataloader_names = {} for idx, name in enumerate(self.val_dataset_names): self.dataloader_names[idx] = name self.test_dataloader_names = {} for idx, name in enumerate(self.test_dataset_names): self.test_dataloader_names[idx] = name print(f"Number of training samples: {len(self.train_annotations)}") print(f"Number of validation samples: {sum([len(annot) for annot in self.val_annotations])} over {len(self.val_annotations)} datasets, {[len(annot) for annot in self.val_annotations]}") print(f"Number of test samples: {sum([len(annot) for annot in self.test_annotations])} over {len(self.test_annotations)} datasets, {[len(annot) for annot in self.test_annotations]}") print(f"Datasets: {self.dataset_names}") @property def names(self): """Mapping mode -> (dataloader_idx -> dataset name). Used by callbacks for log keys.""" return { "val": getattr(self, "dataloader_names", None) or {}, "test": getattr(self, "test_dataloader_names", None) or {}, } def setup(self, stage: str) -> None: if stage != 'eval': self.train_dataset = TextAudioDataset(annotations=self.train_annotations, target_n_samples=self.target_n_samples, target_sr=self.target_sr, return_audio=self.return_audio, return_text=self.return_text, concept=self.concept, preextracted_features=self.preextracted_features, truncate_preextracted=self.truncate_preextracted[0], root_dir=self.root_dirs[0], new_dir=self.new_dirs[0]) self.val_datasets = [TextAudioDataset(annotations=self.val_annotations[i], target_n_samples=self.target_n_samples, target_sr=self.target_sr, return_audio=self.return_audio, return_text=self.return_text, concept=self.concept, preextracted_features=self.preextracted_features, truncate_preextracted=self.truncate_preextracted[self.val_dataset_indices[i]], root_dir=self.root_dirs[self.val_dataset_indices[i]], new_dir=self.new_dirs[self.val_dataset_indices[i]]) for i in range(len(self.val_annotations))] self.test_datasets = [TextAudioDataset(annotations=self.test_annotations[i], target_n_samples=self.target_n_samples, target_sr=self.target_sr, return_audio=self.return_audio, return_text=self.return_text, concept=self.concept, preextracted_features=self.preextracted_features, truncate_preextracted=self.truncate_preextracted[self.test_dataset_indices[i]], root_dir=self.root_dirs[self.test_dataset_indices[i]], new_dir=self.new_dirs[self.test_dataset_indices[i]]) for i in range(len(self.test_annotations))] def train_dataloader(self): return DataLoader( self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, drop_last=True # Critical for distributed training to avoid hangs ) def val_dataloader(self): return [ DataLoader( val_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, drop_last=False # Don't drop for validation ) for val_dataset in self.val_datasets ] def test_dataloader(self): return [ DataLoader( test_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, drop_last=False # Don't drop for test ) for test_dataset in self.test_datasets ]