How to adapt tensorflow recommender system tutorial to own data? Issues with Dataset and MapDataset

I am working on a recommender system in tensorflow. What I am trying to do is something similar to tensorflow's quickstart example. However I cannot seem to understand how to replace the Dataset structure(s) with my own data correctly, as doing so raises errors either in the dataset mapping phase or in the model fitting phase. I am running Python 3.7.13 on Google Colab and Tensorflow 2.8.0.

So, let's say this is a music recommender. Note that my data is all integer IDs. In order to follow the tutorial, I limit my data in a similar manner

I figured that I can actually load my data with tf.data.Dataset.from_tensor_slices():

rating = tf.data.Dataset.from_tensor_slices(df[['song_id', 'user_id']].values)

songs = tf.data.Dataset.from_tensor_slices(df[['song_id']].values)

This works, so I go on to map the dataset:

rating = rating.map(lambda x:{'song_id':x['song_id'],'user_id':x['user_id']})
songs = songs.map(lambda x: x['song_id'])

However, this raises the following:

TypeError: Only integers, slices (:), ellipsis (...), tf.newaxis (None) and scalar tf.int32/tf.int64 tensors are valid indices, got 'song_id'

I am not sure as to why I need to map the dataset in the first place... I assume it's something tied to the default data structure used in the examples?

So let's say I don't map. I go on using IntegerLookup() instead of StringLookup(mask_token=None) to preprocess my data, since all I have is integers:

user_id_vocabulary = tf.keras.layers.IntegerLookup()
user_id_vocabulary.adapt(rating)
 
songs_vocabulary = tf.keras.layers.IntegerLookup()
songs_vocabulary.adapt(songs)

Then I build the model Class following the tutorial, just changing variable names, and define the users model, the songs model and the retrieval task:

class MyModel(tfrs.Model):
 
  def __init__(
      self,
      user_model: tf.keras.Model,
      song_model: tf.keras.Model,
      task: tfrs.tasks.Retrieval):
    super().__init__()
 
    # Set up user and song representations.
    self.user_model = user_model
    self.song_model = song_model
 
    # Set up a retrieval task.
    self.task = task
 
  def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:
    # Define how the loss is computed.
 
    user_embeddings = self.user_model(features["user_id"])
    song_embeddings = self.song_model(features["song_id"])
 
    return self.task(user_embeddings, song_embeddings)

users_model = tf.keras.Sequential([user_id_vocabulary, tf.keras.layers.Embedding(user_id_vocabulary.vocabulary_size(),64)])
                           
songs_model = tf.keras.Sequential([songs_vocabulary,                                   tf.keras.layers.Embedding(songs_vocabulary.vocabulary_size(),64)])
 
task = tfrs.tasks.Retrieval(metrics=tfrs.metrics.FactorizedTopK(
    rooms.batch(128).map(room_model)))

Lastly, I compile and fit the model:

model = MyModel(users_model,songs_model,task)
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.5))
model.fit(rating.batch(4096), epochs=3)

But this still raises the following on the .fit line:

TypeError: Only integers, slices (:), ellipsis (...), tf.newaxis (None) and scalar tf.int32/tf.int64 tensors are valid indices, got 'song_id'

What am I missing? Thanks in advance!