fix: route MIEB/MAEB task-types to correct LoRA adapter for jina v5 omni

[florian-hoenicke]

Summary

jinaai/jina-embeddings-v5-omni-{nano,small} ship four LoRA adapters (retrieval / clustering / text-matching / classification) and switch between them via the task= arg passed to self.model.encode(...) inside JinaV5OmniWrapper.encode. The initial omni registration in #4604 copied the text wrapper's model_prompts dict, which only had text-task-type keys (Retrieval, Clustering, STS, …). Every MIEB / MAEB task therefore fell through to a hard-coded task="retrieval", encoding clustering, VisualSTS, audio-clustering and audio-pair tasks with the wrong adapter.

Fix

Three changes in this PR:

1. Adapter routing via simplified-task fallback + per-type overrides. JinaV5OmniWrapper.encode now resolves the adapter in two steps: (a) per-MTEB-type override from loader_kwargs.model_prompts, (b) otherwise fall back to task_metadata.simplified_task_type mapped via:

   _SIMPLIFIED_TO_JINA_TASK = {
       "retrieval": "retrieval",
       "clustering": "clustering",
       "classification": "classification",
       "semantic-similarity": "text-matching",
       "pair-classification": "text-matching",
   }

   Each omni ModelMeta keeps only the 10 overrides where empirical routing diverges from the simplified default. New MTEB task types route automatically.

2. Drop retrieval-only "Query: " / "Document: " prefix for non-retrieval adapters. Only the retrieval LoRA was trained with the prefix; injecting it on the other variants collapsed scores (VisualSTS(eng) 0.47→0.83, AROCocoOrder 0.37→0.52, NMSQAPair 0.75→0.83 on small).

3. Drop the nano torch_dtype=float32 override. Forced fp32 on nano (small ran bf16) broke OCR/document tasks (HatefulMemesT2I 0.06→0.77, VidoreArxivQA 0.19→0.75).

Final per-type routing (after empirical validation: 145 tasks × 4 variants × 2 models)

LoRA adapter	MTEB task types routed here	Source
retrieval	Any2AnyRetrieval, Any2AnyMultilingualRetrieval, AudioRetrieval, DocumentUnderstanding, VisionCentricQA, AudioReranking	simplified-fallback (retrieval)
retrieval	ImageClassification, ZeroShotClassification, AudioClassification, AudioZeroshotClassification, AudioMultilabelClassification, VideoClassification, VideoZeroshotClassification	override (classification LoRA empty for non-text modalities)
retrieval	AudioPairClassification	override (+0.06-0.08 vs text-matching on NMSQAPair / CREMADPair / VoxPopuliAccentPair)
clustering	AudioClustering	simplified-fallback (clustering)
clustering	Compositionality	override (+0.01 to +0.27 vs text-matching on AROFlickrOrder / AROCocoOrder / SugarCrepe / Winoground / ImageCoDe; AROVisualRelation / AROVisualAttribution insensitive)
text-matching	VisualSTS(eng), VisualSTS(multi)	simplified-fallback (semantic-similarity)
text-matching	ImageClustering	override (best on 8 of 10 small+nano combinations across CIFAR10/100, TinyImageNet, ImageNet10/Dog15 Clustering; aggregate mean 0.771 vs retrieval 0.763 vs clustering 0.756)
classification	(none — every *Classification MIEB/MAEB type routes to retrieval)	—

Tests

tests/test_models/test_jina_v5_omni_wrapper.py (19 tests):

• prefix gating: retrieval keeps "Query: " / "Document: ", other variants get prompt=""
• nano loader_kwargs no longer pins torch_dtype=float32
• simplified-fallback paths (Any2AnyRetrieval, AudioClustering, VisualSTS(eng), AudioPairClassification)
• override paths (Image/Audio/Zeroshot/VideoClassification, Compositionality, AudioPairClassification, ImageClustering)
• drift guard: _SIMPLIFIED_TO_JINA_TASK keys must equal mteb's SimplifiedTaskType set
• minimality guard: every override must differ from the simplified default

🤖 Generated with Claude Code

[florian-hoenicke]

Task-type → variant assignment rationale

For traceability, here is the complete mapping we use, which mirrors Jina's internal frontier-dashboard routing (harness/src/core/evaluator.py: VARIANT_TASK_TO_MTEB_TYPES and VARIANT_TASK_TO_DATASETS["maeb"]):

Image (MIEB lite / eng / Multilingual)

