| 8 | |
| 9 | @dataclass |
| 10 | class DataCollator: |
| 11 | tokenizer: PreTrainedTokenizerBase |
| 12 | model: Optional[Any] = None |
| 13 | padding: Union[bool, str, PaddingStrategy] = True # ‘longest’ |
| 14 | max_prompt_len: Optional[int] = None |
| 15 | max_ans_len: Optional[int] = None |
| 16 | pad_to_multiple_of: Optional[int] = 1 |
| 17 | label_pad_token_id: int = -100 |
| 18 | return_tensors: str = "pt" |
| 19 | inference: bool = False |
| 20 | demonstrations: Optional[Any] = None |
| 21 | task: str = None |
| 22 | |
| 23 | def __call__(self, batch, return_tensors=None): |
| 24 | if return_tensors is None: |
| 25 | return_tensors = self.return_tensors |
| 26 | model_inputs = self.decoder_call(batch, self.return_tensors) |
| 27 | |
| 28 | return model_inputs |
| 29 | |
| 30 | # only support left padding for now |
| 31 | def tokenize(self, sentence, cutoff_len, add_bos_token=True, add_eos_token=True): |
| 32 | # there's probably a way to do this with the tokenizer settings |
| 33 | # but again, gotta move fast |
| 34 | result = self.tokenizer( |
| 35 | sentence, |
| 36 | truncation=True, |
| 37 | max_length=cutoff_len, |
| 38 | add_special_tokens=False, |
| 39 | padding=False, |
| 40 | return_tensors=None, |
| 41 | ) |
| 42 | |
| 43 | if ( |
| 44 | len(result["input_ids"]) < cutoff_len |
| 45 | and add_eos_token |
| 46 | ): |
| 47 | result["input_ids"].append(self.tokenizer.eos_token_id) |
| 48 | result["attention_mask"].append(1) |
| 49 | |
| 50 | if ( |
| 51 | len(result["input_ids"]) < cutoff_len |
| 52 | and add_bos_token |
| 53 | ): |
| 54 | result["input_ids"] = [self.tokenizer.bos_token_id] + result["input_ids"] |
| 55 | result["attention_mask"] = [1] + result["attention_mask"] |
| 56 | |
| 57 | result["labels"] = result["input_ids"].copy() |
| 58 | |
| 59 | return result |
| 60 | |
| 61 | # support decoder-only models for left padding |
| 62 | def decoder_call(self, batch, return_tensors): |
| 63 | # to fix the bug |
| 64 | sources = [] |
| 65 | gts = [] |
| 66 | tokenized_sources = [] |
| 67 | label_lens = [] |
no outgoing calls
no test coverage detected