Create a pruned version of a SCROLLS task dataset containing only inputs that are less than `max_tokens` when tokenized by each tokenizer
(self)
| 172 | return self.doc_to_text(self._process_doc(sample)[0]) |
| 173 | |
| 174 | def prune(self): |
| 175 | """Create a pruned version of a SCROLLS task dataset containing only inputs |
| 176 | that are less than `max_tokens` when tokenized by each tokenizer |
| 177 | """ |
| 178 | |
| 179 | tokenizers = [AutoTokenizer.from_pretrained(tokenizer) for tokenizer in self.PRUNE_TOKENIZERS] |
| 180 | cache = {} |
| 181 | |
| 182 | def _filter(sample): |
| 183 | text = self._get_prune_text(sample) |
| 184 | cached = cache.get(text, None) |
| 185 | if cached is None: |
| 186 | for tokenizer in tokenizers: |
| 187 | if len(tokenizer(text).input_ids) > self.PRUNE_MAX_TOKENS: |
| 188 | cache[text] = False |
| 189 | return False |
| 190 | cache[text] = True |
| 191 | return True |
| 192 | else: |
| 193 | return cached |
| 194 | |
| 195 | self.dataset = self.dataset.filter(_filter, num_proc=self.PRUNE_NUM_PROC) |
| 196 | |
| 197 | def doc_to_target(self, doc): |
| 198 | return " " + ", ".join(doc["outputs"]) |