MTEB task type	LoRA adapter	Example tasks
Any2AnyRetrieval	retrieval	MSCOCOT2IRetrieval, Flickr30kT2I, BLINK*, OVEN*, METI2IRetrieval, GLDv2*, Imagenet1kZeroShot, …
Any2AnyMultilingualRetrieval	retrieval	XM3600T2IRetrieval, XFlickr30kCoT2IRetrieval, WITT2IRetrieval
DocumentUnderstanding	retrieval	Vidore* (DocVQA, InfoVQA, TabFQuAD, …)
ZeroShotClassification	retrieval	Caltech101ZeroShot, MNISTZeroShot, RESISC45ZeroShot, …
ImageClassification	retrieval	Caltech101, DTD, Food101Classification, Imagenet1k, …
VisionCentricQA	retrieval	CVBenchCount, CVBenchDepth, CVBenchDistance, CVBenchRelation
ImageClustering	clustering	CIFAR10Clustering, CIFAR100Clustering, TinyImageNetClustering, ImageNet10Clustering, ImageNetDog15Clustering
Compositionality	clustering	AROCocoOrder, AROFlickrOrder, AROVisualAttribution, AROVisualRelation, SugarCrepe, Winoground, ImageCoDe
VisualSTS(eng)	text-matching	STS12-16VisualSTS, VisualSTS17Eng, VisualSTS-b-Eng
VisualSTS(multi)	text-matching	VisualSTS17Multilingual, VisualSTS-b-Multilingual

Audio (MAEB beta)

MTEB task type	LoRA adapter	Example tasks
AudioRetrieval	retrieval	ClothoT2A, FleursT2A, GigaSpeechT2A, SpokenSQuADT2A, JamAlt*, MACST2A, UrbanSound8KT2A, …
AudioClassification	retrieval	BeijingOpera, BirdCLEF, CREMA_D, FSD2019Kaggle, GTZANGenre, IEMOCAPGender, MridinghamTonic, VoxCelebSA
AudioZeroshotClassification	retrieval	RavdessZeroshot, SpeechCommandsZeroshotv0.02, CommonLanguageAgeDetection, VoxPopuliLanguageID
AudioReranking	retrieval	GTZANAudioReranking
AudioMultilabelClassification	retrieval	SIBFLEURS
AudioClustering	clustering	CREMA_DClustering, VehicleSoundClustering, VoxPopuliGenderClustering
AudioPairClassification	text-matching	CREMADPairClassification, NMSQAPairClassification, VoxPopuliAccentPairClassification

Why no classification adapter

Every task whose MTEB type contains "Classification" (image or audio, zeroshot or supervised) is routed to the retrieval adapter. That matches our internal training/eval pipeline where the classification LoRA's task list is empty and classification eval is performed via the retrieval adapter's contrastive embeddings.

Verification

After this patch, the wrapper logs Using prompt_name=ImageClustering for task=CIFAR10Clustering prompt_type=... (or the equivalent for VisualSTS / AudioClustering / AudioPairClassification) instead of No combination of task name and prompt type was found. We're now re-running affected tasks (12 image clustering + 9 image VisualSTS + 3 audio clustering + 3 audio pair-classification × 2 models = 54 tasks) and will update embeddings-benchmark/results with the corrected JSONs once they land.

🤖 Generated with Claude Code

[Samoed]

I think would be better to change in wrapper task.metadata.simplified_task_type

[florian-hoenicke]

Added two follow-up fixes after comparing our MTEB outputs against the model author's evaluation harness:

1. Retrieval-only prompt prefix. JinaV5OmniWrapper.encode was unconditionally prepending "Query: " / "Document: " for every task type. Only the retrieval LoRA adapter was trained with that prefix; clustering, text-matching, and classification variants were trained without it. The wrapper now only adds the prefix when task == "retrieval". This matches the upstream training-time convention (instructions={"query":"","document":""} for non-retrieval variants). Observed effect on small without the gating: VisualSTS(eng) 0.47 vs 0.83, AROCocoOrder 0.37 vs 0.52, NMSQAPairClassification 0.75 vs 0.83.

2. nano torch_dtype=float32 override. The nano ModelMeta was forcing fp32, while small ran in bf16. The forced upcast caused large divergence on OCR / document retrieval tasks: HatefulMemesT2IRetrieval 0.06 vs 0.77, VidoreArxivQARetrieval 0.19 vs 0.75. Dropping the override makes nano load with the same dtype handling as small.

