code: def test_get_nvs_vocab_slice_ids():
      # Batch size 2
      # Note: the first 4 tokens are special tokens (PAD, UNK etc.)
      #                             0    1    2    3    4     5    6     7     8    9
      nvs_prediction = pt.tensor(
          [
              [0.1, 0.1, 0.1, 0.1, 0.7, 0.0, 0.8, 0.0, 0.0, 0.0],
              [0.1, 0.1, 0.1, 0.1, 0.55, 0.0, 0.49, 0.05, 0.0, 0.0],
          ]
      )
      expected_bow = pt.tensor([0, 1, 2, 3, 4, 6, C.EOS_ID, C.EOS_ID])
      (
          bow,
          output_vocab_size,
      ) = sockeye.beam_search._get_nvs_vocab_slice_ids(
          nvs_thresh=0.5, nvs_prediction=nvs_prediction
      )
      assert output_vocab_size == expected_bow.shape[0]
      pt.testing.assert_close(bow, expected_bow)
  
      # Batch size 1
      #                             0    1    2    3    4     5    6     7     8    9
      nvs_prediction = pt.tensor(
          [[0.1, 0.1, 0.1, 0.1, 0.7, 0.0, 0.0, 0.8, 0.0, 0.0]]
      )
      expected_bow = pt.tensor([0, 1, 2, 3, 4, 7, C.EOS_ID, C.EOS_ID])
      (
          bow,
          output_vocab_size,
      ) = sockeye.beam_search._get_nvs_vocab_slice_ids(
          nvs_thresh=0.5, nvs_prediction=nvs_prediction
      )
      assert output_vocab_size == expected_bow.shape[0]
      pt.testing.assert_close(bow, expected_bow)
  
      # Batch size 1 + higher thresh
      #                             0    1    2    3    4     5    6     7     8    9
      nvs_prediction = pt.tensor(
          [[0.1, 0.1, 0.1, 0.1, 0.7, 0.0, 0.0, 0.8, 0.0, 0.0]]
      )
      expected_bow = pt.tensor(
          [0, 1, 2, 3, C.EOS_ID, C.EOS_ID, C.EOS_ID, C.EOS_ID]
      )
      (
          bow,
          output_vocab_size,
      ) = sockeye.beam_search._get_nvs_vocab_slice_ids(
          nvs_thresh=0.9, nvs_prediction=nvs_prediction
      )
      assert output_vocab_size == expected_bow.shape[0]
      pt.testing.assert_close(bow, expected_bow)
  
      # Batch size 2 + target prefix
      # Note: the first 4 tokens are special tokens (PAD, UNK etc.)
      #                             0    1    2    3    4     5    6     7     8    9
      nvs_prediction = pt.tensor(
          [
              [0.1, 0.1, 0.1, 0.1, 0.7, 0.0, 0.8, 0.0, 0.0, 0.0],
              [0.1, 0.1, 0.1, 0.1, 0.55, 0.0, 0.49, 0.05, 0.0, 0.0],
          ]
      )
      target_prefix = pt.tensor([[8, 8], [8, 8]])
      expected_bow = pt.tensor([0, 1, 2, 3, 4, 6, 8, C.EOS_ID])
      (
          bow,
          output_vocab_size,
      ) = sockeye.beam_search._get_nvs_vocab_slice_ids(
          nvs_thresh=0.5,
          nvs_prediction=nvs_prediction,
          target_prefix=target_prefix,
      )
      assert output_vocab_size == expected_bow.shape[0]
      pt.testing.assert_close(bow, expected_bow)
  
      # Batch size 2 + blocking lexicon
      # Note: the first 4 tokens are special tokens (PAD, UNK etc.)
      #                             0    1    2    3    4     5    6     7     8    9
      nvs_prediction = pt.tensor(
          [
              [0.1, 0.1, 0.1, 0.1, 0.7, 0.0, 0.8, 0.0, 0.0, 0.0],
              [0.1, 0.1, 0.1, 0.1, 0.55, 0.0, 0.49, 0.05, 0.0, 0.0],
          ]
      )
      expected_bow = pt.tensor(
          [0, 1, 2, 3, 4, C.EOS_ID, C.EOS_ID, C.EOS_ID]
      )
      restrict_lexicon = sockeye.lexicon.StaticBlockLexicon(
          np.array([6])
      )
      (
          bow,
          output_vocab_size,
      ) = sockeye.beam_search._get_nvs_vocab_slice_ids(
          nvs_thresh=0.5,
          nvs_prediction=nvs_prediction,
          restrict_lexicon=restrict_lexicon,
      )
      assert output_vocab_size == expected_bow.shape[0]
      pt.testing.assert_close(bow, expected_bow)
  
imports: from typing import List, Optional
  from typing import Tuple
  import numpy
  import pytest
  import torch
  import numpy
  import sockeye.beam_search
  import sockeye.constants
  import sockeye.lexicon
  import sockeye.utils
label_desc: Use image node_id 14 for calling a global function with example usage: test_get_nvs_vocab_slice_ids() without return types