| import random |
| import copy |
| import json |
| import torch |
| import transformers |
| import numpy as np |
| import pickle as pkl |
| from torch.utils.data import Dataset |
| from dataclasses import dataclass, field |
| from typing import Dict, Optional, Sequence, List |
|
|
| IGNORE_INDEX = -100 |
| MAX_LENGTH = 2048 |
|
|
|
|
| @dataclass |
| class DataArguments: |
| data_path: str = field(default='./MusicCaps', metadata={"help": "Path to the training data."}) |
| feat_folder: Optional[str] = field(default='./MusicCaps/music_feat') |
|
|
| def preprocess_v1(sources: str, tokenizer: transformers.PreTrainedTokenizer, metadata, |
| prompt_pattern="USER: <Speech><SpeechHere></Speech> Describe the music in detail.\nASSISTANT:\n") -> Dict: |
| sources = sources.split('\n') |
| clips, duration, caption = metadata['clips'], metadata['duration'], [] |
| length = 0 |
| for l, c in zip(clips, sources): |
| caption.append( |
| f'From {int(length / duration * 100)} to {int((length + l) / duration * 100)},' |
| + ','.join(c.split(',')[1:]) |
| ) |
| length += l |
|
|
| targets = prompt_pattern + '\n'.join(caption) |
|
|
| targets_left, targets_right = targets.split('<SpeechHere>') |
| targets_right = tokenizer(targets_right, return_tensors="pt", add_special_tokens=False).input_ids[0] |
|
|
| sources_left, sources_right = prompt_pattern.split('<SpeechHere>') |
| sources_left = tokenizer(sources_left, return_tensors="pt", add_special_tokens=False).input_ids[0] |
| sources_right_length = tokenizer(sources_right, return_tensors="pt", add_special_tokens=False).input_ids.shape[-1] |
|
|
| sources_right = copy.deepcopy(targets_right) |
|
|
| targets_left = torch.LongTensor([IGNORE_INDEX] * len(sources_left)) |
| targets_right[:sources_right_length] = IGNORE_INDEX |
|
|
| sources_right, targets_right = sources_right[:MAX_LENGTH], targets_right[:MAX_LENGTH] |
|
|
| return dict(input_ids=(sources_left, sources_right), labels=(targets_left, targets_right)) |
|
|
|
|
| def preprocess(sources: str, tokenizer: transformers.PreTrainedTokenizer, metadata, |
| prompt_pattern="USER: <Speech><SpeechHere></Speech> Describe the music in detail.\nASSISTANT:\n") -> Dict: |
| targets = prompt_pattern + sources |
|
|
| targets_left, targets_right = targets.split('<SpeechHere>') |
| targets_right = tokenizer(targets_right, return_tensors="pt", add_special_tokens=False).input_ids[0] |
|
|
| sources_left, sources_right = prompt_pattern.split('<SpeechHere>') |
| sources_left = tokenizer(sources_left, return_tensors="pt", add_special_tokens=False).input_ids[0] |
| sources_right_length = tokenizer(sources_right, return_tensors="pt", add_special_tokens=False).input_ids.shape[-1] |
|
|
| sources_right = copy.deepcopy(targets_right) |
|
|
| targets_left = torch.LongTensor([IGNORE_INDEX] * len(sources_left)) |
| targets_right[:sources_right_length] = IGNORE_INDEX |
|
|
| sources_right, targets_right = sources_right[:MAX_LENGTH], targets_right[:MAX_LENGTH] |
|
|
| return dict(input_ids=(sources_left, sources_right), labels=(targets_left, targets_right)) |
|
|
|
|
| class LazySupervisedDataset(Dataset): |
| """Dataset for supervised fine-tuning.""" |
|
|
| def __init__(self, data_path, tokenizer, data_args): |
| super(LazySupervisedDataset, self).__init__() |
|
|
| self.tokenizer = tokenizer |
| self.list_data_dict = json.load(open(data_path, "r")) |
| self.data_args = data_args |
|
|
| def __len__(self): |
| return len(self.list_data_dict) |
|
|
| def __getitem__(self, i): |
| source = copy.deepcopy(self.list_data_dict[i]) |
|
|
| feature_path = '{}/{}.pkl'.format(self.data_args.feat_folder, source['id']) |
| music = pkl.load(open(feature_path, 'rb')) |
| speech = torch.from_numpy(music['speech']) |
| audio = torch.from_numpy(music['audio']) |
|
|
| captions = source['caption'] |
| if not isinstance(captions, str): |
| weights = np.asarray([len(c) for c in captions]) |
| weights = weights / weights.sum() |
| captions = random.choices(captions, weights, k=1)[0] |
|
|
| data_dict = preprocess(captions, self.tokenizer, source['meta']) |
| |
| data_dict['speeches'] = speech |
| data_dict['audios'] = audio |
| return data_dict |
|
|
|
|
| @dataclass |
| class DataCollatorForSupervisedDataset(object): |
| """Collate examples for supervised fine-tuning.""" |
|
|
| tokenizer: transformers.PreTrainedTokenizer |
|
|
| def __call__(self, instances): |
| input_ids, labels, speeches, audios = tuple( |
| [instance[key] for instance in instances] for key in ("input_ids", "labels", "speeches", "audios")) |
| batch = dict(input_ids=input_ids, labels=labels, speeches=speeches, audios=audios) |
| return batch |
|
|
|
|
| def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict: |
| """Make dataset and collator for supervised fine-tuning.""" |
| train_dataset = LazySupervisedDataset(tokenizer=tokenizer, data_path=data_args.data_path, data_args=data_args) |
| data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) |
| return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator) |
|
|
