| 112 | self.start_time = 0 |
| 113 | |
| 114 | def pad_seq(self, prompt_tokens: List[int], label: int, extra: dict = None) -> Dict[str, List[int]]: |
| 115 | total_length = len(prompt_tokens) |
| 116 | assert total_length <= self._max_seq_len, f"padding sequence: {total_length} > {self._max_seq_len}" |
| 117 | pad_len = self._max_seq_len - total_length |
| 118 | input_ids = prompt_tokens + [self._pad_token] * pad_len |
| 119 | attention_mask = [1] * len(prompt_tokens) + [0] * pad_len |
| 120 | label = [label] |
| 121 | |
| 122 | return { |
| 123 | "input_ids": input_ids, |
| 124 | "attention_mask": attention_mask, |
| 125 | "length": [len(prompt_tokens)], |
| 126 | "labels": label |
| 127 | } |
| 128 | def process_sample(self, sample: LabelSample) -> Iterable[Dict[str, List[int]]]: |
| 129 | """ |
| 130 | Process a sample. |