build up

DG_dataset.py

@@ -0,0 +1,298 @@
+import torch
+from datasets import load_dataset, Dataset
+from itertools import chain
+from typing import List, Optional
+from transformers import AutoTokenizer
+from transformers.testing_utils import CaptureLogger
+from transformers.utils.logging import get_logger
+
+
+class DGDataset:
+    def __init__(
+        self, 
+        dataset: str = "blended_skill_talk",
+        task: str = "seq2seq",
+        tokenizer: AutoTokenizer = None,
+        max_source_length: int = 512,
+        max_target_length: int = 512,
+        padding: str = "max_length",
+        ignore_pad_token_for_loss: bool = True,
+        preprocessing_num_workers: int = None,
+        overwrite_cache: bool = True,
+    ):
+        self.dataset = dataset
+        self.task = task
+        self.tokenizer = tokenizer
+        self.max_source_length = max_source_length
+        self.max_target_length = max_target_length
+        self.padding = padding
+        self.ignore_pad_token_for_loss = ignore_pad_token_for_loss
+        self.preprocessing_num_workers = preprocessing_num_workers
+        self.overwrite_cache = overwrite_cache
+        # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
+        self.tok_logger = get_logger("transformers.tokenization_utils_base")
+
+
+    def prepare_context(self, instance: dict):
+        if self.dataset == 'blended_skill_talk':
+            num_entries = len(instance["free_messages"])
+            total_entries = num_entries
+            if self.task == 'seq2seq':
+                persona_pieces = f"<PS>{instance['personas'][1]}"
+                if instance['context'] == "wizard_of_wikipedia":
+                    additional_context_pieces = f"<CTX>{instance['additional_context']}."
+                else:
+                    additional_context_pieces = ""
+                context = persona_pieces + additional_context_pieces
+            else:
+                num_entries = min(num_entries, 2)
+                context = ''
+            prev_utt_pc = [sent for sent in instance["previous_utterance"] if sent != '']
+
+        elif self.dataset == 'conv_ai_2':
+            total_entries = len(instance['dialog'])
+            num_entries = total_entries//2
+            if self.task == 'seq2seq':
+                user_profile = ' '.join([''.join(x) for x in instance['user_profile']])
+                persona_pieces = f"<PS>{user_profile}"
+                context = persona_pieces
+            else:
+                num_entries = min(num_entries, 2)
+                context = ''
+            prev_utt_pc = []
+
+        elif self.dataset == 'empathetic_dialogues':
+            total_entries = len(instance['dialog'])
+            num_entries = total_entries//2
+            if self.task == 'seq2seq':
+                persona_pieces = f"<PS>{instance['prompt']}"
+                additional_context_pieces = f"<CTX>{instance['context']}."
+                context = persona_pieces + additional_context_pieces
+            else:
+                num_entries = min(num_entries, 2)
+                context = ''
+            prev_utt_pc = []
+
+        elif self.dataset == 'AlekseyKorshuk/persona-chat':
+            total_entries = len(instance['utterances'])
+            num_entries = total_entries//2
+            if self.task == 'seq2seq':
+                user_profile = ' '.join(instance['personality'])
+                persona_pieces = f"<PS>{user_profile}"
+                context = persona_pieces
+            else:
+                num_entries = min(num_entries, 2)
+                context = ''
+            prev_utt_pc = []
+
+        else:
+            raise ValueError("Dataset not supported.")
+        return num_entries, total_entries, context, prev_utt_pc
+
+
+    def prepare_entry(
+        self, 
+        instance: dict, 
+        entry_idx: int, 
+        context: str, 
+        prev_utt_pc: List[str], 
+        total_entries: int,
+    ):
+        if self.dataset == 'blended_skill_talk':
+            free_message = instance['free_messages'][entry_idx]
+            guided_message = instance['guided_messages'][entry_idx]
+            references = [values[entry_idx] for key, values in instance['suggestions'].items()]
+
+        elif self.dataset == 'conv_ai_2':
+            free_message = instance['dialog'][entry_idx*2]['text']
+            if entry_idx*2+1 >= total_entries:
+                guided_message = None
+            else:
+                guided_message = instance['dialog'][entry_idx*2+1]['text']
+            references = []
+
+        elif self.dataset == 'empathetic_dialogues':
+            free_message = instance['dialog'][entry_idx*2]['text']
+            if entry_idx*2+1 >= total_entries:
+                guided_message = None
+            else:
+                guided_message = instance['dialog'][entry_idx*2+1]['text']
+            references = []
+
+        elif self.dataset == 'AlekseyKorshuk/persona-chat':
+            free_message = instance['utterances'][entry_idx*2]['history'][-1]
+            if entry_idx*2+1 >= total_entries:
+                guided_message = None
+            else:
+                guided_message = instance['utterances'][entry_idx*2+1]['history'][-1]
+            references = instance['utterances'][entry_idx*2]['candidates']
+
+        else:
+            raise ValueError("Dataset not supported.")
+
+        if not prev_utt_pc:
+            original_context = context
+        else:
+            sp_token = '<SEP>' if self.task == 'seq2seq' else ' '
+            original_context = context + sp_token + sp_token.join(prev_utt_pc)
+
+        references.append(guided_message)
+        return free_message, guided_message, original_context, references
+
+
+    def tokenize_and_align_labels(self, instance: dict):
+        num_entries, total_entries, context, prev_utt_pc = self.prepare_context(instance)
+        inputs, labels = [], []
+        for entry_idx in range(num_entries):
+            free_message, guided_message, original_context, references = self.prepare_entry(
+                instance, 
+                entry_idx, 
+                context, 
+                prev_utt_pc,
+                total_entries,
+            )
+            if guided_message is None:
+                continue
+            # Input & Output
+            if self.task == 'seq2seq':
+                text = original_context + self.tokenizer.eos_token + free_message
+            else:
+                text = original_context + free_message + guided_message
+
+            inputs.append(text)
+            labels.append(guided_message)
+            prev_utt_pc += [
+                free_message,
+                guided_message,
+            ]
+
+        if not inputs:
+            return {"input_ids": [], "labels": [], "attention_mask": []}
+
+        if self.task == 'seq2seq':
+            inputs = self.tokenizer(inputs, max_length=self.max_source_length, padding=self.padding, truncation=True)
+            # Setup the tokenizer for targets
+            with self.tokenizer.as_target_tokenizer():
+                labels = self.tokenizer(labels, max_length=self.max_target_length, padding=self.padding, truncation=True)
+
+            # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 
+            # when we want to ignore padding in the loss.
+            if self.padding == "max_length" and self.ignore_pad_token_for_loss:
+                labels["input_ids"] = [
+                    [(l if l != self.tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
+                ]
+            inputs["labels"] = labels["input_ids"]
+            return inputs
+        else:
+            with CaptureLogger(self.tok_logger) as cl:
+                inputs = self.tokenizer(
+                    inputs, 
+                    return_tensors="pt",
+                    max_length=self.max_source_length, 
+                    padding=self.padding, 
+                    truncation=True,
+                )
+                labels = self.tokenizer(
+                    labels, 
+                    return_tensors="pt",
+                    max_length=self.max_target_length, 
+                    padding=self.padding, 
+                    truncation=True,
+                )
+
+            new_inputs = inputs.copy()
+            for k, v1 in inputs.items():
+                v2 = labels[k]
+                new_inputs[k] = torch.cat((v1, v2), dim=1)
+
+            new_labels = torch.cat((-100*torch.ones_like(inputs["input_ids"]), labels["input_ids"]), dim=1)
+            new_inputs["labels"] = new_labels
+
+            # clm input could be much much longer than block_size
+            if "Token indices sequence length is longer than the" in cl.out:
+                self.tok_logger.warning(
+                    "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits"
+                    " before being passed to the model."
+                )
+            return new_inputs
+
+
+    def group_texts(self, examples):
+        # ['input_ids', 'attention_mask', 'labels']
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        return concatenated_examples
+
+
+    def group_ED(self, dataset: Dataset):
+        results = {
+            'conv_id': [], 
+            'prompt': [],
+            'dialog': [], 
+            'context': [],
+        }
+        for i, instance in enumerate(dataset):
+            if instance['utterance_idx'] == 1:
+                results['conv_id'].append(instance['conv_id'])
+                results['dialog'].append([])
+                results['prompt'].append(instance['prompt'])
+                results['context'].append(instance['context'])
+
+            response = {'text': instance['utterance'], 'speaker_idx': instance['speaker_idx']}
+            results['dialog'][-1].append(response)
+        return Dataset.from_dict(results)
+
+
+    def preprocess(self, dataset: Dataset):
+        if self.dataset == "empathetic_dialogues":
+            dataset = self.group_ED(dataset)
+
+        dataset = dataset.map(
+            self.tokenize_and_align_labels,
+            batched=False,
+            num_proc=self.preprocessing_num_workers,
+            remove_columns=dataset.column_names,
+            load_from_cache_file=not self.overwrite_cache,
+        )
+        dataset = dataset.map(
+            self.group_texts,
+            batched=True,
+            num_proc=self.preprocessing_num_workers,
+            load_from_cache_file=not self.overwrite_cache,
+        )
+        return dataset
+
+
+
+
+if __name__ == "__main__":
+    from transformers import AutoTokenizer
+    from datasets import load_dataset
+
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+    tokenizer.pad_token = tokenizer.eos_token
+    data_names = [
+        "conv_ai_2",
+        "empathetic_dialogues",
+        "AlekseyKorshuk/persona-chat",
+        "blended_skill_talk",
+    ]
+    task = "seq2seq"
+    max_length = 256
+
+    for data_name in data_names:
+        train_dataset = load_dataset(data_name)["train"]
+        dg = DGDataset(
+            dataset=data_name,
+            task=task,
+            tokenizer=tokenizer,
+            max_source_length=max_length,
+            max_target_length=max_length,
+        )
+        print('{}: {}'.format(data_name, train_dataset))
+        train_dataset = dg.preprocess(train_dataset)
+        print("processed dataset: ", train_dataset)
+        print("processed dataset[0]: ", train_dataset[0])
+
+
+
+