埋め込みの続編。gpt-oss-20bの埋め込み空間で king - man + woman をやってみよう。
モデルは mlx-community/gpt-oss-20b-MXFP4-Q8 を使う。ただし、このモデルは推論用には微妙である(そもそもQ8なのにQ4とほとんどサイズが変わらない)。
model, tokenizer = load("mlx-community/gpt-oss-20b-MXFP4-Q8")
prompt = "queen"
token_ids = tokenizer.encode(prompt) # [153556]
tokens = tokenizer.convert_ids_to_tokens(token_ids) # ['queen']
トークン列を埋め込みベクトルに変換するための関数を探すために、Pythonに model と打ち込んでみる:
Model(
(model): GptOssMoeModel(
(embed_tokens): QuantizedEmbedding(201088, 2880, group_size=64, bits=8, mode=affine)
(norm): RMSNorm(2880, eps=1e-05)
...以下略...
これから model.model.embed_tokens だと推量する:
embeddings = model.model.embed_tokens(token_ids) queen = embeddings[0] # 2880次元のベクトル mx.linalg.norm(queen) # array(123, dtype=bfloat16)
同じことを 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([[ 0.9999998 , 0.43785974, 0.5198675 , 0.35447446, -0.11281953],
[ 0.43785974, 0.9999999 , 0.19880976, 0.4180091 , 0.6151444 ],
[ 0.5198675 , 0.19880976, 1. , 0.33824253, 0.4844749 ],
[ 0.35447446, 0.4180091 , 0.33824253, 0.99999976, 0.30324462],
[-0.11281953, 0.6151444 , 0.4844749 , 0.30324462, 1.0000004 ]],
dtype=float32)
king - man + woman は queen より king に近く、それよりさらに woman に近い。