Example Using Gemma Pre-trained Embedding?

[rlcauvin]

The pre-trained embeddings example using the Universal Sentence Encoder is very helpful. I'm wondering how to do the same thing, but with a Gemma model loaded from Keras Hub.

[rstz]

Great suggestion. I know it's possible (we've experimented with it in the past), including an example in the docs is a good idea. If anyone wants to write something up, feel free to open a PR!

[rlcauvin]

I would be happy to write something up, except that I don't yet understand how to properly generate embeddings using the Gemma model.

I find it handy to load the Universal Sentence Encoder as a Keras layer and output embeddings this way:

use_layer = tfh.KerasLayer("https://tfhub.dev/google/universal-sentence-encoder/4")
use_layer(np.array(["i ate a lemon", "i ate an orange", "money doesn't grow on trees"]))

<tf.Tensor: shape=(3, 512), dtype=float32, numpy=
array([[-0.04931547, -0.00642865,  0.0258174 , ..., -0.02363558,
         0.07640501, -0.06680857],
       [-0.04869086, -0.00624514,  0.00913326, ...,  0.02101696,
         0.06567055, -0.07759373],
       [-0.03008169, -0.08008684, -0.00631541, ...,  0.02706281,
        -0.06680012, -0.05502592]], dtype=float32)>

I tried to do something similar with the Gemma model in Keras Hub, but the shape of the embeddings it outputs is different:

class GemmaEncoder(tf.keras.layers.Layer):

  def __init__(self):
    
    super().__init__()
    self.gemma_lm = keras_hub.models.GemmaCausalLM.from_preset("gemma2_2b_en")

  def call(self, inputs):
    
    preprocessed = self.gemma_lm.preprocessor.generate_preprocess(inputs)
    embeddings = self.gemma_lm.backbone.token_embedding(preprocessed["token_ids"])
    
    return embeddings

gse_layer = GemmaEncoder()
gse_layer(np.array(["i ate a lemon", "i ate an orange", "money doesn't grow on trees"]))

<tf.Tensor: shape=(3, 1024, 2304), dtype=float32, numpy=
array([[[ 0.00012589, -0.00585938,  0.02258301, ...,  0.01464844,
         -0.00799561, -0.01226807],
        [-0.02355957,  0.04345703, -0.03442383, ...,  0.01831055,
         -0.04418945, -0.00515747],
        [ 0.06201172, -0.00576782, -0.02099609, ...,  0.02075195,
          0.00189209,  0.01251221],
        ...,
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484],
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484],
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484]],

       [[ 0.00012589, -0.00585938,  0.02258301, ...,  0.01464844,
         -0.00799561, -0.01226807],
        [-0.02355957,  0.04345703, -0.03442383, ...,  0.01831055,
         -0.04418945, -0.00515747],
        [ 0.06201172, -0.00576782, -0.02099609, ...,  0.02075195,
          0.00189209,  0.01251221],
        ...,
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484],
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484],
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484]],

       [[ 0.00012589, -0.00585938,  0.02258301, ...,  0.01464844,
         -0.00799561, -0.01226807],
        [-0.01257324, -0.01000977, -0.00083542, ...,  0.02880859,
         -0.02294922,  0.01940918],
        [ 0.00115204, -0.01818848, -0.02246094, ..., -0.04296875,
          0.02893066,  0.03295898],
        ...,
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484],
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484],
        [ 0.03417969, -0.03198242,  0.07324219, ...,  0.02001953,
          0.04931641, -0.03271484]]], dtype=float32)>

I'm not sure how I should reshape (or add some sort of pooling layer to?) the Gemma embeddings.