| 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 |
|
|
| |
| import utils |
|
|
|
|
| LOGGER = utils.get_logger(__name__) |
|
|
|
|
| class DatasetIterator: |
| HDR_MAGIC = b"LITPKDS" |
| HDR_SIZE = 24 |
| 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("<Q", f.read(8)) |
| assert version == (1,) |
| (dtype_code,) = struct.unpack("<B", f.read(1)) |
| dtype = self.DTYPES[dtype_code] |
| (chunk_size,) = struct.unpack("<Q", f.read(8)) |
| return dtype, chunk_size |
|
|
| def _close_mmaps(self): |
| for mmap in self._mmaps: |
| mmap._mmap.close() |
|
|
| def _load_n_chunks(self): |
| self._close_mmaps() |
| self._mmaps = [] |
| self._buffers = [] |
| if self._n_chunks > 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) |
| |
| dataset[:, 0] = vocab_size - 2 |
| dataset[:, -1] = vocab_size - 1 |
|
|
| for i in range(dataset_size): |
| |
| 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 = [] |
|
|
| |
| 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 |
| ]): |
| |
| 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 = { |
| '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() |
|
|
| |
| 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): |
| |
| concatenated_examples = list( |
| itertools.chain(* examples['input_ids'])) |
| total_length = len(concatenated_examples) |
| if insert_special_tokens: |
| |
| new_block_size = block_size - 2 |
| else: |
| new_block_size = block_size |
| total_length = ( |
| total_length // new_block_size) * new_block_size |
| |
| 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': |
| |
| |
| 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) |
| |
| 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: |
| |
| 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 |
| 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 |
| dataset = generate_synthetic_dataset( |
| name='swissroll', |
| train_dataset_size=100000, |
| validation_dataset_size=1024, |
| seq_len=2, |
| 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']: |
| |
| 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']]} |
| |
| 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) |
| |
| valid_loader.tokenizer = tokenizer |
|
|
| return train_loader, valid_loader |
|
|
|
|
| |
|
|
|
|
| class RandomFaultTolerantSampler(torch.utils.data.RandomSampler): |
|
|
| def __init__(self, *args, generator=None, **kwargs): |
| |
| |
| |
| |
| 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.restarting = True |
|
|
| |
| |
|
|
| 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 |
|
|
| |
| |
| def __iter__(self): |
| if self.shuffle: |
| |
| g = torch.Generator() |
| g.manual_seed(self.seed + self.epoch) |
| |
| indices = torch.randperm( |
| len(self.dataset), generator=g).tolist() |
| else: |
| |
| indices = list(range(len(self.dataset))) |
|
|
| if not self.drop_last: |
| |
| 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: |
| |
| indices = indices[:self.total_size] |
| assert len(indices) == self.total_size |
|
|
| |
| 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] |
|
|
| |
| prompt_ids = [tokenizer.encode(prompt) for prompt in prompts] |
| completion_ids = [tokenizer.encode(completion) for completion in completions] |
| |
| |
| if max_length is None: |
| max_length = max([p.shape[0] + c.shape[0] for p, c in zip(prompt_ids, completion_ids)]) |
| |
| |
| 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) |
| |
| |
| for i, (p_ids, c_ids) in enumerate(zip(prompt_ids, completion_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 |
| |
| |
| 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 |
| 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): |
| |
| |
| |
| 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 |
| )) |
|
|
| dataset.append(dict( |
| question=question + thought, |
| answer=answer |
| )) |
|
|
| 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') |
| |
| if mode == "train": |
| return dataset["train"] |
| elif mode == "validation": |
| return None |
| |
| |
| 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)) |
|
|
| |
| |
| |
| 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: |
| |
| 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() |
| |
| 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: |
| 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): |
| |
| string = string.replace("s '", "s'") |
| string = re.sub(r"/' [0-9]/", r"/'[0-9]/", string) |
| |
| string = string.replace(" @-@ ", "-") |
| string = string.replace(" @,@ ", ",") |
| string = string.replace(" @.@ ", ".") |
| |
| string = string.replace(" : ", ": ") |
| string = string.replace(" ; ", "; ") |
| string = string.replace(" . ", ". ") |
| string = string.replace(" ! ", "! ") |
| string = string.replace(" ? ", "? ") |
| string = string.replace(" , ", ", ") |
| |
| 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) |
| |
| 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("<unk>", "?") |
| 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): |
| |
| 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') |
|
|
| |
| with open(path, 'r') as f: |
| vocab = json.load(f) |
| assert isinstance(vocab, list) |
| if max_vocab is not None: |
| vocab = vocab[:max_vocab] |
|
|
| |
| 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'] |
|
|
| |
| assert len(ranks) == len(vocab) |
| assert set(ranks.values()) == set(range(len(ranks))) |
|
|
| return ranks |
|
|
|
|
| class CustomTikTokenizer(MegatronTokenizer): |
| SPECIAL_TOKENS = ['<unk>', '<s>', '</s>', '<mask>'] |
|
|
| 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 |
| 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'<SPECIAL_{i}>' 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) |
| |
| 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['<unk>'] |
| self._bos_id = special_tokens['<s>'] |
| self._eos_id = special_tokens['</s>'] |
| self._mask_id = special_tokens['<mask>'] |
| |
| |
| self.mask_token = '<mask>' |
| self.mask_token_id = self._mask_id |
| self.bos_token_id = self._bos_id |
| self.eos_token_id = self._eos_id |
|
|
| |
| self._model = tiktoken.Encoding( |
| name=pathlib.Path(path).parent.name, |
| pat_str=pattern, |
| mergeable_ranks=token_to_id_sans_special, |
| special_tokens=special_tokens) |
|
|
| |
| |
| |
| 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: |
| |
| |
| |
| |
|
|
| 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] |
