| 342 | |
| 343 | |
| 344 | class BeamHypotheses: |
| 345 | def __init__(self, num_beams: int, max_length: int, length_penalty: float, early_stopping: bool): |
| 346 | """ |
| 347 | Initialize n-best list of hypotheses. |
| 348 | """ |
| 349 | self.max_length = max_length - 1 # ignoring bos_token |
| 350 | self.length_penalty = length_penalty |
| 351 | self.early_stopping = early_stopping |
| 352 | self.num_beams = num_beams |
| 353 | self.beams = [] |
| 354 | self.worst_score = 1e9 |
| 355 | |
| 356 | def __len__(self): |
| 357 | """ |
| 358 | Number of hypotheses in the list. |
| 359 | """ |
| 360 | return len(self.beams) |
| 361 | |
| 362 | def add(self, hyp: torch.LongTensor, sum_logprobs: float, mems=None): |
| 363 | """ |
| 364 | Add a new hypothesis to the list. |
| 365 | """ |
| 366 | score = sum_logprobs / (max(hyp.shape[-1], 1) ** self.length_penalty) |
| 367 | if len(self) < self.num_beams or score > self.worst_score: |
| 368 | self.beams.append((score, hyp, mems)) |
| 369 | if len(self) > self.num_beams: |
| 370 | sorted_next_scores = sorted([(s, idx) for idx, (s, _, _) in enumerate(self.beams)]) |
| 371 | del self.beams[sorted_next_scores[0][1]] |
| 372 | self.worst_score = sorted_next_scores[1][0] |
| 373 | else: |
| 374 | self.worst_score = min(score, self.worst_score) |
| 375 | |
| 376 | def is_done(self, best_sum_logprobs: float, cur_len: int) -> bool: |
| 377 | """ |
| 378 | If there are enough hypotheses and that none of the hypotheses being generated can become better than the worst |
| 379 | one in the heap, then we are done with this sentence. |
| 380 | """ |
| 381 | |
| 382 | if len(self) < self.num_beams: |
| 383 | return False |
| 384 | elif self.early_stopping: |
| 385 | return True |
| 386 | else: |
| 387 | cur_score = best_sum_logprobs / cur_len ** self.length_penalty |
| 388 | ret = self.worst_score >= cur_score |
| 389 | return ret |
| 390 | |
| 391 | |
| 392 | class LogitsProcessor(ABC): |