Moved offsetting of topk scores out of the (traced) TopK module. This…

… allows sending requests of variable (#1031)

batch size to the same Translator/Model/BeamSearch instance.

CHANGELOG.md

@@ -11,6 +11,13 @@ Note that Sockeye has checks in place to not translate with an old model that wa
 
 Each version section may have subsections for: _Added_, _Changed_, _Removed_, _Deprecated_, and _Fixed_.
 
+## [3.1.6]
+
+### Changed
+
+- Moved offsetting of topk scores out of the (traced) TopK module. This allows sending requests of variable
+  batch size to the same Translator/Model/BeamSearch instance.
+
 ## [3.1.5]
 
 ### Changed

sockeye/__init__.py

@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '3.1.5'
+__version__ = '3.1.6'

sockeye/beam_search.py

@@ -410,12 +410,11 @@ def __init__(self, k: int) -> None:
         super().__init__()
         self.k = k
 
-    def forward(self, scores, offset):
+    def forward(self, scores):
         """
         Get the lowest k elements per sentence from a `scores` matrix.
 
         :param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
-        :param offset: Array to add to the hypothesis indices for offsetting in batch decoding.
         :return: The row indices, column indices and values of the k smallest items in matrix.
         """
         batch_times_beam, vocab_size = scores.size()
@@ -430,9 +429,6 @@ def forward(self, scores, offset):
 
         best_hyp_indices, best_word_indices = indices.div(vocab_size, rounding_mode='floor'), indices.fmod(vocab_size)
 
-        if batch_size > 1:
-            # Offsetting the indices to match the shape of the scores matrix
-            best_hyp_indices = best_hyp_indices + offset
         return best_hyp_indices, best_word_indices, values
 
 
@@ -909,8 +905,11 @@ def forward(self,
 
                 if self._traced_top is None:
                     logger.debug("Tracing _top")
-                    self._traced_top = pt.jit.trace(self._top, (scores, offset))
-                best_hyp_indices, best_word_indices, scores_accumulated = self._traced_top(scores, offset)
+                    self._traced_top = pt.jit.trace(self._top, (scores,))
+                best_hyp_indices, best_word_indices, scores_accumulated = self._traced_top(scores)
+                if batch_size > 1:
+                    # Offsetting the indices to match the shape of the scores matrix
+                    best_hyp_indices = best_hyp_indices + offset
 
             # Map from restricted to full vocab ids if needed
             if restrict_lexicon:

test/unit/test_beam_search.py

@@ -140,7 +140,10 @@ def test_topk_func(batch_size, beam_size, target_vocab_size):
     np_hyp, np_word, np_values = numpy_topk(scores, k=beam_size, offset=offset)
 
     topk = sockeye.beam_search.TopK(k=beam_size)
-    pt_hyp, pt_word, pt_values = topk(pt.tensor(scores), pt.tensor(offset))
+    pt_hyp, pt_word, pt_values = topk(pt.tensor(scores))
+    if batch_size > 1:
+        # Offsetting the indices to match the shape of the scores matrix
+        pt_hyp += pt.tensor(offset)
     assert onp.allclose(pt_hyp.detach().numpy(), np_hyp)
     assert onp.allclose(pt_word.detach().numpy(), np_word)
     assert onp.allclose(pt_values.detach().numpy(), np_values)