| 6 | |
| 7 | |
| 8 | class PromptDatasetProcessor(object): |
| 9 | def __init__( |
| 10 | self, |
| 11 | tokenize: Callable, |
| 12 | pad_token: int, |
| 13 | keep_order: bool = False, |
| 14 | max_seq_len: int = 2048, |
| 15 | sliding_stride: int = 200, |
| 16 | discard_overlong: bool = True, |
| 17 | eod_token: int = None, |
| 18 | preprocess: Callable = None, |
| 19 | ): |
| 20 | super(PromptDatasetProcessor, self).__init__() |
| 21 | self._keep_order = keep_order |
| 22 | self._max_seq_len = max_seq_len |
| 23 | self._sliding_stride = sliding_stride |
| 24 | self._tokenize = tokenize |
| 25 | self._pad_token = pad_token |
| 26 | self._discard_overlong = discard_overlong |
| 27 | self._eod_token = eod_token |
| 28 | self._preprocess = preprocess |
| 29 | |
| 30 | self.doc_processed = 0 |
| 31 | self.doc_generated = 0 |
| 32 | self.start_time = 0 |
| 33 | |
| 34 | def pad_seq(self, prompt_tokens: List[int], code_tokens: List[int], extra: dict = None) -> Dict[str, List[int]]: |
| 35 | total_length = len(prompt_tokens) + len(code_tokens) |
| 36 | assert total_length <= self._max_seq_len, f"padding sequence: {total_length} > {self._max_seq_len}" |
| 37 | pad_len = self._max_seq_len - total_length |
| 38 | input_ids = prompt_tokens + code_tokens + [self._pad_token] * pad_len |
| 39 | attention_mask = [1] * len(prompt_tokens) + [1] * len(code_tokens) + [0] * pad_len |
| 40 | labels = [-100] * len(prompt_tokens) + code_tokens + [-100] * pad_len |
| 41 | |
| 42 | return { |
| 43 | "input_ids": input_ids, |
| 44 | "attention_mask": attention_mask, |
| 45 | "labels": labels, |
| 46 | } |
| 47 | |
| 48 | def process_sample(self, sample: PromptSample) -> Iterable[Dict[str, List[int]]]: |
| 49 | """ |
| 50 | Process a sample. |
| 51 | """ |
| 52 | prompt_tokens = self._tokenize(sample.prompt) |
| 53 | code_tokens = self._tokenize(sample.code) |
| 54 | |
| 55 | if self._eod_token is not None: |
| 56 | code_tokens.append(self._eod_token) |
| 57 | |
| 58 | if len(prompt_tokens) + len(code_tokens) > self._max_seq_len: |
| 59 | if self._discard_overlong: |
| 60 | return |
| 61 | for p, t in sliding_window(prompt_tokens, code_tokens, self._max_seq_len, self._sliding_stride, self._sliding_stride): |
| 62 | yield self.pad_seq(p, t) |
| 63 | else: |
| 64 | yield self.pad_seq(prompt_tokens, code_tokens, extra=sample.extra) |
| 65 | |