How to make predictions by retrieval models with multiple features?

`
class UserModel(tf.keras.Model):

def init(self,use_timestamps,use_distance):
super().init()

self._use_timestamps = use_timestamps
self._use_distance = use_distance

self.user_embedding = tf.keras.Sequential([
    tf.keras.layers.StringLookup(
    vocabulary=unique_user_ids, mask_token=None),
    tf.keras.layers.Embedding(len(unique_user_ids) + 1, 32)
])

max_tokens = 10_000

self.Preference1_embedding = tf.keras.Sequential([
  tf.keras.layers.StringLookup(
      vocabulary=unique_Preference1, mask_token=None),
  tf.keras.layers.Embedding(len(unique_Preference1) + 1, 32)
])

self.Preference1_vectorizer = tf.keras.layers.TextVectorization(max_tokens=max_tokens)

self.Preference1_text_embedding = tf.keras.Sequential([
  self.Preference1_vectorizer,
  tf.keras.layers.Embedding(max_tokens, 32, mask_zero=True),
  tf.keras.layers.GlobalAveragePooling1D(),
])

self.Preference1_vectorizer.adapt(user.map(lambda x: x["Preference1"]))



self.Preference2_embedding = tf.keras.Sequential([
  tf.keras.layers.StringLookup(
      vocabulary=unique_Preference2, mask_token=None),
  tf.keras.layers.Embedding(len(unique_Preference2) + 1, 32)
])

self.Preference2_vectorizer = tf.keras.layers.TextVectorization(
max_tokens=max_tokens)

self.Preference2_text_embedding = tf.keras.Sequential([
  self.Preference2_vectorizer,
  tf.keras.layers.Embedding(max_tokens, 32, mask_zero=True),
  tf.keras.layers.GlobalAveragePooling1D(),
])

self.Preference2_vectorizer.adapt(user.map(lambda x: x["Preference2"]))


if use_timestamps:
  self.timestamp_embedding = tf.keras.Sequential([
      tf.keras.layers.Discretization(timestamp_buckets.tolist()),
      tf.keras.layers.Embedding(len(timestamp_buckets) + 1, 32),
  ])
  self.normalized_timestamp = tf.keras.layers.Normalization(
      axis=None
  )

  self.normalized_timestamp.adapt(timestamps)

if use_distance:
  self.distance_embedding = tf.keras.Sequential([
      tf.keras.layers.Discretization(Distance_buckets.tolist()),
      tf.keras.layers.Embedding(len(Distance_buckets) + 1, 32),
  ])
  self.normalized_distance = tf.keras.layers.Normalization(
      axis=None
  )

  self.normalized_timestamp.adapt(Distance)

def call(self, inputs):

if not self._use_timestamps:

return self.user_embedding(inputs["user_id"])

return tf.concat([
    self.user_embedding(inputs["UserID"]),
    self.Preference1_embedding(inputs["Preference1"]),
    self.Preference1_text_embedding(inputs["Preference1"]),
    self.Preference2_embedding(inputs["Preference2"]),
    self.Preference2_text_embedding(inputs["Preference2"]),       
    self.timestamp_embedding(inputs["Time"]),
    tf.reshape(self.normalized_timestamp(inputs["Time"]), (-1, 1)),
    self.distance_embedding(inputs["Distance"]),
    tf.reshape(self.normalized_timestamp(inputs["Distance"]), (-1, 1))],
    axis=1)

class LocationModel(tf.keras.Model):

def init(self):
super().init()

max_tokens = 10_000

self.location_embedding = tf.keras.Sequential([
  tf.keras.layers.StringLookup(
  vocabulary=unique_location_ids, mask_token=None),
  tf.keras.layers.Embedding(len(unique_location_ids) + 1, 32)
])

self.locationTag_embedding = tf.keras.Sequential([tf.keras.layers.StringLookup(
  vocabulary=unique_locationTag, mask_token=None),
  tf.keras.layers.Embedding(len(unique_locationTag) + 1, 32)
])



self.locationTag_vectorizer = tf.keras.layers.TextVectorization(
    max_tokens=max_tokens)

self.locationTag_text_embedding = tf.keras.Sequential([
  self.locationTag_vectorizer,
  tf.keras.layers.Embedding(max_tokens, 32, mask_zero=True),
  tf.keras.layers.GlobalAveragePooling1D(),
])

self.locationTag_vectorizer.adapt(location.map(lambda x: x["Tag of Location visited"]))


self.Description_embedding = tf.keras.Sequential([tf.keras.layers.StringLookup(vocabulary=unique_Description, mask_token=None),
  tf.keras.layers.Embedding(len(unique_Description) + 1, 32)])

self.Description_vectorizer = tf.keras.layers.TextVectorization(
    max_tokens=max_tokens)

self.Description_text_embedding = tf.keras.Sequential([
  self.Description_vectorizer,
  tf.keras.layers.Embedding(max_tokens, 32, mask_zero=True),
  tf.keras.layers.GlobalAveragePooling1D(),
])

self.Description_vectorizer.adapt(location.map(lambda x: x["Description of Location visited"]))

def call(self, inputs):
return tf.concat([
self.location_embedding(inputs['Name of Location']),
self.locationTag_embedding(inputs['Tag of Location visited']),
self.locationTag_text_embedding(inputs['Tag of Location visited']),
self.Description_embedding(inputs['Description of Location visited']),
self.Description_text_embedding(inputs['Description of Location visited']),
], axis=1)
class TravelModel(tfrs.models.Model):

def init(self, use_timestamps,use_distance):
super().init()

self.query_model = tf.keras.Sequential([
  UserModel(use_timestamps,use_distance),
  tf.keras.layers.Dense(32)
])

self.candidate_model = tf.keras.Sequential([
  LocationModel(),
  tf.keras.layers.Dense(32)
])


self.task = tfrs.tasks.Retrieval(
    metrics=tfrs.metrics.FactorizedTopK(
        candidates=location.batch(128).map(self.candidate_model),
    ),
)

def compute_loss(self, features, training=False):

query_embeddings = self.query_model({
    "UserID": features["UserID"],
    "Time": features["Time"],
    "Distance": features["Distance"],
    "Preference1": features["Preference1"],
    "Preference2": features["Preference2"]
})


location_embeddings = self.candidate_model({
    "Name of Location": features["Name of Location"],
    "Tag of Location visited": features["Tag of Location visited"],
    "Description of Location visited": features["Description of Location visited"]
})

return self.task(query_embeddings, location_embeddings)

tf.random.set_seed(42)
shuffled = data.shuffle(100_000, seed=42, reshuffle_each_iteration=False)

train = shuffled.take(int(len(data)*0.8))
test = shuffled.skip(int(len(data)*0.8)).take(int(len(data)*0.2)+1)

cached_train = train.shuffle(100_000).batch(2048)
cached_test = test.batch(4096).cache()

model = TravelModel(use_timestamps=True,use_distance=True)
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1))

model.fit(cached_train, epochs=30)

train_accuracy = model.evaluate(
cached_train, return_dict=True)["factorized_top_k/top_100_categorical_accuracy"]
test_accuracy = model.evaluate(
cached_test, return_dict=True)["factorized_top_k/top_100_categorical_accuracy"]

print(f"Top-100 accuracy (train): {train_accuracy:.2f}.")
print(f"Top-100 accuracy (test): {test_accuracy:.2f}.")

`

Start with the UserModel, LocationModel, and TravelModel definitions in the issue, then run the shown training and evaluation flow to reproduce the retrieval behavior with timestamp, distance, preference, and location features. Trace how query and candidate embeddings are passed to tfrs.tasks.Retrieval, and consider the work complete when multi-feature prediction behavior is clearly demonstrated and documented.