Restore the original batch_id before calling the callback

include/ctranslate2/generation.h

@@ -111,6 +111,22 @@ namespace ctranslate2 {
     }
   };
 
+  template <typename Options>
+  Options restore_batch_ids_in_callback(Options options, const std::vector<size_t>& example_index) {
+    if (options.callback) {
+      std::function<bool(GenerationStepResult)> wrapped_callback =
+        [&example_index, callback = std::move(options.callback)]
+        (GenerationStepResult step_result) {
+          step_result.batch_id = example_index[step_result.batch_id];
+          return callback(std::move(step_result));
+        };
+
+      options.callback = std::move(wrapped_callback);
+    }
+
+    return options;
+  }
+
   class ResolveEndToken {
   private:
     const Vocabulary& _vocabulary;

python/tests/test_translator.py

@@ -228,6 +228,36 @@ def _callback(step_result):
     assert len(hypotheses) == 3
 
 
+def test_callback_batch_id():
+    # The method will internally sort the input from longest to shortest,
+    # but we check that the returned batch ids match the user input.
+
+    source = [
+        ["ن"] * 1,
+        ["ن"] * 2,
+        ["ن"] * 3,
+    ]
+
+    target_prefix = [
+        ["a"],
+        ["b"],
+        ["c"],
+    ]
+
+    def _callback(step_result):
+        assert step_result.token == target_prefix[step_result.batch_id][0]
+        return True
+
+    translator = _get_transliterator()
+    translator.translate_batch(
+        source,
+        target_prefix,
+        max_batch_size=2,
+        beam_size=1,
+        callback=_callback,
+    )
+
+
 def test_file_translation(tmp_dir):
     input_path = str(tmp_dir.join("input.txt"))
     output_path = str(tmp_dir.join("output.txt"))

src/generator.cc

@@ -15,7 +15,9 @@ namespace ctranslate2 {
       batch_type,
       [options](models::SequenceGeneratorReplica& generator, const Batch& batch) {
         spdlog::debug("Running batch generation on {} examples", batch.num_examples());
-        auto results = generator.generate(batch.get_stream(0), options);
+        auto results = generator.generate(
+          batch.get_stream(0),
+          restore_batch_ids_in_callback(options, batch.example_index));
         spdlog::debug("Finished batch generation");
         return results;
       });

src/translator.cc

@@ -184,7 +184,9 @@ namespace ctranslate2 {
                   const Batch& batch,
                   const TranslationOptions& options) {
     spdlog::debug("Running batch translation on {} examples", batch.num_examples());
-    auto results = model.translate(batch.get_stream(0), batch.get_stream(1), options);
+    auto results = model.translate(batch.get_stream(0),
+                                   batch.get_stream(1),
+                                   restore_batch_ids_in_callback(options, batch.example_index));
     spdlog::debug("Finished batch translation");
     return results;
   }