| 261 | } |
| 262 | |
| 263 | void BiasedDecoder::decode(const dim_t cur_batch_size, |
| 264 | const size_t step, |
| 265 | const std::vector<dim_t>& batch_offset, |
| 266 | const std::vector<std::vector<bool>>& beams_diverged_from_prefix, |
| 267 | const StorageView& logits, |
| 268 | StorageView& log_probs) { |
| 269 | const dim_t num_beams = logits.dim(0); |
| 270 | const Device device = logits.device(); |
| 271 | const DataType dtype = logits.dtype(); |
| 272 | |
| 273 | if (_spare_beam.dtype() != dtype || _spare_beam.device() != device) { |
| 274 | _spare_beam = StorageView(device, dtype); |
| 275 | } |
| 276 | |
| 277 | std::vector<StorageView> logit_beam_view_storage(num_beams, StorageView(device, dtype)); |
| 278 | std::vector<StorageView*> logit_beam_views(num_beams); |
| 279 | std::vector<StorageView> log_prob_beam_view_storage(num_beams, StorageView(device, dtype)); |
| 280 | std::vector<StorageView*> log_prob_beam_views(num_beams); |
| 281 | for (dim_t i = 0; i < num_beams; ++i) { |
| 282 | logit_beam_views[i] = &(logit_beam_view_storage[i]); |
| 283 | log_prob_beam_views[i] = &(log_prob_beam_view_storage[i]); |
| 284 | } |
| 285 | ops::Split(0, /*no_copy=*/true)(logits, logit_beam_views); |
| 286 | log_probs.resize_as(logits); |
| 287 | log_probs.reshape(logits.shape()); |
| 288 | ops::Split(0, /*no_copy=*/true)(log_probs, log_prob_beam_views); |
| 289 | |
| 290 | // Scalar's need to be allocated on CPUs. |
| 291 | StorageView scalar_discount(1 - _prefix_bias_beta, Device::CPU); |
| 292 | assert (num_beams % cur_batch_size == 0); |
| 293 | const dim_t cur_beam_size = num_beams / cur_batch_size; |
| 294 | for (dim_t b = 0; b < num_beams; ++b) { |
| 295 | StorageView &logit_beam = *(logit_beam_views[b]); |
| 296 | StorageView &log_prob_beam = *(log_prob_beam_views[b]); |
| 297 | const dim_t index_batch = b / cur_beam_size; |
| 298 | const dim_t index_beam = b % cur_beam_size; |
| 299 | const auto& prefix = _prefix_ids[batch_offset[index_batch]]; |
| 300 | if (static_cast<size_t>(step) < prefix.size() |
| 301 | && !beams_diverged_from_prefix[index_batch][index_beam]) { |
| 302 | ops::SoftMax()(logit_beam, log_prob_beam); |
| 303 | ops::Mul()(log_prob_beam, |
| 304 | scalar_discount.to(log_prob_beam.dtype()), |
| 305 | _spare_beam); |
| 306 | const size_t biased_word_id = prefix[step]; |
| 307 | StorageView spare_scalar_view; |
| 308 | TYPE_DISPATCH( |
| 309 | _spare_beam.dtype(), |
| 310 | spare_scalar_view = StorageView({1}, _spare_beam.data<T>() + biased_word_id, device)); |
| 311 | const StorageView spare_scalar_copy(spare_scalar_view); |
| 312 | StorageView beta_scalar; |
| 313 | TYPE_DISPATCH( |
| 314 | _spare_beam.dtype(), |
| 315 | // Scalar's need to be allocated on CPUs. |
| 316 | beta_scalar = StorageView(static_cast<T>(_prefix_bias_beta), Device::CPU)); |
| 317 | ops::Add()(spare_scalar_copy, beta_scalar, spare_scalar_view); |
| 318 | ops::Log()(_spare_beam, log_prob_beam); |
| 319 | } else { |
| 320 | ops::LogSoftMax()(logit_beam, log_prob_beam); |