|
"""Simple Dataset Reader.""" |
|
|
|
import random |
|
from typing import Dict, List, Optional, Union |
|
|
|
import torch |
|
from datasets import Dataset, DatasetDict |
|
from transformers import AutoTokenizer |
|
|
|
from opencompass.openicl.icl_prompt_template import PromptTemplate |
|
from opencompass.registry import ICL_DATASET_READERS |
|
from opencompass.utils.types import (_check_dataset, _check_str, |
|
_check_type_list) |
|
|
|
|
|
@ICL_DATASET_READERS.register_module() |
|
class DatasetReader: |
|
"""In-conext Learning Dataset Reader Class Generate an DatasetReader |
|
instance through 'dataset'. |
|
|
|
Attributes: |
|
dataset (:obj:`Dataset` or :obj:`DatasetDict`): The dataset to be read. |
|
input_columns (:obj:`List[str]` or :obj:`str`): A list of column names |
|
(a string of column name) in the dataset that represent(s) the |
|
input field. |
|
output_column (:obj:`str`): A column name in the dataset that |
|
represents the prediction field. |
|
input_template (:obj:`PromptTemplate`, optional): An instance of the |
|
:obj:`PromptTemplate` class, used to format the input field |
|
content during the retrieval process. (in some retrieval methods) |
|
output_template (:obj:`PromptTemplate`, optional): An instance of the |
|
:obj:`PromptTemplate` class, used to format the output field |
|
content during the retrieval process. (in some learnable retrieval |
|
methods) |
|
train_split (str): The name of the training split. Defaults to 'train'. |
|
train_range (int or float or str, optional): The size of the partial |
|
training dataset to load. |
|
If None, the entire training dataset will be loaded. |
|
If int or float, the random partial dataset will be loaded with the |
|
specified size. |
|
If str, the partial dataset will be loaded with the |
|
specified index list (e.g. "[:100]" for the first 100 examples, |
|
"[100:200]" for the second 100 examples, etc.). Defaults to None. |
|
test_split (str): The name of the test split. Defaults to 'test'. |
|
test_range (int or float or str, optional): The size of the partial |
|
test dataset to load. |
|
If None, the entire test dataset will be loaded. |
|
If int or float, the random partial dataset will be loaded with the |
|
specified size. |
|
If str, the partial dataset will be loaded with the |
|
specified index list (e.g. "[:100]" for the first 100 examples, |
|
"[100:200]" for the second 100 examples, etc.). Defaults to None. |
|
""" |
|
dataset = None |
|
input_template = None |
|
output_template = None |
|
|
|
def __init__(self, |
|
dataset: Union[Dataset, DatasetDict, str], |
|
input_columns: Union[List[str], str], |
|
output_column: Optional[str], |
|
input_template: Optional[PromptTemplate] = None, |
|
output_template: Optional[PromptTemplate] = None, |
|
train_split: str = 'train', |
|
train_range: Optional[Union[int, float, str]] = None, |
|
test_split: str = 'test', |
|
test_range: Optional[Union[int, float, str]] = None) -> None: |
|
self.input_columns = _check_type_list(input_columns, [List, str]) |
|
if isinstance(self.input_columns, str): |
|
self.input_columns = self.input_columns.split() |
|
self.output_column = None |
|
if output_column: |
|
self.output_column = _check_str(output_column) |
|
|
|
train_range = _check_type_list(train_range, [None, int, float, str]) |
|
test_range = _check_type_list(test_range, [None, int, float, str]) |
|
|
|
if input_template is not None: |
|
self.input_template = PromptTemplate._check_prompt_template( |
|
input_template) |
|
if output_template is not None: |
|
self.output_template = PromptTemplate._check_prompt_template( |
|
output_template) |
|
|
|
self.dataset = _check_dataset(dataset) |
|
if isinstance(self.dataset, Dataset): |
|
self.dataset = DatasetDict({ |
|
'train': self.dataset, |
|
'test': self.dataset |
|
}) |
|
|
|
|
|
for origin_split, mapped_split, split_range in [[ |
|
train_split, 'train', train_range |
|
], [test_split, 'test', test_range]]: |
|
self.dataset[mapped_split] = load_partial_dataset( |
|
self.dataset[origin_split], size=split_range) |
|
|
|
def generate_input_field_prompt(self, entry: Dict) -> str: |
|
"""Generate a prompt for the input field based on the provided |
|
:obj:`entry` data. |
|
|
|
Args: |
|
entry (:obj:`Dict`): A piece of data to be used for generating the |
|
prompt. |
|
|
|
Returns: |
|
:obj:`str`: The generated prompt. |
|
""" |
|
prompt = None |
|
if self.input_template is None: |
|
prompt = ' '.join([str(entry[ctx]) for ctx in self.input_columns]) |
|
else: |
|
prompt = self.input_template.generate_item(entry) |
|
return prompt |
|
|
|
def generate_input_field_corpus(self, |
|
dataset: Union[Dataset, DatasetDict], |
|
split: Optional[str] = None) -> List[str]: |
|
"""Generate corpus for input field. |
|
|
|
Args: |
|
dataset (:obj:`Dataset` or :obj:`DatasetDict`): A |
|
:obj:`datasets.Dataset` or :obj:`datasets.DatasetDict` |
|
instance. |
|
split (:obj:`str`, optional): The split of the dataset to use. If |
|
:obj:`None`, the entire dataset will be used. Defaults to |
|
``None``. |
|
|
|
Returns: |
|
:obj:`List[str]`: A list of generated input field prompts. |
|
""" |
|
if split is not None: |
|
dataset = dataset[split] |
|
corpus = [] |
|
for entry in dataset: |
|
corpus.append(self.generate_input_field_prompt(entry)) |
|
return corpus |
|
|
|
def generate_output_field_prompt(self, entry: Dict) -> str: |
|
"""Generate a prompt for the output field based on the provided |
|
:obj:`entry` data. |
|
|
|
Args: |
|
entry (:obj:`Dict`): A piece of data to be used for generating the |
|
prompt. |
|
|
|
Returns: |
|
:obj:`str`: The generated prompt. |
|
""" |
|
prompt = None |
|
if self.output_template is None: |
|
prompt = str(entry[self.output_column]) |
|
else: |
|
prompt = self.output_template.generate_item(entry) |
|
return prompt |
|
|
|
def generate_output_field_corpus(self, |
|
dataset: Union[Dataset, DatasetDict], |
|
split: Optional[str] = None) -> List[str]: |
|
"""Generate corpus for output field. |
|
|
|
Args: |
|
dataset (:obj:`Dataset` or :obj:`DatasetDict`): A |
|
:obj:`datasets.Dataset` or :obj:`datasets.DatasetDict` |
|
instance. |
|
split (:obj:`str`, optional): The split of the dataset to use. |
|
If :obj:`None`, the entire dataset will be used. Defaults to |
|
``None``. |
|
|
|
Returns: |
|
:obj:`List[str]`: A list of generated output field prompts. |
|
""" |
|
if split is not None: |
|
dataset = dataset[split] |
|
corpus = [] |
|
for entry in dataset: |
|
corpus.append(self.generate_output_field_prompt(entry)) |
|
return corpus |
|
|
|
def generate_input_output_field_prompt(self, entry: Dict) -> str: |
|
"""Generate a prompt for the input-output field based on the |
|
provided:obj:`entry` data. |
|
|
|
Args: |
|
entry (:obj:`Dict`): A piece of data to be used for generating the |
|
prompt. |
|
|
|
Returns: |
|
:obj:`str`: The generated prompt. |
|
""" |
|
prompt = None |
|
if self.input_output_template is None: |
|
prompt = ' '.join([entry[ctx] for ctx in self.input_columns] + |
|
[str(entry[self.output_column])]) |
|
else: |
|
prompt = self.input_output_template.generate_item(entry) |
|
return prompt |
|
|
|
def _check_dataset_reader(obj) -> 'DatasetReader': |
|
if isinstance(obj, DatasetReader): |
|
return obj |
|
else: |
|
raise TypeError(f'Expected a DatasetReader object, but got {obj}') |
|
|
|
def __len__(self): |
|
return len(self.dataset) |
|
|
|
def __getitem__(self, idx): |
|
return self.dataset[idx] |
|
|
|
def __repr__(self): |
|
return (f'DatasetReader({{\n dataset: {self.dataset},' |
|
f'\n input_columns: {self.input_columns},\n' |
|
f' output_columns: {self.output_column}\n}})') |
|
|
|
|
|
def load_partial_dataset( |
|
dataset: Dataset, |
|
size: Optional[Union[int, float, str]] = None) -> Dataset: |
|
"""Load a partial dataset. |
|
|
|
Args: |
|
dataset (Dataset): A :obj:`datasets.Dataset` instance. |
|
size (int or float or (int, int), optional): The size of the partial |
|
dataset to load. If None, the entire dataset will be loaded. |
|
If int or float, the random partial dataset will be loaded with the |
|
specified size. If str, the partial dataset will be loaded with the |
|
specified index list (e.g. "[:100]" for the first 100 examples, |
|
"[100:200]" for the second 100 examples, etc.). Defaults to None. |
|
""" |
|
total_size = len(dataset) |
|
index_list = list(range(total_size)) |
|
if isinstance(size, (int, float)): |
|
if size >= total_size or size <= 0: |
|
return dataset |
|
if size > 0 and size < 1: |
|
size = int(size * total_size) |
|
rand = random.Random(x=size) |
|
rand.shuffle(index_list) |
|
dataset = dataset.select(index_list[:size]) |
|
elif isinstance(size, str): |
|
dataset = dataset.select(eval(f'index_list{size}')) |
|
return dataset |
|
|
|
|
|
class DatasetEncoder(torch.utils.data.Dataset): |
|
|
|
def __init__(self, |
|
datalist: List, |
|
model_name=None, |
|
tokenizer=None) -> None: |
|
self.datalist = datalist |
|
if model_name is None and tokenizer is None: |
|
raise ValueError('model_name and tokenizer could not both be None') |
|
if tokenizer is not None: |
|
self.tokenizer = tokenizer |
|
else: |
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
self.tokenizer.pad_token = self.tokenizer.eos_token |
|
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id |
|
self.tokenizer.padding_side = 'left' |
|
self.encode_dataset = [] |
|
self.init_dataset() |
|
self.datalist_length = len(self.encode_dataset) |
|
|
|
def init_dataset(self): |
|
for idx, data in enumerate(self.datalist): |
|
tokenized_data = self.tokenizer.encode_plus(data, |
|
truncation=True, |
|
return_tensors='pt', |
|
verbose=False) |
|
self.encode_dataset.append({ |
|
'input_ids': |
|
tokenized_data.input_ids[0], |
|
'attention_mask': |
|
tokenized_data.attention_mask[0], |
|
'metadata': { |
|
'id': idx, |
|
'len': len(tokenized_data.input_ids[0]), |
|
'text': data |
|
} |
|
}) |
|
|
|
def __len__(self): |
|
return self.datalist_length |
|
|
|
def __getitem__(self, idx): |
|
return self.encode_dataset[idx] |
|
|
