How to Extract Recommendations for a List of Users from Saved Building Deep Retrieval Models?

Hi,

I can get a list of recommendations for a list of users

scann_index= tfrs.layers.factorized_top_k.ScaNN(model.user_model, k=100)
scann_index.index(items.batch(100).map(model.candidate_model), items)

scann_index({"user_id": np.array(['42','24']), "timestamp": np.array([1616689590,1616689590])})

but it doesn't work in saved model

# Save the index.
  scann_index.save(
      path,
      options=tf.saved_model.SaveOptions(namespace_whitelist=["Scann"])
  )

  # Load it back; can also be done in TensorFlow Serving.
  loaded = tf.keras.models.load_model(path)

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (3 total):
    * {'user_id': <tf.Tensor 'queries_1:0' shape=(2,) dtype=string>, 'timestamp': <tf.Tensor 'queries:0' shape=(2,) dtype=int64>}
    * None
    * False
  Keyword arguments: {}

Expected these arguments to match one of the following 4 option(s):

Option 1:
  Positional arguments (3 total):
    * {'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='queries/user_id'), 'timestamp': TensorSpec(shape=(), dtype=tf.float32, name='queries/timestamp')}
    * None
    * False
  Keyword arguments: {}

Option 2:
  Positional arguments (3 total):
    * {'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='user_id'), 'timestamp': TensorSpec(shape=(), dtype=tf.float32, name='timestamp')}
    * None
    * False
  Keyword arguments: {}

Option 3:
  Positional arguments (3 total):
    * {'timestamp': TensorSpec(shape=(), dtype=tf.float32, name='queries/timestamp'), 'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='queries/user_id')}
    * None
    * True
  Keyword arguments: {}

Option 4:
  Positional arguments (3 total):
    * {'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='user_id'), 'timestamp': TensorSpec(shape=(), dtype=tf.float32, name='timestamp')}
    * None
    * True
  Keyword arguments: {}

Start from the ScaNN indexing, save/load snippets in the issue and inspect the SavedModel signature options in the reported error. Reproduce the call with a batch of user_id and timestamp values, then determine whether the loaded model accepts that input shape and dtype. Done means the saved model can return recommendations for a list of users, or the limitation is documented clearly.