tests/test_models/test_jina_v5_omni_wrapper.py covers both: stubs the underlying SentenceTransformer to capture the task and prompt kwargs the wrapper passes, and asserts (a) retrieval keeps "Query: "/"Document: ", (b) clustering / text-matching / classification get prompt="", (c) nano loader_kwargs no longer pins float32. Verified bidirectional: 5 of the 7 tests fail on the unfixed wrapper.

[florian-hoenicke]

@Samoed — good call, just pushed ee555e9d to do exactly that.

The wrapper now resolves the Jina LoRA adapter in two steps:

1. Look up the concrete MTEB type in model_prompts (per-type override).
2. If not present, fall back to task_metadata.simplified_task_type, mapped to a Jina variant via:

   _SIMPLIFIED_TO_JINA_TASK = {
       "retrieval": "retrieval",
       "clustering": "clustering",
       "classification": "classification",
       "semantic-similarity": "text-matching",
       "pair-classification": "text-matching",
   }

Each omni ModelMeta model_prompts shrinks from 27 entries to 8 — only the genuine overrides where our harness routing diverges from simplified_task_type:

MTEB type	simplified_task_type	Jina LoRA (override)	reason
ImageClassification, ZeroShotClassification, AudioClassification, AudioZeroshotClassification, AudioMultilabelClassification, VideoClassification, VideoZeroshotClassification	classification	retrieval	the classification LoRA's task list is empty for non-text modalities — eval runs through the retrieval adapter's contrastive embeddings
Compositionality	pair-classification	clustering	harness routes it through the clustering adapter

Text routings (Classification → classification, PairClassification → text-matching, etc.) are now handled automatically by the simplified fallback; same for Any2AnyRetrieval, ImageClustering, VisualSTS(eng/multi), AudioPairClassification, etc.

Added two guard tests:

• test_simplified_to_jina_task_covers_all_simplified_types — fails on drift if mteb adds a new SimplifiedTaskType.
• test_omni_overrides_are_minimal — fails if any override actually matches the simplified default (forces the dict to stay minimal).

17/17 tests passing; the 4 newly-added fallback tests confirm new MTEB task types route correctly without touching the wrapper.

[Samoed]

I don't think you should delete old prompts

[florian-hoenicke]

@Samoed — restored the 10 text task-type entries in model_prompts on both omni-small and omni-nano in 488e4e1. Verified each entry resolves to the same Jina LoRA adapter via either the explicit dict or the _SIMPLIFIED_TO_JINA_TASK fallback, so this is a no-op behaviorally but acts as a safety net against upstream drift in _TASKTYPE2SIMPLIFIEDTASKTYPE. The redundant-entry guard test is removed; 18/18 remaining tests pass.

Is this what you had in mind?

[Samoed]

I'm not sure if we want to keep these tests

[florian-hoenicke]

These were added as regression guards for two real bugs we hit:

• the prefix-injection bug (which collapsed VisualSTS17Eng from 0.87 → 0.47 silently)
• the simplified-task-type fallback added in ee555e9d per your earlier suggestion

The core 6 tests verify the prefix gate per variant; the rest cover the simplified-fallback path, the override path, and two drift guards (one that the override dict only contains genuine divergences from the simplified default, one that _SIMPLIFIED_TO_JINA_TASK covers every SimplifiedTaskType mteb defines).

I can trim to a minimal 3-test set (retrieval keeps prefix / non-retrieval drops prefix / drift guard) if you prefer, or drop the whole file. Let me know which way you'd like to go.

[Samoed]

I simplified your tests. Can you look into them? Also @KennethEnevoldsen What do you think?

[florian-hoenicke]

The simplification looks great — much cleaner with MockSentenceTransformer / MockRetrievalTask and the parametrized form. All 17 cases pass locally. Happy with this; the per-call logger.info is a nice touch too.

[KennethEnevoldsen]

hmm generally we haven't yet maintained tests for models specific implementations. These are luckily fast.

I would be fine with adding these as a way to experiment more with covering central models with (fast) tests. However I could also see the argument for the reverse. Leaning accept and merge.

[KennethEnevoldsen]

Looks good - SimplifiedType was mostly intended as a documentation variable, but the use here doesn't seem problematic to me.

[KennethEnevoldsen]

hmm generally we haven't yet maintained tests for models specific implementations. These are luckily fast.

I would be fine with adding these as a way to experiment more with covering central models with (fast) tests. However I could also see the argument for the reverse. Leaning accept and merge.

[Samoed]

SimplifiedType was mostly intended as a documentation variable

FYI it also used in Bidir models