Class DatasetEncoder

opencompass/openicl/icl_dataset_reader.py:247–289

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]