|
from typing import Optional, Iterator, Callable, Any |
|
|
|
import torch |
|
from datasets import load_dataset, concatenate_datasets |
|
from transformers import AutoTokenizer |
|
|
|
|
|
def load_text_dataset(tokenizer: AutoTokenizer, |
|
kind: str, |
|
path: str, |
|
name: Optional[str]=None, |
|
data_dir: Optional[str]=None, |
|
data_files: Optional[str]=None, |
|
keep_in_memory: bool=False, |
|
revision: Optional[str]=None, |
|
split: str='train', |
|
num_proc: Optional[int]=None, |
|
format: Optional[Callable|str]=None) -> Any: |
|
assert isinstance(format, str) or callable(format), f'{path=} {format=}' |
|
assert kind == 'base' |
|
|
|
dataset = load_dataset(path=path, |
|
name=name, |
|
data_dir=data_dir, |
|
data_files=data_files, |
|
keep_in_memory=keep_in_memory, |
|
revision=revision, |
|
split=split, |
|
trust_remote_code=True, |
|
num_proc=num_proc) |
|
|
|
EOS_TOKEN = tokenizer.eos_token |
|
|
|
def format_dataset(batch): |
|
nonlocal EOS_TOKEN |
|
nonlocal format |
|
texts: list = [] |
|
rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())] |
|
|
|
if callable(format): |
|
for row in rows: |
|
|
|
text = format(row) |
|
|
|
if not text: |
|
text = '[NONE]' |
|
|
|
text += EOS_TOKEN |
|
texts.append(text) |
|
else: |
|
for row in rows: |
|
|
|
text = format.format(**row) |
|
|
|
if not text: |
|
text = '[NONE]' |
|
|
|
text += EOS_TOKEN |
|
texts.append(text) |
|
|
|
return {'text': texts} |
|
|
|
dataset = dataset.map(format_dataset, batched=True) |
|
return dataset |
|
|
|
|
|
def load_chat_dataset(tokenizer: AutoTokenizer, |
|
kind: str, |
|
path: str, |
|
name: Optional[str]=None, |
|
data_dir: Optional[str]=None, |
|
data_files: Optional[str]=None, |
|
keep_in_memory: bool=False, |
|
revision: Optional[str]=None, |
|
split: str='train', |
|
num_proc: Optional[int]=None, |
|
field: Optional[str]=None, |
|
transform: Optional[Callable]=None) -> Any: |
|
assert kind == 'instruct' |
|
|
|
dataset = load_dataset(path=path, |
|
name=name, |
|
data_dir=data_dir, |
|
data_files=data_files, |
|
keep_in_memory=keep_in_memory, |
|
revision=revision, |
|
split=split, |
|
trust_remote_code=True, |
|
num_proc=num_proc) |
|
|
|
EOS_TOKEN = tokenizer.eos_token |
|
|
|
def format_dataset(batch): |
|
nonlocal EOS_TOKEN |
|
nonlocal tokenizer |
|
nonlocal field |
|
nonlocal transform |
|
texts: list = [] |
|
rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())] |
|
|
|
if callable(transform): |
|
for row in rows: |
|
if field: |
|
messages = transform(row[field]) |
|
else: |
|
messages = transform(row) |
|
|
|
text = tokenizer.apply_chat_template(messages, tokenize=False) |
|
text += EOS_TOKEN |
|
texts.append(text) |
|
else: |
|
for row in rows: |
|
if field: |
|
messages = row[field] |
|
else: |
|
raise ValueError(field) |
|
|
|
text = tokenizer.apply_chat_template(messages, tokenize=False) |
|
text += EOS_TOKEN |
|
texts.append(text) |
|
|
|
return {'text': texts} |
|
|
|
dataset = dataset.map(format_dataset, batched=True) |
|
return dataset |
|
|
