| 39 | const float disable_value = std::numeric_limits<float>::lowest()); |
| 40 | |
| 41 | void add(dim_t batch_id, dim_t token_id) { |
| 42 | const auto flat_index = batch_id * _vocabulary_size + token_id; |
| 43 | |
| 44 | if (_logits_data) { |
| 45 | // On CPU we directly assign the value. |
| 46 | _logits_data[flat_index] = _disable_value; |
| 47 | |
| 48 | } else { |
| 49 | // On GPU we prepare a list of unique index to disable. |
| 50 | const auto it = std::lower_bound(_flat_indices.begin(), _flat_indices.end(), flat_index); |
| 51 | if (it == _flat_indices.end() || *it != flat_index) |
| 52 | _flat_indices.insert(it, flat_index); |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | // Disable a token for all batches. |
| 57 | void add(dim_t token_id) { |