Coverage for src/lilbee/retrieval/embedding_profiles.py: 100%

1"""Per-family embedding instruction profiles for asymmetric query/document embedding.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6 

7 

8@dataclass(frozen=True) 

9class EmbeddingProfile: 

10 """Instruction prefixes an embedder applies to a query vs a document.""" 

11 

12 query_instruction: str = "" 

13 doc_prefix: str = "" 

14 # META_SCHEMA_VERSION that first stamped this doc_prefix at ingest; an 

15 # older store warns to rebuild rather than silently mixing prefixed 

16 # queries with unprefixed documents. 

17 doc_prefix_since: int = 1 

18 

19 

20# Instruction-tuned embedders (Qwen3-Embedding, *-instruct: e5/gte/...) share the 

21# Instruct/Query format with no document prefix. Base e5 uses query:/passage:, 

22# nomic-embed requires search_query:/search_document: (trained with them), bge 

23# v1/v1.5 wants a query-side instruction only. Everything else (bge-m3, 

24# gte-large, ...) stays symmetric: correct-but-symmetric, never silently wrong. 

25# Order is specific-first. 

26_INSTRUCT = EmbeddingProfile( 

27 query_instruction=( 

28 "Instruct: Given a web search query, retrieve relevant passages that " 

29 "answer the query\nQuery: " 

30 ), 

31) 

32_E5 = EmbeddingProfile(query_instruction="query: ", doc_prefix="passage: ") 

33_NOMIC = EmbeddingProfile( 

34 query_instruction="search_query: ", doc_prefix="search_document: ", doc_prefix_since=2 

35) 

36_BGE_V1 = EmbeddingProfile( 

37 query_instruction="Represent this sentence for searching relevant passages: " 

38) 

39_SYMMETRIC = EmbeddingProfile() 

40_FAMILY_PROFILES: tuple[tuple[str, EmbeddingProfile], ...] = ( 

41 ("qwen3-embedding", _INSTRUCT), 

42 # "instructor" contains "instruct" but is a different dialect: the Instructor 

43 # family (hkunlp/instructor-*, a t5encoder the engine can load) prefixes a 

44 # "Represent the ... for retrieval:" instruction on the document as well as 

45 # the query. Handing it the Instruct/Query query prefix with no doc prefix 

46 # would be the asymmetric-in-the-wrong-dialect case this module promises not 

47 # to produce, so it takes the symmetric fallback until the real prefixes are 

48 # wired. 

49 ("instructor", _SYMMETRIC), 

50 ("instruct", _INSTRUCT), # any instruction-tuned embedder: Instruct/Query, no doc prefix 

51 ("multilingual-e5", _E5), 

52 ("e5-", _E5), 

53 ("nomic-embed", _NOMIC), 

54 ("bge-m3", _SYMMETRIC), # genuinely symmetric; must not fall through to bge v1 

55 ("bge-", _BGE_V1), 

56) 

57 

58 

59def resolve_embedding_profile(model_ref: str) -> EmbeddingProfile: 

60 """Profile for *model_ref*, or a symmetric (empty) profile when unrecognized.""" 

61 ref = model_ref.lower() 

62 for needle, profile in _FAMILY_PROFILES: 

63 if needle in ref: 

64 return profile 

65 return EmbeddingProfile()