Coverage for src/local_deep_research/embeddings/sentence_transformer_models.py: 100%

1"""Curated Sentence Transformer model identities. 

2 

3The model key is the stable value stored in user settings and collection 

4metadata. Only newly introduced keys may pin a different Hugging Face artifact: 

5legacy keys must keep their historic floating-revision semantics because 

6collection metadata does not yet persist an artifact revision. 

7 

8Pinned ``revision`` values are immutable Hugging Face commit hashes 

9that identify the exact model artifact bytes fetched at download time. 

10""" 

11 

12from dataclasses import dataclass 

13from types import MappingProxyType 

14from typing import Mapping 

15 

16 

17from ..constants import ( 

18 DEFAULT_LOCAL_SEARCH_MODEL as DEFAULT_SENTENCE_TRANSFORMER_MODEL, 

19) 

20 

21 

22@dataclass(frozen=True) 

23class SentenceTransformerModelSpec: 

24 """Artifact policy and display metadata for a curated model.""" 

25 

26 repository: str 

27 revision: str | None 

28 dimensions: int 

29 description: str 

30 max_seq_length: int 

31 # Require behavior-changing metadata before an offline load. Missing 

32 # Transformer or tokenizer configuration can silently change truncation 

33 # or tokenizer selection without raising a missing-weight error. 

34 required_cache_files: tuple[str, ...] = ( 

35 "config.json", 

36 "modules.json", 

37 "1_Pooling/config.json", 

38 "tokenizer_config.json", 

39 "sentence_bert_config.json", 

40 ) 

41 # Tokenizer vocabulary, as groups of which at least one file must be 

42 # cached. Without it the offline constructor does not fail: it builds a 

43 # placeholder tokenizer from the special tokens alone and embeds every 

44 # text as the same few ids, which the missing-weight repair never sees. 

45 # Legacy WordPiece/MPNet repositories publish both the fast 

46 # ``tokenizer.json`` and the slow ``vocab.txt``; either one alone rebuilds 

47 # the same tokenizer, so a cache holding only one of them still loads. 

48 required_cache_file_alternatives: tuple[tuple[str, ...], ...] = ( 

49 ("tokenizer.json", "vocab.txt"), 

50 ) 

51 

52 def display_metadata(self) -> dict[str, int | str]: 

53 """Return the legacy metadata shape consumed by provider APIs.""" 

54 

55 return { 

56 "dimensions": self.dimensions, 

57 "description": self.description, 

58 "max_seq_length": self.max_seq_length, 

59 } 

60 

61 

62SENTENCE_TRANSFORMER_MODELS: Mapping[str, SentenceTransformerModelSpec] = ( 

63 MappingProxyType( 

64 { 

65 DEFAULT_SENTENCE_TRANSFORMER_MODEL: SentenceTransformerModelSpec( 

66 repository=DEFAULT_SENTENCE_TRANSFORMER_MODEL, 

67 revision="e7f32e3c00f91d699e8c43b53106206bcc72bb22", # DevSkim: ignore DS173237 - pinned public Hugging Face commit SHA, not a secret 

68 dimensions=768, 

69 description=( 

70 "Recommended default with higher-quality English and " 

71 "long-context document retrieval." 

72 ), 

73 max_seq_length=8192, 

74 # This pinned artifact has no sentence_bert_config.json; 

75 # its token limit is inferred from model and tokenizer config. 

76 required_cache_files=( 

77 "config.json", 

78 "modules.json", 

79 "1_Pooling/config.json", 

80 "tokenizer_config.json", 

81 ), 

82 # ModernBERT ships a BPE ``tokenizer.json`` and no vocab.txt. 

83 required_cache_file_alternatives=(("tokenizer.json",),), 

84 ), 

85 "all-MiniLM-L6-v2": SentenceTransformerModelSpec( 

86 repository="sentence-transformers/all-MiniLM-L6-v2", 

87 revision=None, 

88 dimensions=384, 

89 description="Fast, lightweight model. Good for general use.", 

90 max_seq_length=256, 

91 ), 

92 "all-mpnet-base-v2": SentenceTransformerModelSpec( 

93 repository="sentence-transformers/all-mpnet-base-v2", 

94 revision=None, 

95 dimensions=768, 

96 description="Higher quality, slower. Best accuracy.", 

97 max_seq_length=384, 

98 ), 

99 "multi-qa-MiniLM-L6-cos-v1": SentenceTransformerModelSpec( 

100 repository=("sentence-transformers/multi-qa-MiniLM-L6-cos-v1"), 

101 revision=None, 

102 dimensions=384, 

103 description="Optimized for question-answering tasks.", 

104 max_seq_length=512, 

105 ), 

106 "paraphrase-multilingual-MiniLM-L12-v2": ( 

107 SentenceTransformerModelSpec( 

108 repository=( 

109 "sentence-transformers/" 

110 "paraphrase-multilingual-MiniLM-L12-v2" 

111 ), 

112 revision=None, 

113 dimensions=384, 

114 description="Supports multiple languages.", 

115 max_seq_length=128, 

116 # An XLM-R SentencePiece tokenizer, not WordPiece: there 

117 # is no vocab.txt fallback, so require ``tokenizer.json``. 

118 required_cache_file_alternatives=(("tokenizer.json",),), 

119 ) 

120 ), 

121 } 

122 ) 

123) 

124 

125 

126def get_sentence_transformer_model_spec( 

127 model_name: str, 

128) -> SentenceTransformerModelSpec | None: 

129 """Resolve either a stable catalog key or its full repository name.""" 

130 

131 spec = SENTENCE_TRANSFORMER_MODELS.get(model_name) 

132 if spec is not None: 

133 return spec 

134 

135 return next( 

136 ( 

137 candidate 

138 for candidate in SENTENCE_TRANSFORMER_MODELS.values() 

139 if candidate.repository == model_name 

140 ), 

141 None, 

142 )