Gemini API Embedding で詰まった3か所 — 壁紙アプリのカテゴリ自動分類で踏んだエラーと対処

from google import genai
from google.genai import types
 
client = genai.Client(api_key="YOUR_GEMINI_API_KEY")
 
# インデックスと一緒に固定する2つ。ここを触るときは再生成を覚悟する
EMBED_MODEL = "gemini-embedding-001"
EMBED_DIM = 768
 
# ❌ 固定次元のモデルに別の次元を渡すと INVALID_ARGUMENT
resp = client.models.embed_content(
    model="text-embedding-004",          # 768 次元固定
    contents="美しい山の景色",
    config=types.EmbedContentConfig(output_dimensionality=3072),
)
 
# ✅ 次元を選べるモデルで、インデックス作成時と同じ値を使い回す
resp = client.models.embed_content(
    model=EMBED_MODEL,
    contents="美しい山の景色",
    config=types.EmbedContentConfig(
        output_dimensionality=EMBED_DIM,
        task_type="RETRIEVAL_DOCUMENT",
    ),
)
vector = resp.embeddings[0].values

import numpy as np
 
def l2_normalize(vec: list[float]) -> list[float]:
    """切り詰めた埋め込みを単位ベクトルに揃える"""
    arr = np.asarray(vec, dtype=np.float32)
    norm = float(np.linalg.norm(arr))
    if norm == 0.0:
        return arr.tolist()
    return (arr / norm).tolist()
 
 
# 手元で確認したいとき: 3072 のままなら 1.0 前後、切り詰めると 1.0 から外れる
raw = resp.embeddings[0].values
print(round(float(np.linalg.norm(raw)), 4))
print(round(float(np.linalg.norm(l2_normalize(raw))), 4))  # 1.0

import asyncio
import json
import random
from pathlib import Path
 
from google import genai
from google.genai import types
 
client = genai.Client(api_key="YOUR_GEMINI_API_KEY")
 
CHECKPOINT = Path("embed_checkpoint.jsonl")
MAX_CONCURRENCY = 4
 
 
def load_done_ids() -> set[str]:
    if not CHECKPOINT.exists():
        return set()
    return {
        json.loads(line)["id"]
        for line in CHECKPOINT.read_text(encoding="utf-8").splitlines()
        if line.strip()
    }
 
 
async def embed_one(item: dict, sem: asyncio.Semaphore, max_retries: int = 5) -> dict | None:
    async with sem:
        for attempt in range(max_retries):
            try:
                resp = await client.aio.models.embed_content(
                    model=EMBED_MODEL,
                    contents=item["text"],
                    config=types.EmbedContentConfig(
                        output_dimensionality=EMBED_DIM,
                        task_type="RETRIEVAL_DOCUMENT",
                    ),
                )
                vec = l2_normalize(resp.embeddings[0].values)
                return {"id": item["id"], "vector": vec}
            except Exception as e:
                msg = str(e)
                if "429" in msg or "RESOURCE_EXHAUSTED" in msg:
                    # フルジッター: 0〜2^attempt 秒の一様乱数。同時再開の波を崩す
                    await asyncio.sleep(random.uniform(0, 2 ** attempt))
                    continue
                raise
        return None
 
 
async def embed_all(items: list[dict]) -> int:
    done = load_done_ids()
    todo = [i for i in items if i["id"] not in done]
    sem = asyncio.Semaphore(MAX_CONCURRENCY)
    written = 0
 
    with CHECKPOINT.open("a", encoding="utf-8") as f:
        tasks = [embed_one(i, sem) for i in todo]
        for coro in asyncio.as_completed(tasks):
            rec = await coro
            if rec is None:
                continue
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            f.flush()  # 落ちた瞬間までの分を残す
            written += 1
 
    return written

# ❌ task_type 未指定 — 汎用のベクトル空間になる
resp = client.models.embed_content(
    model=EMBED_MODEL,
    contents="ナチュラルな壁紙を探しています",
    config=types.EmbedContentConfig(output_dimensionality=EMBED_DIM),
)
 
# ✅ 検索クエリ側は RETRIEVAL_QUERY
query = client.models.embed_content(
    model=EMBED_MODEL,
    contents="ナチュラルな壁紙を探しています",
    config=types.EmbedContentConfig(
        output_dimensionality=EMBED_DIM,
        task_type="RETRIEVAL_QUERY",
    ),
)
 
# ✅ 格納するドキュメント側は RETRIEVAL_DOCUMENT
doc = client.models.embed_content(
    model=EMBED_MODEL,
    contents="緑の葉と木漏れ日が差し込む自然の景色",
    config=types.EmbedContentConfig(
        output_dimensionality=EMBED_DIM,
        task_type="RETRIEVAL_DOCUMENT",
    ),
)

import numpy as np
 
# 手で作る正解セット。20問あれば変更の向きは十分に見える
GOLD = [
    {"query": "自然の風景", "relevant": ["w-102", "w-118", "w-233"]},
    {"query": "夜の街の明かり", "relevant": ["w-045", "w-291"]},
    # ... 20問程度まで
]
 
 
def search(query_vec: list[float], index: dict[str, list[float]], k: int = 5) -> list[str]:
    q = np.asarray(query_vec, dtype=np.float32)
    scored = [
        (doc_id, float(np.dot(q, np.asarray(vec, dtype=np.float32))))
        for doc_id, vec in index.items()
    ]
    scored.sort(key=lambda x: x[1], reverse=True)
    return [doc_id for doc_id, _ in scored[:k]]
 
 
def recall_at_k(index: dict[str, list[float]], embed_query, k: int = 5) -> float:
    scores = []
    for case in GOLD:
        ranked = search(embed_query(case["query"]), index, k)
        hit = len(set(ranked) & set(case["relevant"]))
        scores.append(hit / len(case["relevant"]))
    return sum(scores) / len(scores)

# 壁紙アプリのバックエンドで固定している設定（参考）
EMBEDDING_CONFIG = {
    "model": EMBED_MODEL,                # 変えるならインデックス再生成
    "output_dimensionality": EMBED_DIM,  # 同上
    "normalize": True,                   # 切り詰めたベクトルは正規化して格納
    "query_task_type": "RETRIEVAL_QUERY",
    "doc_task_type": "RETRIEVAL_DOCUMENT",
    "max_concurrency": 4,                # 429 が増えたら下げる
    "max_retries": 5,                    # フルジッター付き
}