(weighted_scores, top_p=0.8, top_k=25)
| 168 | return top_ids |
| 169 | |
| 170 | def nucleus_sampling(weighted_scores, top_p=0.8, top_k=25): |
| 171 | prob, indices = [], [] |
| 172 | cum_prob = 0.0 |
| 173 | sorted_value, sorted_idx = weighted_scores.softmax(dim=0).sort(descending=True, stable=True) |
| 174 | for i in range(len(sorted_idx)): |
| 175 | # sampling both top-p and numbers. |
| 176 | if cum_prob < top_p and len(prob) < top_k: |
| 177 | cum_prob += sorted_value[i] |
| 178 | prob.append(sorted_value[i]) |
| 179 | indices.append(sorted_idx[i]) |
| 180 | else: |
| 181 | break |
| 182 | prob = torch.tensor(prob).to(weighted_scores) |
| 183 | indices = torch.tensor(indices, dtype=torch.long).to(weighted_scores.device) |
| 184 | top_ids = indices[prob.multinomial(1, replacement=True)] |
| 185 | return top_ids |
| 186 | |
| 187 | |
| 188 | def random_sampling(weighted_scores, decoded_tokens): |
no outgoing calls
no test coverage detected