埋め込みの続編。Gemma 3の埋め込み空間で king - man + woman をやってみよう。
MacなのでMLXを使うことにする。モデルとトークナイザを取り出すのは簡単。例えば queen という語についてやってみよう。
model, tokenizer = load("mlx-community/gemma-3-27b-it-8bit")
prompt = "queen"
token_ids = tokenizer.encode(prompt) # [2, 91024]
tokens = tokenizer.convert_ids_to_tokens(token_ids) # ['<bos>', 'queen']
トークン列を埋め込みベクトルに変換するための関数を探すために、Pythonに model と打ち込んでみる:
Model(
(language_model): Model(
(model): Gemma3Model(
(embed_tokens): QuantizedEmbedding(262208, 5376, group_size=64, bits=8, mode=affine)
(layers.0): TransformerBlock(
...以下略...
これから model.language_model.model.embed_tokens だと推量する:
embeddings = model.language_model.model.embed_tokens(token_ids)
queen = embeddings[1] # 5376次元のベクトル
mx.linalg.norm(queen) # array(1.01562, dtype=bfloat16)
# mx.save("queen.npy", queen) # 念のため保存したいとき
同じことを man、woman、king についても行う。
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
X = np.vstack([man.astype(mx.float32),
woman.astype(mx.float32),
king.astype(mx.float32),
queen.astype(mx.float32),
(king - man + woman).astype(mx.float32)])
cosine_similarity(X)
結果は次の通り:
array([[ 1.000002 , 0.27174592, 0.21292192, 0.10807453, -0.35068962],
[ 0.27174592, 1.0000023 , 0.10984613, 0.18116248, 0.55315906],
[ 0.21292192, 0.10984613, 1.0000019 , 0.32620484, 0.5962215 ],
[ 0.10807453, 0.18116248, 0.32620484, 1.0000018 , 0.26477924],
[-0.35068962, 0.55315906, 0.5962215 , 0.26477924, 1.000003 ]],
dtype=float32)
やはり king - man + woman は queen より king に近い。