# -*- coding: utf-8 -*- from __future__ import annotations from copy import deepcopy from dataclasses import dataclass from typing import Any, Dict, Iterable, List, Union import numpy as np import torch from datasets import Dataset, IterableDataset from flame.logging import get_logger from transformers import PreTrainedTokenizer logger = get_logger(__name__) class HuggingfaceDataset(IterableDataset): def __init__( self, dataset: Dataset, tokenizer: PreTrainedTokenizer, context_len: int = 2048, rank: int = 0, world_size: int = 1, buffer_size: int = 1024 ) -> HuggingfaceDataset: self.dataset = dataset self.tokenizer = tokenizer self.data = dataset.shard(world_size, rank) self.context_len = context_len self.rank = rank self.world_size = world_size self.buffer_size = buffer_size if tokenizer.vocab_size < torch.iinfo(torch.int16).max: self.dtype = torch.int16 elif tokenizer.vocab_size < torch.iinfo(torch.int32).max: self.dtype = torch.int32 else: self.dtype = torch.int64 self.states = None self.buffer = torch.tensor([], dtype=self.dtype) self.tokens = [] self.rand_id = 0 self.token_id = 0 self.rng_state = None self._epoch = 0 def __iter__(self): g = torch.Generator() g.manual_seed(self._epoch + self.rank) if self.rng_state is not None: g.set_state(self.rng_state) rand_it = self.randint(0, self.buffer_size, g=g) if self.states is not None: self.data.load_state_dict(self.states) # max number of tokens allowed in the chunk buffer n_tokens = self.buffer_size * self.context_len while True: for sample in self.tokenize(self.data): # keep appending the samples to the token buffer self.tokens += sample # if the token buffer is full, start sampling # NOTE: we first convert the token ids to a tensor of shape [n_chunks, context_len] for efficiency if len(self.buffer) == 0 and len(self.tokens) >= n_tokens: self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=self.dtype).view(self.buffer_size, -1) self.tokens = self.tokens[n_tokens:] if len(self.buffer) == self.buffer_size: yield from self.sample(rand_it) n_chunks = len(self.tokens) // self.context_len # handle the left tokens in the buffer if n_chunks > 0: n_tokens = n_chunks * self.context_len indices = torch.randperm(n_chunks, generator=g).tolist() self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=torch.long).view(n_chunks, -1) self.tokens = self.tokens[n_tokens:] for i in indices: yield {'input_ids': self.buffer[i]} def tokenize(self, data, batch_size: int = 64): texts, states = [], [] for sample in data: texts.append(sample['text']) states.append(self.data.state_dict()) if len(texts) == batch_size: for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']): self.states = s yield tokenized texts, states = [], [] if len(texts) > 0: for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']): self.states = s yield tokenized def sample(self, indices): n_tokens = (len(self.tokens) // self.context_len) * self.context_len while self.token_id < n_tokens: i = next(indices) start, end = self.token_id, self.token_id + self.context_len self.token_id += self.context_len yield {'input_ids': self.buffer[i].to(torch.long)} self.buffer[i] = torch.tensor(self.tokens[start:end], dtype=self.dtype) self.token_id = 0 self.tokens = self.tokens[n_tokens:] def randint( self, low: int, high: int, batch_size: int = 1024, g: torch.Generator = torch.Generator() ) -> Iterable[int]: indices = torch.empty(batch_size, dtype=torch.long) while True: # record the generator states before sampling self.rng_state = g.get_state() indices = torch.randint(low, high, (batch_size,), out=indices, generator=g) for i in indices[self.rand_id:].tolist(): self.rand_id += 1 yield i self.rand_id = 0 def set_epoch(self, epoch): self._epoch = epoch if hasattr(self.dataset, "set_epoch"): self.dataset.set_epoch(epoch) def state_dict(self): return { 'states': self.states, 'buffer': self.buffer.clone(), 'tokens': deepcopy(self.tokens), 'rand_id': self.rand_id, 'token_id': self.token_id, 'rng_state': self.rng_state, 'epoch': self._epoch } def load_state_dict(self, state_dict): self.states = state_dict['states'] self.buffer = state_dict['buffer'].clone() self.tokens = deepcopy(state_dict['tokens']) self.rand_id = state_dict['rand_id'] self.token_id = state_dict['token_id'] self.rng_state = state_dict['rng_state'].clone() if state_dict['rng_state'] is not None else None self._epoch = state_dict['epoch'] @dataclass class DataCollatorForLanguageModeling: """ Data collator used for language modeling. Args: tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]): The tokenizer used for encoding the data. varlen (`bool`): Whether to return sequences with variable lengths. If `True`, the offsets indicating the start and end of each sequence will be returned. For example, if the sequence lengths are `[4, 8, 12]`, the returned `input_ids` will be a long flattened tensor of shape `[1, 24]`, with `offsets` being `[0, 4, 12, 24]`. If `False`, the `input_ids` with shape `[batch_size, seq_len]` will be returned directly. return_tensors (`str`): The type of Tensor to return. Allowable values are "pt". """ tokenizer: PreTrainedTokenizer varlen: bool = False return_tensors: str = "pt" def __call__( self, examples: List[Union[List[int], Dict[str, Any]]] ) -> Dict[str, Any]: if not isinstance(examples[0], Dict): examples = [{'input_ids': example} for example in examples] def tensorize(example: Dict[str, Any]) -> Dict[str, Any]: tensorized = {} for key in ['input_ids', 'offsets']: if key not in example: continue if isinstance(example[key], List): tensorized[key] = torch.tensor(example[key], dtype=torch.long) elif isinstance(example[key], np.ndarray): tensorized[key] = torch.from_numpy(example[key]) else: tensorized[key] = example[key] return tensorized examples = list(map(tensorize, examples)) if not self.varlen: length_of_first = examples[0]['input_ids'].size(0) # Check if padding is necessary. if all(example['input_ids'].size(0) == length_of_first for example in examples): batch = { 'input_ids': torch.stack([example['input_ids'] for example in examples], dim=0), } else: # If yes, check if we have a `pad_token`. if self.tokenizer._pad_token is None: raise ValueError( f"You are attempting to pad samples but the tokenizer you are using " f"({self.tokenizer.__class__.__name__}) does not have a pad token." ) batch = self.tokenizer.pad(examples, return_tensors=self.return_tensors, return_attention_mask=False) else: if len(examples) > 1: raise ValueError("The batch size must be 1 for variable length inputs.") batch = { 'input_ids': torch.cat([example['input_ids'] for example in examples], dim=0).unsqueeze(0) } if 'offsets' in examples[0]: batch['offsets'] = torch.cat([example['offsets'] for example in examples], dim=0).unsqueeze(0) else: # determine boundaries by bos/eos positions if self.tokenizer.add_bos_token: offsets = [] if batch['input_ids'][0, 0] != self.tokenizer.bos_token_id: offsets.append(torch.tensor([0], dtype=torch.long)) offsets.append(torch.where(batch['input_ids'].eq(self.tokenizer.bos_token_id))[1]) offsets.append(torch.tensor([len(batch['input_ids'][0])], dtype=torch.long)) batch['offsets'] = torch.cat(offsets, dim=0) elif self.tokenizer.add_eos_token: offsets = [torch.tensor([0], dtype=torch.long)] offsets.append(torch.where(batch['input_ids'].eq(self.tokenizer.eos_token_id))[1] + 1) if batch['input_ids'][0, -1] != self.tokenizer.eos_token_id: offsets.append(torch.tensor([len(batch['input_ids'][0])], dtype=torch.long)) batch['offsets'] = torch.cat(offsets, dim=0) else: raise ValueError("You must allow the tokenizer to add either a bos or eos token as separators.") labels = batch['input_ids'].clone() if self.tokenizer.pad_token_id is not None: labels[labels == self.tokenizer.pad_token_id] = -100 batch["labels"] = labels return batch