Building the query and item towers

Each tower is a small Keras model. A StringLookup turns customer or article ids into integer indices. An Embedding layer turns indices into dense vectors. Optional Dense layers compose those embeddings with context features. The output is a single vector per input.

recsys/training/two_tower.py
python
class QueryTower(tf.keras.Model):
    def __init__(self, user_ids, emb_dim, **kwargs):
        super().__init__(**kwargs)
        self.user_embedding = tf.keras.Sequential([
            StringLookup(vocabulary=user_ids, mask_token=None),
            tf.keras.layers.Embedding(len(user_ids) + 1, emb_dim),
        ])
        self.normalized_age = Normalization(axis=None)
        self.fnn = tf.keras.Sequential([
            tf.keras.layers.Dense(emb_dim, activation="relu"),
            tf.keras.layers.Dense(emb_dim),
        ])

    def call(self, inputs):
        return self.fnn(tf.concat([
            self.user_embedding(inputs["customer_id"]),
            tf.reshape(self.normalized_age(inputs["age"]), (-1, 1)),
            tf.reshape(inputs["month_sin"], (-1, 1)),
            tf.reshape(inputs["month_cos"], (-1, 1)),
        ], axis=1))

The query tower projects (customer_id, age, month_sin, month_cos) into a 16-dimensional vector.

recsys/training/two_tower.py
python
class ItemTower(tf.keras.Model):
    def __init__(self, item_ids, garment_groups, index_groups, emb_dim, **kwargs):
        super().__init__(**kwargs)
        self.item_embedding = tf.keras.Sequential([
            StringLookup(vocabulary=item_ids, mask_token=None),
            tf.keras.layers.Embedding(len(item_ids) + 1, emb_dim),
        ])
        self.garment_group_tokenizer = StringLookup(vocabulary=garment_groups, mask_token=None)
        self.index_group_tokenizer = StringLookup(vocabulary=index_groups, mask_token=None)
        self.fnn = tf.keras.Sequential([
            tf.keras.layers.Dense(emb_dim, activation="relu"),
            tf.keras.layers.Dense(emb_dim),
        ])

The item tower has the same shape but consumes article id, garment group, and index group.

Flashcards: Flashcards

Loading practice…

Quiz: Quiz

Loading practice…

Validation checklist: Before you train

Loading practice…