code: def test_topk_lexicon():
      lexicon = [
          "a\ta\t-0.6931471805599453",
          "a\tb\t-1.2039728043259361",
          "a\tc\t-1.6094379124341003",
          "b\tb\t0.0",
      ]
      vocab_list = ["a", "b", "c"]
      vocab = dict(
          (y, x) for (x, y) in enumerate(C.VOCAB_SYMBOLS + vocab_list)
      )
      k = 2
      lex = sockeye.lexicon.TopKLexicon(vocab, vocab)
  
      # Create from known lexicon
      with TemporaryDirectory(prefix="test_topk_lexicon.") as work_dir:
          # Write fast_align format lex table
          input_lex_path = os.path.join(work_dir, "input.lex")
          with open(input_lex_path, "w") as out:
              for line in lexicon:
                  print(line, file=out)
          # Use fast_align lex table to build top-k lexicon
          lex.create(input_lex_path, k)
  
          # Test against known lexicon
          expected = np.zeros(
              (len(C.VOCAB_SYMBOLS) + len(vocab_list), k),
              dtype=np.int32,
          )
          # a -> special + a b
          expected[len(C.VOCAB_SYMBOLS), :2] = [
              len(C.VOCAB_SYMBOLS),
              len(C.VOCAB_SYMBOLS) + 1,
          ]
          # b -> special + b
          expected[len(C.VOCAB_SYMBOLS) + 1, :1] = [
              len(C.VOCAB_SYMBOLS) + 1
          ]
          assert np.all(lex.lex == expected)
  
          # Test save/load
          expected_sorted = np.sort(expected, axis=1)
          json_lex_path = os.path.join(work_dir, "lex.json")
          lex.save(json_lex_path)
          lex.load(json_lex_path)
          assert np.all(lex.lex == expected_sorted)
  
          # Test lookup
          trg_ids = lex.get_allowed_trg_ids(
              np.array([[vocab["a"], vocab["c"]]], dtype=np.int32)
          )
          expected = np.array(
              [
                  vocab[symbol]
                  for symbol in C.VOCAB_SYMBOLS + ["a", "b"]
              ],
              dtype=np.int32,
          )
          assert np.all(trg_ids == expected)
  
          trg_ids = lex.get_allowed_trg_ids(
              np.array([[vocab["b"]]], dtype=np.int32)
          )
          expected = np.array(
              [vocab[symbol] for symbol in C.VOCAB_SYMBOLS + ["b"]],
              dtype=np.int32,
          )
          assert np.all(trg_ids == expected)
  
          trg_ids = lex.get_allowed_trg_ids(
              np.array([[vocab["c"]]], dtype=np.int32)
          )
          expected = np.array(
              [vocab[symbol] for symbol in C.VOCAB_SYMBOLS],
              dtype=np.int32,
          )
          assert np.all(trg_ids == expected)
  
          # Test load with smaller k
          small_k = k - 1
          lex.load(json_lex_path, k=small_k)
          assert lex.lex.shape[1] == small_k
          trg_ids = lex.get_allowed_trg_ids(
              np.array([[vocab["a"]]], dtype=np.int32)
          )
          expected = np.array(
              [vocab[symbol] for symbol in C.VOCAB_SYMBOLS + ["a"]],
              dtype=np.int32,
          )
          assert np.all(trg_ids == expected)
  
          # Test load with larger k
          large_k = k + 1
          lex.load(json_lex_path, k=large_k)
          assert lex.lex.shape[1] == k
          trg_ids = lex.get_allowed_trg_ids(
              np.array([[vocab["a"], vocab["c"]]], dtype=np.int32)
          )
          expected = np.array(
              [
                  vocab[symbol]
                  for symbol in C.VOCAB_SYMBOLS + ["a", "b"]
              ],
              dtype=np.int32,
          )
          assert np.all(trg_ids == expected)
  
imports: import os
  from tempfile import TemporaryDirectory
  import numpy
  import sockeye.constants
  import sockeye.lexicon
label_desc: Use image node_id 1 for calling a global function with example usage: test_topk_lexicon() without return types