Repository: OpenNMT/CTranslate2 (C++)

Issue:
Broken prefix bias decoding when enable disable_unk after 2.15.0
Hi @guillaumekln,
I recently encountered an issue after upgrading ctranslate2 to the latest version. The beam search with prefix bias no long work with `--disable_unk`.
After some debugging and testing, I can confirm this bug comes from release 2.15.0, maybe specifically from #764.

## To Reproduce
Take fairseq en-de WMT16 model and convert it to ctranslate2 format based on the [doc here](https://opennmt.net/CTranslate2/guides/fairseq.html#wmt16-english-german)

Before the 2.15.0 release(tested with 2.14.0):
```python
In [1] import ctranslate2
In [2] translator = ctranslate2.Translator("en-de.wmt16/model_ct2/", device="cpu", compute_type="default")
In [3] text="Le@@ ad researchers say this may bring early detection of cancer , tu@@ ber@@ cu@@ lo@@ sis , H
   ...: IV and mal@@ aria to patients in low-@@ income countries , where the survival rates for ill@@ nesses 
   ...: such as breast cancer can be half those of ri@@ cher countries ."

In [4] translator.translate_batch([['Le@@', 'ad', 'researchers', 'say', 'this', 'may']], target_prefix=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies']], prefix_bias_beta=0.2)
Out[4] [TranslationResult(hypotheses=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies', 'kann']], scores=[], attention=[])]
In [5] translator.translate_batch([['Le@@', 'ad', 'researchers', 'say', 'this', 'may']], target_prefix=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies']], prefix_bias_beta=0.2, disable_unk=True)
Out[5] [TranslationResult(hypotheses=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies', 'kann']], scores=[], attention=[])]
```

From 2.15.0 onwards(validated with 2.15.0 and 2.21.1)
```python
In [4] translator.translate_batch([['Le@@', 'ad', 'researchers', 'say', 'this', 'may']], target_prefix=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies']], prefix_bias_beta=0.2)
Out[4] [TranslationResult(hypotheses=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies', 'kann']], scores=[], attention=[])]
In [5] translator.translate_batch([['Le@@', 'ad', 'researchers', 'say', 'this', 'may']], target_prefix=[['Führ@@', 'ende', 'Forscher', 'sagen', 'dies']], prefix_bias_beta=0.2, disable_unk=True)
Out[5] [TranslationResult(hypotheses=[['dam@@', 'alige', 'Spitzen@@', 'for@@', 'scher', 'sagen']], scores=[], attention=[])]
```

As you can see, once `disable_unk` is enabled, prefix bias decoding `prefix_bias_beta` providing `target_prefix` no longer works properly after 2.15.0, while it worked perfectly with 2.14.0.

Any idea?


Discussion:
Hi,

Thank you for reporting and locating the issue. This is indeed a bug introduced by the PR you referenced.

I will look how to best fix this issue.
Thanks for the fix! Having verified with some examples, it can reproduce the same result as pre 2.15.
Any chance for a patch release including this fix?
There will be a new version by the end of the week.



Patch:
diff --git a/include/ctranslate2/decoding.h b/include/ctranslate2/decoding.h
index 1d0b7395d..892720310 100644
--- a/include/ctranslate2/decoding.h
+++ b/include/ctranslate2/decoding.h
@@ -72,19 +72,21 @@ namespace ctranslate2 {
 
   class BiasedDecoder {
   public:
-    BiasedDecoder() = default;
+    BiasedDecoder(const float prefix_bias_beta,
+                  const std::vector<std::vector<size_t>>& prefix_ids,
+                  const std::vector<size_t>* output_ids_map = nullptr);
 
     void
-    decode(const float prefix_bias_beta,
-           const dim_t cur_batch_size,
+    decode(const dim_t cur_batch_size,
            const size_t step,
            const std::vector<dim_t>& batch_offset,
            const std::vector<std::vector<bool>>& beams_diverged_from_prefix,
-           const std::vector<std::vector<size_t>>& prefix_ids,
            const StorageView& logits,
            StorageView& log_probs);
   private:
     StorageView _spare_beam;
+    const float _prefix_bias_beta;
+    std::vector<std::vector<size_t>> _prefix_ids;
   };
 
 
diff --git a/src/decoding.cc b/src/decoding.cc
index 924f111e0..7ee735e59 100644
--- a/src/decoding.cc
+++ b/src/decoding.cc
@@ -302,13 +302,34 @@ namespace ctranslate2 {
       result.attention = index_vector(result.attention, idx);
   }
 
+  BiasedDecoder::BiasedDecoder(const float prefix_bias_beta,
+                               const std::vector<std::vector<size_t>>& prefix_ids,
+                               const std::vector<size_t>* output_ids_map)
+    : _prefix_bias_beta(prefix_bias_beta)
+    , _prefix_ids(prefix_ids)
+  {
+    if (output_ids_map) {
+      std::unordered_map<size_t, size_t> output_ids_reverse_map;
+      output_ids_reverse_map.reserve(output_ids_map->size());
+      for (size_t i = 0; i < output_ids_map->size(); ++i)
+        output_ids_reverse_map.emplace((*output_ids_map)[i], i);
+
+      for (auto& ids : _prefix_ids) {
+        for (auto& id : ids) {
+          auto it = output_ids_reverse_map.find(id);
+          if (it == output_ids_reverse_map.end())
+            id = std::numeric_limits<size_t>::max();
+          else
+            id = it->second;
+        }
+      }
+    }
+  }
 
-  void BiasedDecoder::decode(const float prefix_bias_beta,
-                             const dim_t cur_batch_size,
+  void BiasedDecoder::decode(const dim_t cur_batch_size,
                              const size_t step,
                              const std::vector<dim_t>& batch_offset,
                              const std::vector<std::vector<bool>>& beams_diverged_from_prefix,
-                             const std::vector<std::vector<size_t>>& prefix_ids,
                              const StorageView& logits,
                              StorageView& log_probs) {
     const dim_t num_beams = logits.dim(0);
@@ -333,7 +354,7 @@ namespace ctranslate2 {
     ops::Split(0, /*no_copy=*/true)(log_probs, log_prob_beam_views);
 
     // Scalar's need to be allocated on CPUs.
-    StorageView scalar_discount(1 - prefix_bias_beta, Device::CPU);
+    StorageView scalar_discount(1 - _prefix_bias_beta, Device::CPU);
     assert (num_beams % cur_batch_size == 0);
     const dim_t cur_beam_size = num_beams / cur_batch_size;
     for (dim_t b = 0; b < num_beams; ++b) {
@@ -341,9 +362,10 @@ namespace ctranslate2 {
       StorageView &log_prob_beam = *(log_prob_beam_views[b]);
       const dim_t index_batch = b / cur_beam_size;
       const dim_t index_beam = b % cur_beam_size;
-      const auto& prefix = prefix_ids[batch_offset[index_batch]];
+      const auto& prefix = _prefix_ids[batch_offset[index_batch]];
       if (static_cast<size_t>(step) < prefix.size()
-          && !(beams_diverged_from_prefix[index_batch][index_beam])) {
+          && !beams_diverged_from_prefix[index_batch][index_beam]
+          && prefix[step] != std::numeric_limits<size_t>::max()) {
         ops::SoftMax()(logit_beam, log_prob_beam);
         ops::Mul()(log_prob_beam,
                    scalar_discount.to(log_prob_beam.dtype()),
@@ -358,7 +380,7 @@ namespace ctranslate2 {
         TYPE_DISPATCH(
           _spare_beam.dtype(),
           // Scalar's need to be allocated on CPUs.
-          beta_scalar = StorageView(static_cast<T>(prefix_bias_beta), Device::CPU));
+          beta_scalar = StorageView(static_cast<T>(_prefix_bias_beta), Device::CPU));
         ops::Add()(spare_scalar_copy, beta_scalar, spare_scalar_view);
         ops::Log()(_spare_beam, log_prob_beam);
       } else {
@@ -457,7 +479,9 @@ namespace ctranslate2 {
     std::vector<std::vector<bool>> beams_diverged_from_prefix;
     bool bias_towards_prefix = prefix_ids && _prefix_bias_beta > 0;
     if (bias_towards_prefix) {
-      biased_decoder = std::make_unique<BiasedDecoder>();
+      biased_decoder = std::make_unique<BiasedDecoder>(_prefix_bias_beta,
+                                                       *prefix_ids,
+                                                       output_ids_map);
       beams_diverged_from_prefix.resize(batch_size, std::vector<bool>(_beam_size, false));
     }
     const bool use_hard_prefix = prefix_ids && !bias_towards_prefix;
@@ -488,12 +512,10 @@ namespace ctranslate2 {
 
       StorageView log_probs(dtype, device);
       if (bias_towards_prefix) {
-        biased_decoder->decode(_prefix_bias_beta,
-                               cur_batch_size,
+        biased_decoder->decode(cur_batch_size,
                                step,
                                batch_offset,
                                beams_diverged_from_prefix,
-                               *prefix_ids,
                                logits,
                                log_probs);
       } else {


Test patch:
diff --git a/tests/translator_test.cc b/tests/translator_test.cc
index b4cb12eca..46b759ce7 100644
--- a/tests/translator_test.cc
+++ b/tests/translator_test.cc
@@ -491,13 +491,46 @@ TEST_P(BiasedDecodingDeviceFPTest, OneBatchOneBeam) {
     const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{false}};
     const std::vector<dim_t> batch_offset = {0};
     const std::vector<std::vector<size_t>> prefix_ids = {{0}};
-    ctranslate2::BiasedDecoder biased_decoder;
-    biased_decoder.decode(prefix_bias_beta,
-                          batch_size,
+    ctranslate2::BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids);
+    biased_decoder.decode(batch_size,
+                          step,
+                          batch_offset,
+                          beams_diverged_from_prefix,
+                          logits.to(device).to(dtype),
+                          log_probs);
+
+    expect_storage_eq(log_probs.to_float(), expected_log_probs, 0.01);
+}
+
+TEST_P(BiasedDecodingDeviceFPTest, OneBatchOneBeamOutputMap) {
+    const Device device = GetParam().first;
+    const DataType dtype = GetParam().second;
+    const dim_t vocab_size = 2;
+    const dim_t batch_size = 1;
+    const dim_t beam_size = 1;
+    const float prefix_bias_beta = 0.35;
+
+    StorageView logits({batch_size * beam_size, 1, vocab_size},
+                       std::vector<float>{4, 6});
+    StorageView softmax;
+    ops::SoftMax()(logits, softmax);
+    std::vector<float> expected_log_probs_vec = {
+      std::log((1-prefix_bias_beta) * softmax.at<float>(0) + prefix_bias_beta),
+      std::log((1-prefix_bias_beta) * softmax.at<float>(1)),
+    };
+    StorageView expected_log_probs(logits.shape(), expected_log_probs_vec, device);
+
+    StorageView log_probs(device, dtype);
+    const size_t step = 0;
+    const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{false}};
+    const std::vector<dim_t> batch_offset = {0};
+    const std::vector<std::vector<size_t>> prefix_ids = {{64}};
+    const std::vector<size_t> output_ids_map = {64, 65};
+    ctranslate2::BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids, &output_ids_map);
+    biased_decoder.decode(batch_size,
                           step,
                           batch_offset,
                           beams_diverged_from_prefix,
-                          prefix_ids,
                           logits.to(device).to(dtype),
                           log_probs);
 
@@ -543,13 +576,11 @@ TEST_P(BiasedDecodingDeviceFPTest, TwoBatchesTwoBeams) {
     const size_t step = 0;
     const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{false, false}, {false, false}};
     const std::vector<dim_t> batch_offset = {0, 1};
-    BiasedDecoder biased_decoder;
-    biased_decoder.decode(prefix_bias_beta,
-                          batch_size,
+    BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids);
+    biased_decoder.decode(batch_size,
                           step,
                           batch_offset,
                           beams_diverged_from_prefix,
-                          prefix_ids,
                           logits.to(device).to(dtype),
                           log_probs);
 
@@ -573,13 +604,11 @@ TEST_P(BiasedDecodingDeviceFPTest, BeamDiverged) {
     const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{true}};
     const std::vector<dim_t> batch_offset = {0};
     const std::vector<std::vector<size_t>> prefix_ids = {{0}};
-    BiasedDecoder biased_decoder;
-    biased_decoder.decode(prefix_bias_beta,
-                          batch_size,
+    BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids);
+    biased_decoder.decode(batch_size,
                           step,
                           batch_offset,
                           beams_diverged_from_prefix,
-                          prefix_ids,
                           logits.to(dtype),
                           log_probs);
 
@@ -603,13 +632,11 @@ TEST_P(BiasedDecodingDeviceFPTest, TimeStepPastPrefix) {
     const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{false}};
     const std::vector<dim_t> batch_offset = {0};
     const std::vector<std::vector<size_t>> prefix_ids = {{0}};
-    BiasedDecoder biased_decoder;
-    biased_decoder.decode(prefix_bias_beta,
-                          batch_size,
+    BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids);
+    biased_decoder.decode(batch_size,
                           step,
                           batch_offset,
                           beams_diverged_from_prefix,
-                          prefix_ids,
                           logits.to(dtype),
                           log_probs);
 
@@ -638,13 +665,11 @@ TEST_P(BiasedDecodingDeviceFPTest, NonZeroTimestepBias) {
     const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{false}};
     const std::vector<dim_t> batch_offset = {0};
     const std::vector<std::vector<size_t>> prefix_ids = {{0, 1, 0}};
-    BiasedDecoder biased_decoder;
-    biased_decoder.decode(prefix_bias_beta,
-                          batch_size,
+    BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids);
+    biased_decoder.decode(batch_size,
                           step,
                           batch_offset,
                           beams_diverged_from_prefix,
-                          prefix_ids,
                           logits.to(device).to(dtype),
                           log_probs);
 
@@ -668,13 +693,11 @@ TEST_P(BiasedDecodingDeviceFPTest, NonZeroTimestepDiverge) {
     const std::vector<std::vector<bool>> beams_diverged_from_prefix = {{true}};
     const std::vector<dim_t> batch_offset = {0};
     const std::vector<std::vector<size_t>> prefix_ids = {{0, 1, 0}};
-    BiasedDecoder biased_decoder;
-    biased_decoder.decode(prefix_bias_beta,
-                          batch_size,
+    BiasedDecoder biased_decoder(prefix_bias_beta, prefix_ids);
+    biased_decoder.decode(batch_size,
                           step,
                           batch_offset,
                           beams_diverged_from_prefix,
-                          prefix_ids,
                           logits.to(dtype),
                           log_probs);