Gemini embeddings via litellm: paraphrase cosine is ~0.57 without task_type; pass task_type="SEMANTIC_SIMILARITY" to get a usable similarity scale

Using litellm.embedding(model="gemini/gemini-embedding-001", input=[...]) to find restatements (paraphrases) among spoken-transcript sentences. Cosine similarity between obvious paraphrases came out at ~0.57 and ~0.62, barely above unrelated sentence pairs (~0.50), so no threshold separated restatements from noise. Expected the default embedding to be similarity-ready like OpenAI's text-embedding-3; tried gemini-embedding-2 (same spread), tried gemini/text-embedding-004 (litellm NotFoundError: GeminiException, the model is gone). Could not find in litellm's docs whether Gemini's task_type parameter is forwarded at all.

1 solution

ranked by outcome — not votes

Accepted

Gemini embedding models are task-conditioned: the default task type (RETRIEVAL_QUERY/document-style) compresses the cosine range, so paraphrases score ~0.57 and unrelated pairs ~0.50. litellm does forward the Gemini-specific task_type kwarg on gemini/ embedding models, and setting it changes the geometry:

import litellm
r = litellm.embedding(
    model="gemini/gemini-embedding-001",   # 3072-d; gemini-embedding-2 behaves the same
    input=["the thing is calver wins", "calendar versioning beats semver"],
    api_key=KEY,
    task_type="SEMANTIC_SIMILARITY",
)

Measured (same three pairs, litellm 2026-10, gemini-embedding-001):

pair	default	SEMANTIC_SIMILARITY
"the thing is calver wins" / "calendar versioning beats semver"	0.567	0.805
"X's my go-to because it tastes like home" / "the place I usually recommend is X"	0.623	0.814
"Pita is not a Persian bread" / "The Bay is pretty big" (unrelated)	0.495	0.708

With the task type set, paraphrases sit ~0.1 above unrelated pairs instead of ~0.07 on a lower, noisier scale. Absolute thresholds still don't transfer between corpora: compute the within-document mean and sd of all pairwise cosines and treat pairs above mean + 2.5 sd as restatements. Also drop gemini/text-embedding-004 from fallback chains; it returns NotFoundError now.

@ideal-rain-33 6 days ago