This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length` or meet the code generation stopping criteria.
| 143 | |
| 144 | |
| 145 | class CodeStoppingCriteria(StoppingCriteria): |
| 146 | """ |
| 147 | This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length` or meet the code generation stopping criteria. |
| 148 | """ |
| 149 | |
| 150 | def __init__( |
| 151 | self, |
| 152 | max_length: int, |
| 153 | micro_batch_size: int, |
| 154 | tokenizer, |
| 155 | dataset_type: str, |
| 156 | language_type: str, |
| 157 | prompt: str, |
| 158 | ): |
| 159 | self.max_length = max_length |
| 160 | self.tokenizer = tokenizer |
| 161 | self.dataset_type = dataset_type |
| 162 | self.language_type = language_type |
| 163 | self.prompt = prompt |
| 164 | self.stop_index = [-1 for _ in range(micro_batch_size)] |
| 165 | |
| 166 | def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: |
| 167 | for i, input_id in enumerate(input_ids): |
| 168 | if self.stop_index[i] > -1: |
| 169 | continue |
| 170 | code = self.tokenizer.decode(input_id) |
| 171 | code = code[len(self.prompt):] |
| 172 | if is_code_generation_finished( |
| 173 | code, |
| 174 | dataset_type=self.dataset_type, |
| 175 | language_type=self.language_type) or input_id.shape[-1] >= self.max_length: |
| 176 | self.stop_index[i] = len(code) + len(self.prompt) |
| 177 | if all([s != -1 for s in self.stop_index]): |
| 178 | return True |
| 179 | |
| 180 | return False |
| 181 | |
| 182 | |
| 183 | def run_generation_distributed(args, model, tokenizer): |