embeddinggemma-2を使って画像をembeddingしてみる

$ uv init --python 3.14 --bare
$ uv add sentence-transformers transformers

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("google/embeddinggemma-2")

query = "What causes the northern lights?"
document = "The northern lights are caused by charged particles from the sun."

query_emb = model.encode(query, prompt_name="SearchQuery")
doc_emb = model.encode(document, prompt_name="Document")
print(model.similarity(query_emb, doc_emb))

$ uv run quick_start.py
ValueError: The checkpoint you are trying to load has model type `embedding_gemma2` but Transformers does not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date.

You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is very new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git`

$ uv add git+https://github.com/huggingface/transformers.git
Resolved 62 packages in 1.94s
Uninstalled 1 package in 181ms
Installed 1 package in 82ms
 - transformers==5.18.0
 + transformers==5.19.0.dev0 (from git+https://github.com/huggingface/transformers.git@e598fbad926d80bc356c0c4d96532030e3dfdb62)

$ uv run quick_start.py
ImportError: EmbeddingGemma2Processor requires the PIL library but it was not found in your environment. You can install it with pip:
`pip install pillow`. Please note that you may need to restart your runtime after installation.

To install the required dependencies, run:
pip install -U "sentence-transformers[image]"

$ uv add sentence-transformers[image]
Resolved 64 packages in 1.61s
Installed 2 packages in 17ms
 + pillow==12.3.0
 + torchvision==0.29.1

$ uv run quick_start.py
tensor([[0.8713]])

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("google/embeddinggemma-2")

embedded = model.encode({"image": "images/01_a_chihuahua.png"})
print(embedded)
print(type(embedded))
print(embedded.shape)
print(embedded.dtype)

$ uv run image_embedding.py
[ 4.71191406e-02  2.94494629e-03 -3.63769531e-02 -1.47247314e-03
 -2.29492188e-02  1.69677734e-02  5.07354736e-04 -1.73339844e-02
 ... (中略) ...
 -2.34603882e-04 -2.18505859e-02  1.98974609e-02  2.09045410e-03]
<class 'numpy.ndarray'>
(768,)
float32

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("google/embeddinggemma-2")

query = "search dog image"
documents = [
    {"image": "images/01_a_chihuahua.png"},
    {"image": "images/02_three_miniature_dachshunds.png"},
    {"image": "images/03_a_cat.png"},
]

emb_query = model.encode_query(query)
emb_document = model.encode_document(documents)

similarities = model.similarity(emb_query, emb_document)
print(similarities)

$ uv run calculate_similarities.py
tensor([[0.7282, 0.7152, 0.6756]])

$ aws s3vectors create-vector-bucket \
      --vector-bucket-name test-bucket
{
    "vectorBucketArn": "arn:aws:s3vectors:ap-northeast-1:000000000000:bucket/test-bucket"
}

$ aws s3vectors create-index \
      --vector-bucket-name test-bucket \
      --index-name test-index \
      --data-type float32 \
      --dimension 768 \
      --distance-metric cosine
{
    "indexArn": "arn:aws:s3vectors:ap-northeast-1:000000000000:bucket/test-bucket/index/test-index"
}

import boto3
from sentence_transformers import SentenceTransformer

client = boto3.client("s3vectors")

model = SentenceTransformer("google/embeddinggemma-2")

documents = [
    {"image": "images/01_a_chihuahua.png"},
    {"image": "images/02_three_miniature_dachshunds.png"},
    {"image": "images/03_a_cat.png"},
]

emb_documents = model.encode_document(documents)

vectors = [
    {"key": doc["image"], "data": {"float32": emb.tolist()}}
    for doc, emb in zip(documents, emb_documents)
]

client.put_vectors(
    vectorBucketName="test-bucket",
    indexName="test-index",
    vectors=vectors,
)

$ uv add boto3
$ uv run put_vectors.py

import boto3
from sentence_transformers import SentenceTransformer

client = boto3.client("s3vectors")

model = SentenceTransformer("google/embeddinggemma-2")

query = "search dog image"

emb = model.encode_query(query)

resp = client.query_vectors(
    vectorBucketName="test-bucket",
    indexName="test-index",
    topK=3,
    queryVector={"float32": emb.tolist()},
    returnDistance=True,
)

print(resp["vectors"])

$ uv run query_vectors.py
[{'distance': 0.27181023359298706, 'key': 'images/01_a_chihuahua.png'},
 {'distance': 0.28478217124938965,
  'key': 'images/02_three_miniature_dachshunds.png'},
 {'distance': 0.3243723511695862, 'key': 'images/03_a_cat.png'}]