| 201 | } |
| 202 | |
| 203 | mx::array SuppressBlank::apply(const mx::array &Logits, |
| 204 | const mx::array &Tokens) { |
| 205 | if (Tokens.shape(1) == SampleBegin) { |
| 206 | return Logits + Mask; |
| 207 | } |
| 208 | return Logits; |
| 209 | } |
| 210 | |
| 211 | // SuppressTokens implementation |
| 212 | SuppressTokens::SuppressTokens(const std::vector<int> &SuppressTokens, |