Extract sparse representation of distribution with probs > threshold
(self, knn_probs)
| 153 | return knn_probs |
| 154 | |
| 155 | def sparsify_distribution(self, knn_probs): |
| 156 | """Extract sparse representation of distribution with probs > threshold""" |
| 157 | batch_size = knn_probs.shape[0] |
| 158 | |
| 159 | id_cnt_list = [] |
| 160 | token_id_list = [] |
| 161 | prob_list = [] |
| 162 | |
| 163 | for b in range(batch_size): |
| 164 | # Find indices where probability > threshold |
| 165 | valid_mask = knn_probs[b] > self.threshold |
| 166 | valid_ids = torch.nonzero(valid_mask).squeeze(-1) |
| 167 | valid_probs = knn_probs[b][valid_ids] |
| 168 | |
| 169 | # Sort by probability (descending) |
| 170 | sorted_indices = torch.argsort(valid_probs, descending=True) |
| 171 | sorted_ids = valid_ids[sorted_indices] |
| 172 | sorted_probs = valid_probs[sorted_indices] |
| 173 | |
| 174 | id_cnt_list.append(len(sorted_ids)) |
| 175 | token_id_list.append(sorted_ids.to(self.device)) |
| 176 | prob_list.append(sorted_probs.to(self.device).to(torch.float16)) # Convert to float16 |
| 177 | |
| 178 | return id_cnt_list, token_id_list, prob_list |
| 179 | |
| 180 | def _save_step_data(self, id_cnt, token_id, prob, label): |
| 181 | """Save data for current step using streaming Arrow format""" |