| 946 | |
| 947 | |
| 948 | class BeamHypotheses(object): |
| 949 | def __init__(self, num_beams, max_length, length_penalty, early_stopping): |
| 950 | """ |
| 951 | Initialize n-best list of hypotheses. |
| 952 | """ |
| 953 | self.max_length = max_length - 1 # ignoring bos_token |
| 954 | self.length_penalty = length_penalty |
| 955 | self.early_stopping = early_stopping |
| 956 | self.num_beams = num_beams |
| 957 | self.beams = [] |
| 958 | self.worst_score = 1e9 |
| 959 | |
| 960 | def __len__(self): |
| 961 | """ |
| 962 | Number of hypotheses in the list. |
| 963 | """ |
| 964 | return len(self.beams) |
| 965 | |
| 966 | def add(self, hyp, sum_logprobs): |
| 967 | """ |
| 968 | Add a new hypothesis to the list. |
| 969 | """ |
| 970 | score = sum_logprobs / len(hyp) ** self.length_penalty |
| 971 | if len(self) < self.num_beams or score > self.worst_score: |
| 972 | self.beams.append((score, hyp)) |
| 973 | if len(self) > self.num_beams: |
| 974 | sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.beams)]) |
| 975 | del self.beams[sorted_scores[0][1]] |
| 976 | self.worst_score = sorted_scores[1][0] |
| 977 | else: |
| 978 | self.worst_score = min(score, self.worst_score) |
| 979 | |
| 980 | def is_done(self, best_sum_logprobs, cur_len): |
| 981 | """ |
| 982 | If there are enough hypotheses and that none of the hypotheses being generated |
| 983 | can become better than the worst one in the heap, then we are done with this sentence. |
| 984 | """ |
| 985 | |
| 986 | if len(self) < self.num_beams: |
| 987 | return False |
| 988 | elif self.early_stopping: |
| 989 | return True |
| 990 | else: |
| 991 | cur_score = best_sum_logprobs / cur_len ** self.length_penalty |
| 992 | ret = self.worst_score >= cur_score |
| 993 | return ret |
no outgoing calls
no test coverage detected