| 47 | |
| 48 | |
| 49 | class EncodeDataset(Dataset): |
| 50 | |
| 51 | def __init__(self, data_args: DataArguments): |
| 52 | self.data_args = data_args |
| 53 | self.encode_data = load_dataset( |
| 54 | self.data_args.dataset_name, |
| 55 | self.data_args.dataset_config, |
| 56 | data_files=self.data_args.dataset_path, |
| 57 | split=self.data_args.dataset_split, |
| 58 | cache_dir=self.data_args.dataset_cache_dir, |
| 59 | ) |
| 60 | if self.data_args.dataset_number_of_shards > 1: |
| 61 | self.encode_data = self.encode_data.shard( |
| 62 | num_shards=self.data_args.dataset_number_of_shards, |
| 63 | index=self.data_args.dataset_shard_index, |
| 64 | ) |
| 65 | |
| 66 | def __len__(self): |
| 67 | return len(self.encode_data) |
| 68 | |
| 69 | def __getitem__(self, item) -> Tuple[str, str]: |
| 70 | text = self.encode_data[item] |
| 71 | text_id = text.get("corpus_id", text.get("docs_id", text.get("paper_id", None))) |
| 72 | formated_text = text['abstract'] |
| 73 | return text_id, formated_text |