import functools import itertools import json import math import os import pathlib import re import shutil import typing import urllib import zipfile import random import struct import glob from tqdm import tqdm import datasets from sklearn.datasets import make_checkerboard, make_swiss_roll, make_circles import fsspec import numpy as np import requests import sentencepiece import tokenizers import torch import transformers from dataclasses import dataclass from functools import partial from abc import ABC, abstractmethod import base64 import collections import tiktoken # from tokenizer import get_tokenizer, TinyLLamaTokenizer, Text8Tokenizer, SyntheticTokenizer, MegatronTokenizer, wt_detokenizer, ptb_detokenizer, lm1b_detokenizer, lambada_detokenizer, scientific_papers_detokenizer import utils LOGGER = utils.get_logger(__name__) class DatasetIterator: HDR_MAGIC = b"LITPKDS" HDR_SIZE = 24 # bytes DTYPES = {1: np.uint8, 2: np.int8, 3: np.int16, 4: np.int32, 5: np.int64, 6: np.float32, 7: np.float64, 8: np.uint16} def __init__(self, filenames, n_chunks, block_size, seed, shuffle, wrap, train_start_file_idx=0): self._seed = seed self._shuffle = shuffle self._rng = np.random.default_rng( seed) if shuffle else None self._block_idxs = None self._wrap = wrap self._filenames = filenames self._file_idx = train_start_file_idx self._n_chunks = n_chunks self._dtype = None self._block_size = block_size self._n_blocks = None self._mmaps = [] self._buffers = [] self._block_idxs = [] self._curr_idx = 0 self._load_n_chunks() def _read_header(self, path): with open(path, "rb") as f: magic = f.read(len(self.HDR_MAGIC)) assert magic == self.HDR_MAGIC, ( "File doesn't match expected format.") version = struct.unpack(" len(self._filenames[self._file_idx:]): self._file_idx = 0 for i in range(self._n_chunks): filename = self._filenames[self._file_idx + i] if self._dtype is None: self._dtype, self._chunk_size = self._read_header( filename) self._n_blocks = self._chunk_size // self._block_size mmap = np.memmap(filename, mode='r', order='C', offset=self.HDR_SIZE) self._mmaps.append(mmap) self._buffers.append(memoryview(mmap)) self._file_idx += self._n_chunks n_all_blocks = self._n_chunks * self._n_blocks if self._shuffle: self._block_idxs = self._rng.permutation(n_all_blocks) else: self._block_idxs = range(n_all_blocks) self._curr_idx = 0 def __del__(self): self._close_mmaps() del self._mmaps del self._buffers def __iter__(self): return self def __next__(self): if self._curr_idx >= len(self._block_idxs): self._load_n_chunks() block_idx = self._block_idxs[self._curr_idx] chunk_id = block_idx // self._n_blocks buffer = self._buffers[chunk_id] elem_id = (block_idx % self._n_blocks) * \ self._block_size offset = np.dtype(self._dtype).itemsize * elem_id arr = np.frombuffer(buffer, dtype=self._dtype, count=self._block_size, offset=offset) self._curr_idx += 1 return { 'input_ids': torch.from_numpy(arr.astype(np.int64)), 'file_idx': self._file_idx, 'curr_idx': self._curr_idx} class CustomDataset(torch.utils.data.IterableDataset): def __init__(self, filenames, n_chunks, block_size, seed=0, shuffle=True, wrap=False, train_start_file_idx=0): self._filenames = filenames self._n_chunks = n_chunks self._block_size = block_size self._seed = seed self._shuffle = shuffle self._wrap = wrap self._train_start_file_idx = train_start_file_idx def __iter__(self): return DatasetIterator( filenames=self._filenames, n_chunks=self._n_chunks, block_size=self._block_size, seed=self._seed, shuffle=self._shuffle, wrap=self._wrap, train_start_file_idx=self._train_start_file_idx) def _generate_synthetic_data(name, dataset_size, seq_len, vocab_size): if name == 'random': dataset = np.zeros((dataset_size, seq_len), dtype=int) # tokens representing sequence boundary dataset[:, 0] = vocab_size - 2 # bos dataset[:, -1] = vocab_size - 1 # eos for i in range(dataset_size): # sample from 0, 1, ..., vocab_size - 3 temp = np.random.randint(vocab_size - 2) for j in reversed(range(1, seq_len - 1)): dataset[i, j] = temp if temp != 0: temp = temp // 4 else: temp = np.random.randint(vocab_size - 2) elif name == 'checkerboard': X, _, _ = make_checkerboard(shape=(2, 2), random_state=0) raise NotImplementedError('Checkerboard not implemented') elif name == 'swissroll': X, _ = make_swiss_roll(n_samples=dataset_size, noise=0.2, hole=False, random_state=0) eps = 1e-3 normalized_data = np.stack([ ( X[:, 0] - X[:, 0].min() ) / ( (X[:, 0].max() - X[:, 0].min() + eps) ), \ ( X[:, 2] - X[:, 2].min() ) / ( (X[:, 2].max() - X[:, 2].min() + eps) ) ], axis=1) data = np.int32(np.floor(normalized_data * vocab_size)) elif name == 'circles': X, _ = make_circles(noise=0.02, factor=0.5, random_state=0) raise NotImplementedError('Circles not implemented') else: raise ValueError(f'Invalid toy data name: {name}') return data def generate_synthetic_dataset(train_dataset_size, validation_dataset_size, name, seq_len, vocab_size): np.random.seed(42) train_data = torch.from_numpy( _generate_synthetic_data(name, train_dataset_size, seq_len, vocab_size)) train_dataset = datasets.Dataset.from_dict({ 'input_ids': train_data, 'attention_mask': torch.ones_like(train_data), }) train_dataset.set_format(type='torch') np.random.seed(41) validation_data = torch.from_numpy( _generate_synthetic_data(name, validation_dataset_size, seq_len, vocab_size)) validation_dataset = datasets.Dataset.from_dict({ 'input_ids': validation_data, 'attention_mask': torch.ones_like(validation_data), }) validation_dataset.set_format(type='torch') return { 'train': train_dataset, 'validation': validation_dataset, } def get_lambada_test_dataset(): url = "https://openaipublic.blob.core.windows.net/gpt-2/data/lambada_test.jsonl" def read_jsonl_to_list(url): response = requests.get(url, stream=True) data_list = [] # Process each line in the response content for line in response.iter_lines(decode_unicode=True): if line: data = json.loads(line) data_list.append(data) return data_list lambada_data = read_jsonl_to_list(url) dataset = datasets.Dataset.from_list(lambada_data) return dataset def get_text8_dataset(cache_dir, max_seq_length=256, drop_last=True, crop_train=False): """Adapted from: https://github.com/google-research/google-research/blob/master/d3pm/text/datasets.py#L344 Args: cache_dir: str, path to cache directory. max_seq_length: int, maximum length of sequences. (default: 256, as in D3PM codebase.) drop_last: bool, whether to drop the last incomplete batch. (default: True, as in D3PM codebase.) crop_train: bool, whether to subsample contiguous subsequences from training example. serves to make sure transformer models with absolute position embeddings do not have incorrect position-wise marginals. (default: False, but necessary to match D3PM AR) Returns: dataset: dataset.DatasetDict, with keys 'train', 'valid', 'test'. """ url = 'http://mattmahoney.net/dc/text8.zip' if not crop_train: cache_dir = f'{cache_dir}/text8' else: cache_dir = f'{cache_dir}/text8-crop-train' split_names = ['train', 'validation', 'test'] if not all([ utils.fsspec_exists(os.path.join(cache_dir, split)) for split in split_names ]): # Check if raw data exists raw_cache_dir = os.path.join(cache_dir, 'raw_data') if not all([ utils.fsspec_exists( os.path.join(raw_cache_dir, f'text8.{split}.txt')) for split in split_names ]): if not utils.fsspec_exists( os.path.join(raw_cache_dir, 'text8.zip')): utils.fsspec_mkdirs(raw_cache_dir, exist_ok=True) LOGGER.info( 'Downloading text8 from URL {}.'.format(url)) with (urllib.request.urlopen(url) as in_stream, open(os.path.join(raw_cache_dir, 'text8.zip'), 'wb') as out_file): shutil.copyfileobj(in_stream, out_file) with fsspec.open( os.path.join(raw_cache_dir, 'text8.zip'), 'rb') as f: rawdata = zipfile.ZipFile(f).read( 'text8').decode('utf-8') # Splits taken from D3PM codebase splits = { 'train': rawdata[:90000000], 'validation': rawdata[90000000: 95000000], 'test': rawdata[95000000:], } for split, data in splits.items(): _path = os.path.join(raw_cache_dir, f'text8.{split}.txt') with fsspec.open(_path, 'w') as f: f.write(data) else: splits = {} for split in split_names: _path = os.path.join(raw_cache_dir, f'text8.{split}.txt') with fsspec.open(_path, 'r') as f: splits[split] = f.read() # Chunk and save as datasets.DatasetDict def chunks(lst, n): """Yield successive n-sized chunks from lst.""" for i in range(0, len(lst), n): yield lst[i:i + n] dataset_dict = {} for k, v in splits.items(): if k == 'train' and crop_train == True: chunk_size = 2 * max_seq_length else: chunk_size = max_seq_length text = list(chunks(v, chunk_size)) if drop_last and len(text[-1]) < chunk_size: text = text[:-1] dataset_dict[k] = datasets.Dataset.from_dict( {'text': text}) dataset = datasets.DatasetDict(dataset_dict) dataset.save_to_disk(cache_dir) else: dataset = datasets.load_from_disk(cache_dir) return dataset def _group_texts(examples, block_size, bos, eos, insert_special_tokens): # Concatenate all texts. concatenated_examples = list( itertools.chain(* examples['input_ids'])) total_length = len(concatenated_examples) if insert_special_tokens: # [BOS] and [EOS] to be added new_block_size = block_size - 2 else: new_block_size = block_size total_length = ( total_length // new_block_size) * new_block_size # Split by chunks of max_len. result = {} _values = [] _attn_masks = [] for i in range(0, total_length, new_block_size): if insert_special_tokens: _values.append( [bos] + concatenated_examples[i: i + new_block_size] + [eos]) else: _values.append( concatenated_examples[i: i + new_block_size] ) _attn_masks.append(torch.ones(block_size)) result['input_ids'] = _values result['attention_mask'] = _attn_masks return result def create_dataloader( batch_size: int, block_size: int, filenames: list, n_chunks: int = 8, shuffle: bool = True, seed: int = 12345, pin_memory=True, num_workers=1, train_start_file_idx=0): random.seed(seed) random.shuffle(filenames) dataset = CustomDataset( filenames, n_chunks=n_chunks, block_size=block_size, shuffle=shuffle, seed=seed, train_start_file_idx=train_start_file_idx) return torch.utils.data.DataLoader( dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, persistent_workers=True) def create_dataloaders(config, seed, train_start_file_idx=0): data_dir = pathlib.Path(config.data.cache_dir) if config.data.train == 'slim_pajama': train_filenames = sorted(glob.glob(str(data_dir / 'train*'))) val_filenames = sorted(glob.glob(str(data_dir / 'validation*'))) elif config.data.train == 'nvidia': # data_dir/train/ and data_dir/validation/ contain # filenames with the pattern [0-255]_*.bin assert current_device < total_devices train_filenames = [] bucket_size = 256 // total_devices assert bucket_size * total_devices == 256 for i in range(bucket_size): idx = bucket_size * current_device + i train_filenames.extend( glob.glob(str(data_dir / f'train/{idx}*.bin'))) train_filenames = sorted(train_filenames) val_filenames = sorted(glob.glob(str(data_dir / 'validation/*'))) else: train_filenames = sorted( glob.glob(str(data_dir / f"train*"))) val_filenames = sorted( glob.glob(str(data_dir / f"validation*"))) train_dataloader = create_dataloader( batch_size=config.loader.batch_size, block_size=config.model.length, filenames=train_filenames, n_chunks=config.loader.n_chunks, shuffle=True, seed=seed, num_workers=1, pin_memory=config.loader.pin_memory, split='train', train_start_file_idx=train_start_file_idx) # TODO: shard validation dataloader val_dataloader = create_dataloader( batch_size=config.loader.eval_batch_size, block_size=config.model.length, filenames=val_filenames, n_chunks=config.loader.n_chunks, shuffle=False, seed=seed, num_workers=1, pin_memory=config.loader.pin_memory, split='validation') return train_dataloader, val_dataloader def get_dataset(dataset_name, tokenizer, wrap, mode, cache_dir, insert_eos=True, insert_special_tokens=True, block_size=1024, num_proc=len(os.sched_getaffinity(0)), streaming=False, revision: typing.Optional[str] = None): eos_tag = '' if not insert_eos: eos_tag += '_eosFalse' if not insert_special_tokens: eos_tag += '_specialFalse' if wrap: filename = f'{dataset_name}_{mode}_bs{block_size}_wrapped{eos_tag}.dat' else: filename = f'{dataset_name}_{mode}_bs{block_size}_unwrapped{eos_tag}.dat' _path = os.path.join(cache_dir, filename) if utils.fsspec_exists(_path): LOGGER.info(f'Loading data from: {_path}') return datasets.load_from_disk(_path).with_format('torch') LOGGER.info(f'Generating new data at: {_path}') LOGGER.info(f'{streaming=}') crop_train = dataset_name == 'text8-crop' if mode == 'train' and crop_train: # double block size for sub-sampling block_size *= 2 if dataset_name == 'wikitext103': dataset = datasets.load_dataset( 'wikitext', name='wikitext-103-raw-v1', cache_dir=cache_dir, revision=revision) elif dataset_name == 'wikitext2': dataset = datasets.load_dataset( 'wikitext', name='wikitext-2-raw-v1', cache_dir=cache_dir, revision=revision) elif dataset_name == 'ptb': dataset = datasets.load_dataset( 'ptb_text_only', cache_dir=cache_dir, revision=revision) elif dataset_name == 'lambada': dataset = get_lambada_test_dataset() elif dataset_name == 'text8': assert wrap assert revision is None dataset = get_text8_dataset( cache_dir, max_seq_length=block_size) elif dataset_name == 'text8-crop': assert revision is None dataset = get_text8_dataset( cache_dir, max_seq_length=block_size, crop_train=True) elif dataset_name == 'openwebtext-train': dataset = datasets.load_dataset( 'openwebtext', split='train[:-100000]', cache_dir=cache_dir, revision=revision, streaming=False, num_proc=num_proc, trust_remote_code=True) elif dataset_name == 'openwebtext-valid': dataset = datasets.load_dataset( 'openwebtext', split='train[-100000:]', cache_dir=cache_dir, revision=revision, streaming=False, num_proc=num_proc, trust_remote_code=True) elif dataset_name == 'scientific_papers_arxiv': dataset = datasets.load_dataset( 'scientific_papers', 'arxiv', trust_remote_code=True, cache_dir=cache_dir, streaming=streaming, revision=revision) elif dataset_name == 'scientific_papers_pubmed': dataset = datasets.load_dataset( 'scientific_papers', 'pubmed', trust_remote_code=True, cache_dir=cache_dir, streaming=streaming, revision=revision) elif dataset_name == 'ag_news': dataset = datasets.load_dataset( 'ag_news', cache_dir=cache_dir, streaming=streaming, revision=revision) elif dataset_name == 'random': assert streaming assert wrap # i.e., no pad tokens dataset = generate_synthetic_dataset( name='random', train_dataset_size=100000, validation_dataset_size=1024, seq_len=32, vocab_size=256, ) elif dataset_name == 'swissroll': assert streaming assert wrap # i.e., no pad tokens dataset = generate_synthetic_dataset( name='swissroll', train_dataset_size=100000, validation_dataset_size=1024, seq_len=2, # not used but it's 2D data vocab_size=100, ) else: dataset = datasets.load_dataset( dataset_name, cache_dir=cache_dir, streaming=streaming, trust_remote_code=True, revision=revision) if dataset_name in ['lambada', 'openwebtext-train', 'openwebtext-valid']: data = dataset else: data = dataset[mode] if dataset_name in ['random', 'swissroll']: # already tokenized, no further actions required return data if dataset_name.startswith('wikitext'): detokenizer = wt_detokenizer elif dataset_name == 'ptb': detokenizer = ptb_detokenizer elif dataset_name == 'lm1b': detokenizer = lm1b_detokenizer elif dataset_name == 'lambada': detokenizer = lambada_detokenizer elif dataset_name.startswith('scientific_papers'): detokenizer = scientific_papers_detokenizer else: detokenizer = None def _apply_detokenizer(detokenizer): def detok(text): for i, t in enumerate(text, 0): text[i] = detokenizer(t) return text return detok EOS = tokenizer.encode(tokenizer.eos_token)[0] BOS = tokenizer.encode(tokenizer.bos_token)[0] def preprocess_and_tokenize(example): if dataset_name == 'ptb': text = example['sentence'] elif 'scientific_papers' in dataset_name: text = example['article'] else: text = example['text'] if detokenizer is not None: text = _apply_detokenizer(detokenizer)(text) tokenizer.padding_side = 'right' tokenizer.truncation_side = 'right' if block_size is None: tokens = tokenizer(text, add_special_tokens=False, return_attention_mask=False, return_token_type_ids=False) return tokens if wrap: tokens = tokenizer(text, add_special_tokens=False, return_attention_mask=False, return_token_type_ids=False) if insert_eos: tokens = {'input_ids': [t + [EOS] for t in tokens['input_ids']]} # Still missing BOS, but will be added in group_texts else: tokens = tokenizer(text, max_length=block_size, padding='max_length', truncation=True, add_special_tokens=True, return_attention_mask=True, return_token_type_ids=True) return tokens if streaming: tokenized_dataset = data.map( preprocess_and_tokenize, batched=True) else: tokenized_dataset = data.map( preprocess_and_tokenize, batched=True, num_proc=num_proc, load_from_cache_file=True, desc='Tokenizing') if dataset_name == 'ptb': tokenized_dataset = tokenized_dataset.remove_columns( 'sentence') elif 'scientific_papers' in dataset_name: tokenized_dataset = tokenized_dataset.remove_columns([ 'article', 'abstract', 'section_names']) elif dataset_name == 'ag_news': tokenized_dataset = tokenized_dataset.remove_columns( ['text', 'label']) else: tokenized_dataset = tokenized_dataset.remove_columns( 'text') if not wrap: if not streaming: tokenized_dataset.save_to_disk(_path) return tokenized_dataset.with_format('torch') group_texts = functools.partial( _group_texts, block_size=block_size, bos=BOS, eos=EOS, insert_special_tokens=insert_special_tokens) if streaming: chunked_dataset = tokenized_dataset.map( group_texts, batched=True) else: chunked_dataset = tokenized_dataset.map( group_texts, batched=True, num_proc=num_proc, load_from_cache_file=True, desc='Grouping') chunked_dataset.save_to_disk(_path) chunked_dataset = chunked_dataset.with_format('torch') return chunked_dataset def get_dataloaders(config, tokenizer, skip_train=False, skip_valid=False, valid_seed=None): num_gpus = torch.cuda.device_count() assert (config.loader.global_batch_size == (config.loader.batch_size * config.trainer.num_nodes * num_gpus * config.trainer.accumulate_grad_batches)) if config.loader.global_batch_size % ( num_gpus * config.trainer.accumulate_grad_batches) != 0: raise ValueError( f'Train Batch Size {config.training.batch_size}' f'not divisible by {num_gpus} gpus with accumulation ' f'{config.trainer.accumulate_grad_batches}.') if config.loader.eval_global_batch_size % num_gpus != 0: raise ValueError( f'Eval Batch Size for {config.eval.batch_size} ' f'not divisible by {num_gpus}.') if skip_train: train_set = None else: train_set = get_dataset( config.data.train, tokenizer, mode='train', wrap=config.data.wrap, insert_eos=config.data.insert_train_eos, insert_special_tokens=getattr(config.data, 'insert_train_special', True), cache_dir=config.data.cache_dir, block_size=config.model.length, streaming=config.data.streaming, num_proc=config.loader.num_workers, revision=config.data.get("train_revision", None)) if config.data.valid in ['text8', 'lm1b', 'ag_news']: validation_split = 'test' else: validation_split = 'validation' if skip_valid: valid_set = None else: valid_set = get_dataset( config.data.valid, tokenizer, wrap=config.data.wrap, mode=validation_split, cache_dir=config.data.cache_dir, insert_eos=config.data.insert_valid_eos, insert_special_tokens=getattr(config.data, 'insert_valid_special', True), block_size=config.model.length, streaming=config.data.streaming, num_proc=config.loader.num_workers, revision=config.data.get("valid_revision", None)) if skip_train: train_loader = None else: train_loader = torch.utils.data.DataLoader( train_set, batch_size=config.loader.batch_size, num_workers=config.loader.num_workers, pin_memory=config.loader.pin_memory, shuffle=not config.data.streaming, persistent_workers=True) train_loader.tokenizer = tokenizer if skip_valid: valid_loader = None else: if valid_seed is None: shuffle_valid = False generator = None else: shuffle_valid = True generator = torch.Generator().manual_seed(valid_seed) valid_loader = torch.utils.data.DataLoader( valid_set, batch_size=config.loader.eval_batch_size, num_workers=config.loader.num_workers, pin_memory=config.loader.pin_memory, shuffle=shuffle_valid, generator=generator) # Will be used in generative perplexity calculation valid_loader.tokenizer = tokenizer return train_loader, valid_loader # Samplers adapted from: https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/fault_tolerant_sampler.py class RandomFaultTolerantSampler(torch.utils.data.RandomSampler): def __init__(self, *args, generator=None, **kwargs): # TD [2022-07-17]: We don't force the seed to be zero. We generate random seed, # which should be reproducible if pl.seed_everything was called beforehand. # This means that changing the seed of the experiment will also change the # sampling order. if generator is None: seed = int(torch.empty( (), dtype=torch.int64).random_().item()) generator = torch.Generator().manual_seed(seed) kwargs.pop('shuffle', None) super().__init__(*args, generator=generator, **kwargs) self.counter = 0 self.restarting = False def state_dict(self): return {'random_state': self.generator.get_state(), 'counter': self.counter} def load_state_dict(self, state_dict): self.generator.set_state(state_dict.get('random_state')) self.counter = state_dict['counter'] # self.start_counter = self.counter self.restarting = True # TD [2022-08-28] Setting the len will cause PL to think there are only a few batches left per # epoch, and subsequent epoch will have very few batches. def __iter__(self) -> typing.Iterator[int]: n = len(self.data_source) self.state = self.generator.get_state() indices = torch.randperm( n, generator=self.generator).tolist() if not self.restarting: self.counter = 0 else: indices = indices[self.counter:] self.restarting = False for index in indices: self.counter += 1 yield index self.counter = 0 class FaultTolerantDistributedSampler(torch.utils.data.DistributedSampler): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.counter = 0 self.restarting = False def state_dict(self): return {'epoch': self.epoch, 'counter': self.counter} def load_state_dict(self, state_dict): self.epoch = state_dict['epoch'] self.counter = state_dict['counter'] self.restarting = True # TD [2022-08-28] Setting the len will cause PL to think there are only a few batches left per # epoch, and subsequent epoch will have very few batches. def __iter__(self): if self.shuffle: # deterministically shuffle based on epoch and seed g = torch.Generator() g.manual_seed(self.seed + self.epoch) # type: ignore[arg-type] indices = torch.randperm( len(self.dataset), generator=g).tolist() else: # type: ignore[arg-type] indices = list(range(len(self.dataset))) if not self.drop_last: # add extra samples to make it evenly divisible padding_size = self.total_size - len(indices) if padding_size <= len(indices): indices += indices[:padding_size] else: indices += (indices * math.ceil( padding_size / len(indices)))[:padding_size] else: # remove tail of data to make it evenly divisible. indices = indices[:self.total_size] assert len(indices) == self.total_size # subsample indices = indices[self.rank:self.total_size:self.num_replicas] assert len(indices) == self.num_samples if not self.restarting: self.counter = 0 else: indices = indices[self.counter:] self.restarting = False for index in indices: self.counter += 1 yield index self.counter = 0 def my_collate_fn(batch, tokenizer, keys, max_length=None): """Collate function to process and pad the batch""" prompts = [item[keys[0]] for item in batch] completions = [item[keys[1]] for item in batch] # Tokenize all prompts and completions prompt_ids = [tokenizer.encode(prompt) for prompt in prompts] completion_ids = [tokenizer.encode(completion) for completion in completions] # Calculate max length for padding if max_length is None: max_length = max([p.shape[0] + c.shape[0] for p, c in zip(prompt_ids, completion_ids)]) # Initialize padded tensors batch_size = len(batch) input_ids_padded = torch.full((batch_size, max_length), tokenizer.pad_token_id) prompt_ids_padded = torch.full((batch_size, max_length), tokenizer.pad_token_id) # Fill the padded tensors for i, (p_ids, c_ids) in enumerate(zip(prompt_ids, completion_ids)): # Combine prompt and completion for input_ids input_ids_padded[i, :p_ids.shape[0]] = p_ids input_ids_padded[i, p_ids.shape[0]:p_ids.shape[0] + c_ids.shape[0]] = c_ids # Only prompt for prompt_ids prompt_ids_padded[i, :p_ids.shape[0]] = p_ids return { 'input_ids': input_ids_padded, 'prompt_ids': prompt_ids_padded } class CustomDataset(torch.utils.data.Dataset): def __init__(self, data): self.data = data def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] class CustomGSM8KDataset(torch.utils.data.Dataset): def __init__(self, mode, max_length, tokenizer): if mode == 'train': self.crop = 7_472 - 500 elif mode == 'validation': mode = 'train' self.crop = 500 # same setting as for gsm8k-aug self.dataset_raw = datasets.load_dataset('openai/gsm8k', 'main')[mode] self.max_length = max_length self.tokenizer = tokenizer self.data = self.preprocess_gsm8k() def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def preprocess_gsm8k(self): dataset = [] n_excluded = 0 for i in tqdm(range(self.crop), desc="Preprocessing GSM8K"): question, full_answer = self.dataset_raw[i]['question'], self.dataset_raw[i]['answer'] thought, answer = full_answer.split('####') question = 'Question: ' + question + '\nAnswer: ' thought = thought answer = '####' + answer question_tokens = self.tokenizer.encode(question) thought_tokens = self.tokenizer.encode(thought) answer_tokens = self.tokenizer.encode(answer) length_tokens = len(question_tokens) + len(thought_tokens) + len(answer_tokens) if length_tokens > self.max_length: n_excluded += 1 continue dataset.append(dict( question=question, answer=thought + answer )) print(f"Excluded {n_excluded} examples due to length > {self.max_length}") return CustomDataset(dataset) class CustomGSM8KAugmentedDataset(torch.utils.data.Dataset): def __init__(self, mode, max_length, tokenizer): self.dataset_raw = datasets.load_dataset('whyNLP/gsm8k-aug-nl')[mode] self.max_length = max_length self.tokenizer = tokenizer self.data = self.preprocess_gsm8k_aug() def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def preprocess_gsm8k_aug(self): dataset = [] n_excluded = 0 for i in tqdm(range(len(self.dataset_raw)), desc="Preprocessing GSM8K-Aug"): question, thought, answer = self.dataset_raw[i]['question'], self.dataset_raw[i]['steps'], self.dataset_raw[i]['answer'] question = 'Question: ' + question + '\nAnswer: ' thought = ' '.join(thought) answer = '####' + answer question_tokens = self.tokenizer.encode(question) thought_tokens = self.tokenizer.encode(thought) answer_tokens = self.tokenizer.encode(answer) length_tokens = len(question_tokens) + len(thought_tokens) + len(answer_tokens) if length_tokens > self.max_length: n_excluded += 1 continue dataset.append(dict( question=question, answer=thought + answer )) print(f"Excluded {n_excluded} examples due to length > {self.max_length}") return CustomDataset(dataset) class CustomGSM8KAugmentedSMDMDataset(torch.utils.data.Dataset): def __init__(self, mode, max_length, tokenizer): self.dataset_raw = datasets.load_dataset('whyNLP/gsm8k-aug-nl')[mode] self.max_length = max_length self.tokenizer = tokenizer self.data = self.preprocess_gsm8k_aug() def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def preprocess_gsm8k_aug(self): ### Create an augmented version where the data is like # Question , Answer = Thought # Question , Answer = Thought, Answer dataset = [] n_excluded = 0 for i in tqdm(range(len(self.dataset_raw)), desc="Preprocessing GSM8K-Aug-SMDM"): question, thought, answer = self.dataset_raw[i]['question'], self.dataset_raw[i]['steps'], self.dataset_raw[i]['answer'] question = 'Question: ' + question + '\nAnswer: ' thought = ''.join(thought) answer = '####' + answer question_tokens = self.tokenizer.encode(question) thought_tokens = self.tokenizer.encode(thought) answer_tokens = self.tokenizer.encode(answer) length_tokens = len(question_tokens) + len(thought_tokens) + len(answer_tokens) if length_tokens > self.max_length: n_excluded += 1 continue dataset.append(dict( question=question, answer=thought )) # That is just so weird dataset.append(dict( question=question + thought, answer=answer )) # That is just so weird print(f"Excluded {n_excluded} examples due to length > {self.max_length}") return CustomDataset(dataset) class SFTDataLoader(torch.utils.data.DataLoader): def __init__(self, *args, tokenizer=None, keys=None, max_length=None, **kwargs): super().__init__( *args, **kwargs ) self.tokenizer = tokenizer self.keys = keys self.max_length = max_length self.collate_fn = partial(my_collate_fn, tokenizer=self.tokenizer, keys=self.keys, max_length=self.max_length) def get_sft_dataset(config, dataset_name, cache_dir, tokenizer, mode='train'): global_batch_size = config.sft.loader.global_batch_size if dataset_name == "reversal_curse": dataset = datasets.load_dataset( 'json', data_files=f'{cache_dir}/all_prompts_train.jsonl') # 'json', data_files=f'{cache_dir}/d2p_prompts_train.jsonl') if mode == "train": return dataset["train"] elif mode == "validation": return None # elif dataset_name == "gsm8k-aug": # return datasets.load_dataset("whyNLP/gsm8k-aug")[mode] elif dataset_name == "gsm8k": return CustomGSM8KDataset(mode=mode, max_length=config.sft.data.target_length, tokenizer=tokenizer) elif dataset_name == "gsm8k-aug": return CustomGSM8KAugmentedDataset(mode=mode, max_length=config.sft.data.target_length, tokenizer=tokenizer) elif dataset_name == "gsm8k-aug-smdm": return CustomGSM8KAugmentedSMDMDataset(mode=mode, max_length=config.sft.data.target_length, tokenizer=tokenizer) elif dataset_name == "gsm8k-cat-aug": dataset_train = CustomGSM8KDataset(mode=mode, max_length=config.sft.data.target_length, tokenizer=tokenizer) dataset_aug = CustomGSM8KAugmentedDataset(mode=mode, max_length=config.sft.data.target_length, tokenizer=tokenizer) return torch.utils.data.ConcatDataset([dataset_train, dataset_aug]) else: raise NotImplementedError(f"Dataset {dataset_name} not implemented") def get_sft_keys(config, dataset_name): if dataset_name == "reversal_curse": return ["prompt", "completion"] elif 'gsm8k' in dataset_name: return ["question", "answer"] else: raise NotImplementedError(f"Dataset {dataset_name} not implemented") def get_sft_dataloaders(config, tokenizer, skip_train=False, skip_valid=False): num_gpus = torch.cuda.device_count() assert (config.sft.loader.global_batch_size == (config.sft.loader.batch_size * config.sft.trainer.num_nodes * num_gpus * config.sft.trainer.accumulate_grad_batches)) if config.sft.loader.global_batch_size % ( num_gpus * config.sft.trainer.accumulate_grad_batches) != 0: raise ValueError( f'Train Batch Size {config.sft.training.batch_size}' f'not divisible by {num_gpus} gpus with accumulation ' f'{config.sft.trainer.accumulate_grad_batches}.') if config.sft.loader.eval_global_batch_size % num_gpus != 0: raise ValueError( f'Eval Batch Size for {config.sft.eval.batch_size} ' f'not divisible by {num_gpus}.') train_set = get_sft_dataset( config=config, dataset_name=config.sft.data.train, cache_dir=config.sft.data.cache_dir, tokenizer=tokenizer, mode='train', ) train_keys = get_sft_keys(config, config.sft.data.train) train_dataloader = SFTDataLoader( dataset=train_set, batch_size=config.sft.loader.batch_size, shuffle=config.sft.data.shuffle, num_workers=config.sft.loader.num_workers, pin_memory=config.sft.loader.pin_memory, tokenizer=tokenizer, keys=train_keys, max_length=config.sft.data.target_length) valid_set = get_sft_dataset( config=config, dataset_name=config.sft.data.train, cache_dir=config.sft.data.cache_dir, tokenizer=tokenizer, mode='validation', ) if skip_valid or valid_set is None: valid_dataloader = None else: valid_keys = get_sft_keys(config, config.sft.data.valid) valid_dataloader = SFTDataLoader( dataset=valid_set, batch_size=config.sft.loader.batch_size, shuffle=False, num_workers=config.sft.loader.num_workers, pin_memory=config.sft.loader.pin_memory, tokenizer=tokenizer, keys=valid_keys, max_length=config.sft.data.target_length) return train_dataloader, valid_dataloader PATTERN_TIKTOKEN_V2 = "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" def get_tokenizer(config): if config.data.tokenizer_name_or_path == 'tinyllama': return TinyLLamaTokenizer( pathlib.Path(config.data.tokenizer_path)) elif 'tiktokenizer' in config.data.tokenizer_name_or_path: try: return CustomTikTokenizer( path=config.data.tokenizer_name_or_path, pattern=PATTERN_TIKTOKEN_V2, vocab_size=None, num_special_tokens=1000, special_tokens=None) except FileNotFoundError: try: return CustomTikTokenizer( path=config.data.tokenizer_name_or_path.replace('jmlemercier', 'jlemercier'), pattern=PATTERN_TIKTOKEN_V2, vocab_size=None, num_special_tokens=1000, special_tokens=None) except FileNotFoundError: return CustomTikTokenizer( path=config.data.tokenizer_name_or_path.replace('jlemercier', 'jmlemercier'), pattern=PATTERN_TIKTOKEN_V2, vocab_size=None, num_special_tokens=1000, special_tokens=None) elif config.data.tokenizer_name_or_path == 'text8': tokenizer = Text8Tokenizer() elif config.data.tokenizer_name_or_path == 'bert-base-uncased': tokenizer = transformers.BertTokenizer.\ from_pretrained('bert-base-uncased') elif config.data.tokenizer_name_or_path == 'synthetic': tokenizer = SyntheticTokenizer(vocab_size=256) elif config.data.tokenizer_name_or_path == '2d-toy': tokenizer = SyntheticTokenizer(vocab_size=100) else: tokenizer = transformers.AutoTokenizer.from_pretrained( config.data.tokenizer_name_or_path) if (isinstance(tokenizer, transformers.GPT2TokenizerFast) or isinstance(tokenizer, transformers.GPT2Tokenizer)): tokenizer._tokenizer.post_processor = tokenizers.processors.BertProcessing( (tokenizer.bos_token, tokenizer.bos_token_id), (tokenizer.eos_token, tokenizer.eos_token_id)) # For wrapped batches: # [BOS] sent1 [EOS] sent2-fragment [EOS] # [BOS] sent2-fragment [EOS] sent3 [EOS] if tokenizer.bos_token is None: if tokenizer.cls_token is None: raise AttributeError( 'Tokenizer must have a bos_token or ' f'cls_token: {tokenizer}') tokenizer.bos_token = tokenizer.cls_token if tokenizer.eos_token is None: if tokenizer.sep_token is None: raise AttributeError( 'Tokenizer must have a eos_token ' f'or sep_token: {tokenizer}') tokenizer.eos_token = tokenizer.sep_token if tokenizer.pad_token is None: tokenizer.add_special_tokens({'pad_token': '[PAD]'}) return tokenizer class TinyLLamaTokenizer: def __init__(self, checkpoint_dir: pathlib.Path) -> None: # some checkpoints have both files, `.model` takes precedence if (vocabulary_path := checkpoint_dir / 'tokenizer.model').is_file(): self.processor = sentencepiece.SentencePieceProcessor( model_file=str(vocabulary_path)) self.bos_token_id = self.processor.bos_id() self.eos_token_id = self.processor.eos_id() self.pad_token_id = self.processor.eos_id() self._vocab_size = self.processor.vocab_size() # add a mask token to the vocabulary self.mask_token_id = self._vocab_size self.mask_token = '[MASK]' self._vocab_size += 1 else: raise NotImplementedError @property def vocab_size(self) -> int: return self._vocab_size def __len__(self) -> int: return self._vocab_size def token_to_id(self, token: str) -> int: id_ = self.processor.piece_to_id(token) if id_ is None: raise ValueError( f"token {token!r} not found in the collection.") return id_ def encode(self, string: str, device: typing.Optional[torch.device] = None, bos: bool = False, eos: bool = True, max_length: int = None) -> torch.Tensor: tokens = self.processor.encode(string) if bos: bos_id = self.bos_token_id if bos_id is None: raise NotImplementedError( 'This tokenizer does not defined a bos token') tokens = [bos_id] + tokens if eos: tokens = tokens + [self.eos_token_id] if max_length is not None: tokens = tokens[:max_length] tokens = tokens + [self.pad_token_id] * (max_length - len(tokens)) return torch.tensor(tokens, dtype=torch.int, device=device) @dataclass class BatchEncodeOutput: tokens: torch.Tensor attn_mask: typing.Optional[torch.Tensor] def batch_encode(self, strings: typing.List[str], device: typing.Optional[torch.device] = None, bos: bool = False, eos: bool = True, max_length: int = None) -> BatchEncodeOutput: tokens = [self.processor.encode(string) for string in strings] if bos: bos_id = self.bos_token_id if bos_id is None: raise NotImplementedError( 'This tokenizer does not defined a bos token') tokens = [[bos_id] + t for t in tokens] if eos: tokens = [t + [self.eos_token_id] for t in tokens] if max_length is None: #Pad to the longest sequence in the batch max_length = max(len(t) for t in tokens) attn_mask = torch.ones(len(tokens), max_length, dtype=torch.bool, device=device) for n in range(len(tokens)): tokens[n] = tokens[n][:max_length] tokens[n] = tokens[n] + [self.pad_token_id] * (max_length - tokens[n].shape[-1]) attn_mask[n, - (max_length - tokens[n].shape[-1]): ] = False return self.BatchEncodeOutput( tokens=torch.tensor(tokens, dtype=torch.int, device=device), attn_mask=attn_mask) def batch_encode(self, *args, **kwargs): return self.encode(*args, **kwargs) def decode(self, tensor: torch.Tensor, *ignored_args, **ignored_kwargs) -> str: if tensor.ndim == 0: tokens = [tensor.item()] else: tokens = tensor.tolist() return self.processor.decode(tokens) def batch_decode(self, *args, **kwargs): return self.decode(*args, **kwargs) def wt_detokenizer(string): # contractions string = string.replace("s '", "s'") string = re.sub(r"/' [0-9]/", r"/'[0-9]/", string) # number separators string = string.replace(" @-@ ", "-") string = string.replace(" @,@ ", ",") string = string.replace(" @.@ ", ".") # punctuation string = string.replace(" : ", ": ") string = string.replace(" ; ", "; ") string = string.replace(" . ", ". ") string = string.replace(" ! ", "! ") string = string.replace(" ? ", "? ") string = string.replace(" , ", ", ") # double brackets string = re.sub(r"\(\s*([^\)]*?)\s*\)", r"(\1)", string) string = re.sub(r"\[\s*([^\]]*?)\s*\]", r"[\1]", string) string = re.sub(r"{\s*([^}]*?)\s*}", r"{\1}", string) string = re.sub(r"\"\s*([^\"]*?)\s*\"", r'"\1"', string) string = re.sub(r"'\s*([^']*?)\s*'", r"'\1'", string) # miscellaneous string = string.replace("= = = =", "====") string = string.replace("= = =", "===") string = string.replace("= =", "==") string = string.replace(" " + chr(176) + " ", chr(176)) string = string.replace(" \n", "\n") string = string.replace("\n ", "\n") string = string.replace(" N ", " 1 ") string = string.replace(" 's", "'s") return string def ptb_detokenizer(x): x = x.replace(" 's", "'s") x = x.replace("s ' ", "s' ") x = x.replace(" n't", "n't") x = x.replace(" \n ", "\n") x = x.replace("\\/", "/") for _ in range(10): x = x.replace(" N ", " 1 ") x = x.replace("$ 1", "$1") x = x.replace("# 1", "#1") x = x.replace("", "?") return x def lm1b_detokenizer(x): x = x.replace('http : / / ', 'http://') x = x.replace('https : / / ', 'https://') x = re.sub(r' \'(\w+)', r"'\1", x) x = re.sub(r' (\w+) \. ', r' \1. ', x) x = re.sub(r' (\w+) \.$', r' \1.', x) x = x.replace(' ? ', '? ') x = re.sub(r' \?$', '?', x) x = x.replace(' ! ', '! ') x = re.sub(r' \!$', '!', x) x = x.replace(' , ', ', ') x = x.replace(' : ', ': ') x = x.replace(' ; ', '; ') x = x.replace(' / ', '/') x = re.sub(r'\" ([^\"]+) \"', r'"\1"', x) x = re.sub(r'\' ([^\']+) \'', r"'\1'", x) x = re.sub(r'\( ([^\(\)]+) \)', r"(\1)", x) x = re.sub(r'\[ ([^\[\]]+) \]', r"[\1]", x) x = x.replace('$ ', '$') x = x.replace('£ ', '£') return x def lambada_detokenizer(text): text = text.replace("“", '"') text = text.replace("”", '"') return '\n' + text.strip() def scientific_papers_detokenizer(x): x = wt_detokenizer(x) x = lm1b_detokenizer(x) return x class SyntheticTokenizer( transformers.PreTrainedTokenizer): def __init__( self, vocab_size, bos_token="[BOS]", eos_token="[EOS]", sep_token=None, cls_token=None, pad_token=None, mask_token=None, unk_token=None, **kwargs): self.tokens = [] for i in range(vocab_size): # appending space for readability self.tokens.append(str(i) + " ") self._vocab_str_to_int = { '[BOS]': vocab_size - 2, '[EOS]': vocab_size - 1, ** {ch: i for i, ch in enumerate(self.tokens)}} self._vocab_int_to_str = { v: k for k, v in self._vocab_str_to_int.items()} super().__init__( bos_token=bos_token, eos_token=eos_token, sep_token=sep_token, cls_token=cls_token, pad_token=pad_token, mask_token=mask_token, unk_token=unk_token, **kwargs) @property def vocab_size(self) -> int: return len(self._vocab_str_to_int) def _tokenize(self, text: str, **kwargs) -> typing.List[str]: return list(text.lower()) def _convert_token_to_id(self, token: str) -> int: return self._vocab_str_to_int.get( token, self._vocab_str_to_int['[UNK]']) def _convert_id_to_token(self, index: int) -> str: return self._vocab_int_to_str[index] def convert_tokens_to_string(self, tokens): return ''.join(tokens) def get_vocab(self) -> typing.Dict[str, int]: return self._vocab_str_to_int class Text8Tokenizer(transformers.PreTrainedTokenizer): def __init__( self, bos_token='[BOS]', eos_token='[EOS]', sep_token='[SEP]', cls_token='[CLS]', pad_token='[PAD]', mask_token='[MASK]', unk_token='[UNK]', **kwargs): self.characters = list('abcdefghijklmnopqrstuvwxyz ') self._vocab_str_to_int = { '[CLS]': 0, '[SEP]': 1, '[BOS]': 2, '[EOS]': 3, '[MASK]': 4, '[PAD]': 5, '[RESERVED]': 6, '[UNK]': 7, ** {ch: i + 8 for i, ch in enumerate(self.characters)}} self._vocab_int_to_str = { v: k for k, v in self._vocab_str_to_int.items()} super().__init__( bos_token=bos_token, eos_token=eos_token, sep_token=sep_token, cls_token=cls_token, pad_token=pad_token, mask_token=mask_token, unk_token=unk_token, **kwargs) @property def vocab_size(self) -> int: return len(self._vocab_str_to_int) def _tokenize(self, text: str, **kwargs) -> typing.List[str]: return list(text.lower()) def _convert_token_to_id(self, token: str) -> int: return self._vocab_str_to_int.get( token, self._vocab_str_to_int['[UNK]']) def _convert_id_to_token(self, index: int) -> str: return self._vocab_int_to_str[index] def convert_tokens_to_string(self, tokens): return ''.join(tokens) def get_vocab(self) -> typing.Dict[str, int]: return self._vocab_str_to_int class MegatronTokenizer(ABC): """Abstract class for tokenizer Absent a config or class-specific tracking of which objects are uniquely identifying, we must include all key word arguments as unique identifiers Args: tokenizer_paths (Tuple[str]): All tokenizer source paths or prefixes tokenizer_options (Dict[str, Any]): All tokenizer options """ def __init__(self, *tokenizer_paths, **tokenizer_options): self.unique_identifiers = collections.OrderedDict() self.unique_identifiers["class"] = type(self).__name__ self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths) for option in tokenizer_options: self.unique_identifiers[option] = str(tokenizer_options[option]) self.unique_description = json.dumps(self.unique_identifiers, indent=4) super().__init__() @abstractmethod def tokenize(self, text): """Convert text to embedding ids Args: text (str): The text to convert Returns: numpy.ndarray: The converted embedding ids """ pass def detokenize(self, ids): """Convert embedding ids to text Args: ids (numpy.ndarray): The ids to convert Returns: str: The converted text Raises: NotImplementedError: Non-abstract, optional method """ raise NotImplementedError( f'{type(self).__name__} has no method "detokenize"') def offsets(self, ids, text): """Convert embedding ids to text offsets Args: ids (list[int]): The ids to convert text (str): The text to convert Returns: list[int]: The converted offsets Raises: NotImplementedError: Non-abstract, optional method """ raise NotImplementedError( f'{type(self).__name__} has no method "offsets"') @property @abstractmethod def vocab(self): """Dictionary from vocab text token to id token""" pass @property @abstractmethod def inv_vocab(self): """Dictionary from vocab id token to text token""" pass @property @abstractmethod def vocab_size(self): """The vocabulary size""" pass @property def cls(self): """The CLS token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "cls"') @property def sep(self): """The SEP token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "sep"') @property def pad(self): """The PAD token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "pad"') @property def eod(self): """The EOD token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "eod"') @property def bos(self): """The BOS token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "bos"') @property def eos(self): """The EOS token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "eos"') @property def mask(self): """The MASK token id Raises: NotImplementedError: Non-abstract, optional attribute """ raise NotImplementedError( f'{type(self).__name__} has no attribute "mask"') def reload_mergeable_ranks(path, max_vocab=None): """ Reloads a tokenizer JSON file and converts it to Tiktoken format. """ assert path.endswith('.json') # reload vocab with open(path, 'r') as f: vocab = json.load(f) assert isinstance(vocab, list) if max_vocab is not None: vocab = vocab[:max_vocab] # build ranks ranks: typing.Dict[bytes, int] = {} for i, x in enumerate(vocab): assert x.keys() == {'rank', 'token_bytes', 'token_str'} assert x['rank'] == i merge = base64.b64decode(x['token_bytes']) assert i >= 256 or merge == bytes([i]) ranks[merge] = x['rank'] # sanity check assert len(ranks) == len(vocab) assert set(ranks.values()) == set(range(len(ranks))) return ranks class CustomTikTokenizer(MegatronTokenizer): SPECIAL_TOKENS = ['', '', '', ''] def __init__(self, path, pattern, vocab_size, num_special_tokens, special_tokens): super().__init__( path, pattern=pattern, vocab_size=vocab_size, num_special_tokens=num_special_tokens, special_tokens=special_tokens) if vocab_size is None: vocab_size = 2**17 # Fallback vocab size is 131072. self._vocab_size = vocab_size if special_tokens is None: special_tokens = self.SPECIAL_TOKENS.copy() assert len(special_tokens) == len(set(special_tokens)), ( f'Special tokens should be unique: {special_tokens}') assert (len(special_tokens) <= num_special_tokens < self._vocab_size) assert set(self.SPECIAL_TOKENS) <= set(special_tokens), ( f'Custom special tokens should include {self.SPECIAL_TOKENS}') special_filler = [ f'' for i in range( len(special_tokens), num_special_tokens)] special_tokens = special_tokens + special_filler assert ( len(set(special_tokens)) == len(special_tokens) == num_special_tokens), ( f'Special tokens should be unique: {special_tokens}') inner_vocab_size = self._vocab_size - num_special_tokens token_to_id_sans_special = reload_mergeable_ranks( path, max_vocab=inner_vocab_size) # Create space for special tokens. token_to_id_sans_special = { t: i + num_special_tokens for t, i in token_to_id_sans_special.items()} special_tokens = { t: i for i, t in enumerate(special_tokens)} self._unk_id = special_tokens[''] self._bos_id = special_tokens[''] self._eos_id = special_tokens[''] self._mask_id = special_tokens[''] # Public attributes for backward compatibility. self.mask_token = '' self.mask_token_id = self._mask_id self.bos_token_id = self._bos_id self.eos_token_id = self._eos_id # Create tiktoken model. self._model = tiktoken.Encoding( name=pathlib.Path(path).parent.name, pat_str=pattern, mergeable_ranks=token_to_id_sans_special, special_tokens=special_tokens) # Create final _id_to_token and _token_to_id data # structures with special tokens inserted # at appropriate locations. assert set( token_to_id_sans_special.keys()).isdisjoint( set(special_tokens.keys())) self._token_to_id = token_to_id_sans_special.copy() self._token_to_id.update(special_tokens) self._id_to_token = { v: k for k, v in self._token_to_id.items()} assert (set(range(self._vocab_size)) == set(self._id_to_token.keys())) @property def bos(self) -> int: return self._bos_id @property def eos(self) -> int: return self._eos_id @property def unk(self) -> int: return self._unk_id @property def mask(self) -> int: return self._mask_id @property def eod(self) -> int: return self._eos_id @property def vocab(self): return self._token_to_id @property def inv_vocab(self): return self._id_to_token def tokenize(self, s, bos=False, eos=False): tokens = self._model.encode_ordinary(s) if bos: tokens = [self.bos, *tokens] if eos: tokens = [*tokens, self.eos] return tokens def detokenize(self, tokens): return self._model.decode(tokens) def offsets(self, ids, text): try: return self._model.decode_with_offsets(ids)[1] except UnicodeDecodeError: # Tiktoken has an unnecessary check that raises UnicodeDecodeError # from `text = b"".join(token_bytes).decode("utf-8", errors="strict")` # which is not needed for our use case. So we re-implement it, without # the check. token_bytes = self._model.decode_tokens_bytes(ids) text_len = 0 offsets = [] for token in token_bytes: offsets.append(max(0, text_len - (0x80 <= token[0] < 0xC0))) text_len += sum(1 for c in token if not 0x80 <= c < 0xC0) return offsets @property def vocab_size(self): return self._vocab_size @property def encoder(self): return self._token_to_id @property def decoder(self): return self._id_to_token def __len__(self): return self._vocab_size def decode(self, tokens, **kwargs): if isinstance(tokens, torch.Tensor): tokens = tokens.cpu().numpy().reshape(-1) elif isinstance(tokens, np.ndarray): tokens = tokens.reshape(-1) elif isinstance(tokens, list): tokens = tokens else: raise ValueError(f"Invalid type for tokens: {type(tokens)}") return self.detokenize(tokens) def encode(self, text): return torch.Tensor(self._model.encode_ordinary(text)).type(torch.long) def batch_decode(self, tokens): tokens = tokens.cpu().numpy() return [self.detokenize(token) for token in tokens]