code: def test_decoder_state_generator():
      data = "One Ring to rule them all, One Ring to find them"
      max_seq_len_source = 30
      max_seq_len_target = 30
  
      vocabs = [build_vocab([data])]
      config_embed = sockeye.encoder.EmbeddingConfig(
          vocab_size=len(vocabs[0]), num_embed=16, dropout=0.0
      )
      config_encoder = sockeye.encoder.EncoderConfig(
          model_size=16,
          attention_heads=2,
          feed_forward_num_hidden=16,
          depth_key_value=16,
          act_type="relu",
          num_layers=2,
          dropout_attention=0.0,
          dropout_act=0.0,
          dropout_prepost=0.0,
          positional_embedding_type="fixed",
          preprocess_sequence="n",
          postprocess_sequence="n",
          max_seq_len_source=max_seq_len_source,
          max_seq_len_target=max_seq_len_target,
      )
      config_data = sockeye.data_io.DataConfig(
          data_statistics=None,
          max_seq_len_source=max_seq_len_source,
          max_seq_len_target=max_seq_len_target,
          num_source_factors=0,
          num_target_factors=0,
      )
      config = sockeye.model.ModelConfig(
          config_data=config_data,
          vocab_source_size=len(vocabs[0]),
          vocab_target_size=len(vocabs[0]),
          config_embed_source=config_embed,
          config_embed_target=config_embed,
          config_encoder=config_encoder,
          config_decoder=config_encoder,
      )
  
      with tempfile.TemporaryDirectory() as model_dir, tempfile.TemporaryDirectory() as data_dir:
          params_fname = os.path.join(model_dir, C.PARAMS_BEST_NAME)
  
          # create and save float32 model
          model = sockeye.model.SockeyeModel(config=config)
          assert model.dtype == pt.float32
          for param in model.parameters():
              assert param.dtype == pt.float32
          model.save_config(model_dir)
          model.save_version(model_dir)
          model.save_parameters(params_fname)
          model.eval()
  
          # add dummy sentence to data_path
          data_path = os.path.join(data_dir, "data.txt")
          data_file = open(data_path, "w")
          data_file.write(data + "\n")
          data_file.close()
  
          # create generator from mock model and vocab
          generator = DecoderStateGenerator(
              model,
              vocabs,
              vocabs,
              data_dir,
              max_seq_len_source,
              max_seq_len_target,
              "float32",
              "int32",
              pt.device("cpu"),
          )
          max_seq_len_target = min(
              max_seq_len_target + C.SPACE_FOR_XOS, max_seq_len_target
          )
          generator.num_states = (
              DecoderStateGenerator.probe_token_count(
                  data_path, max_seq_len_target
              )
          )
  
          # test init_store_file
          generator.init_store_file(generator.num_states)
          generator.dimension = 16
  
          # test generate_states_and_store
          data_paths = [data_path]
          generator.generate_states_and_store(data_paths, data_paths, 1)
  
          # check if state and word store files are there
          assert os.path.isfile(get_state_store_path(data_dir))
          assert os.path.isfile(get_word_store_path(data_dir))
  
          # test save_config
          generator.save_config()
  
          # check if the config content makes sense
          config = KNNConfig.load(get_config_path(data_dir))
          assert (
              config.index_size
              == DecoderStateGenerator.probe_token_count(
                  data_path, max_seq_len_target
              )
          )
          assert config.dimension == 16
          assert config.state_data_type == "float32"
          assert config.word_data_type == "int32"
          assert config.index_type == ""
          assert config.train_data_size == -1
  
file_path: sockeye/test/unit/test_knn.py
imports: import numpy
  import os
  import pytest
  import tempfile
  import torch
  import sockeye
  import sockeye.constants
  from sockeye.generate_decoder_states import NumpyMemmapStorage, DecoderStateGenerator
  from sockeye.knn import KNNConfig, FaissIndexBuilder, get_config_path, get_state_store_path, get_word_store_path
  from sockeye.vocab import build_vocab
label_desc: Use image node_id 2 for calling a global function with example usage: test_decoder_state_generator() without return types