title: Saving a model that uses tensorflow.lookup.StaticVocabularyTable in .pb format in Tensorflow 2
body: <p>I am building a model that accepts as input a 2d array of string tokens, then uses a lookup table to get the assigned indices of the input tokens in the vocabulary. The model then uses those indices to compute an embedded representation of the input tokens by fetching associated token embeddings and adding them together. The compounded embedding is then compared agains another matrix using a nearest-neighbors lookup and then the indices of the top-k most similar entries are returned.</p>
  <p>The model is saved in .pb format and is then used in a container running the TensorFlow Serving image for inference.</p>
  <p>At the moment I have something that works just fine in TensorFlow 1.15, however I am trying to migrate my code to TensorFlow 2.4 and can't find a way to make it work.</p>
  <p>Here is a slightly modified version of the code I am working with at the moment in TensorFlow 1.15</p>
  <pre><code>import tensorflow as tf
  
  
  graph = tf.get_default_graph()
  session = tf.Session()
  tf.global_variables_initializer()
  
  
  vocabulary = ['one', 'two', 'three', 'four', 'five', 'six']
  embedding_dimension = 512
  n_tokens = len(vocabulary)
  token_embeddings = np.random.random((n_tokens, embedding_dimension))
  matrix = np.random.random((100, embedding_dimension))
  
  
  lookup_table_initializer = tf.lookup.KeyValueTensorInitializer(vocabulary, np.arange(n_tokens))
  lookup_table = tf.lookup.StaticVocabularyTable(lookup_table_initializer, num_oov_buckets=1)
  
  
  token_embeddings_with_oov_token = np.vstack([token_embeddings, np.zeros(embedding_dimension)])
  token_embeddings_tensor = tf.convert_to_tensor(token_embeddings_with_oov_token, dtype=tf.float32)
  matrix = tf.convert_to_tensor(matrix, dtype=tf.float32)
  
  model_input = tf.placeholder(tf.string, [None, None], name=&quot;input&quot;)
  input_tokens_indices = lookup_table.lookup(model_input)
  
  input_token_indices_one_hot = tf.one_hot(input_tokens_indices, tf.dtypes.cast(value, dtype=np.int32)(lookup_table.size()))
  
  encoded_text = tf.math.reduce_sum(input_token_indices_one_hot, axis=1, keepdims=True)
  embedded_text = tf.linalg.matmul(encoded_text, token_embeddings_tensor)
  
  embedded_text_pooled = tf.math.reduce_sum(embedded_text, axis=1)
  embedded_text_normed = tf.divide(embedded_text_pooled, tf.norm(embedded_text_pooled, ord=2))
  neighbors = tf.linalg.matmul(embedded_text_normed, product_embeddings_tensor, transpose_b=True, name=&quot;output&quot;)
  
  tf.saved_model.simple_save(
      session,
      &quot;model.pb&quot;,
      inputs={&quot;input&quot;: model_input},
      outputs={&quot;output&quot;: neighbors},
      legacy_init_op=tf.tables_initializer(),
  )
  
  </code></pre>
  <p>The issue that I am facing is when converting the above code to TensorFlow 2. First of all, the <code>tf.placeholder</code> is no more and I have read on other posts suggestions to replace that with <code>tf.keras.layers.Input((), dtype=tf.dtypes.string)</code>, however then I get an error when I try to carry out the <code>lookup_table.lookup()</code> step, as apparently I cannot pass a symbolic tensor to that function. As a result I am stuck and do not know which way to proceed to make my model compatible with tf2 and after hours searching online for solutions I can't seem to find something that works.</p>