@@ -108,27 +108,69 @@ class BGEM3Strategy(EmbeddingStrategy):
108108 def tier_name (self ) -> str :
109109 return MAIN_TIER
110110
111+ def _ensure_loaded (self ):
112+ if self ._model is None :
113+ import torch
114+ device , dtype = self ._resolve_device_and_dtype ()
115+ logger .info ("Loading %s model: %s" , self .tier_name , self .model_name )
116+ try :
117+ from FlagEmbedding import BGEM3FlagModel
118+ use_fp16 = dtype != torch .float32
119+ self ._model = BGEM3FlagModel (self .model_name , use_fp16 = use_fp16 , device = device )
120+ self ._is_flag_model = True
121+ except ImportError :
122+ logger .warning ("FlagEmbedding not installed. Falling back to SentenceTransformer. Sparse embeddings disabled." )
123+ from sentence_transformers import SentenceTransformer
124+ self ._is_flag_model = False
125+ try :
126+ self ._model = SentenceTransformer (
127+ self .model_name ,
128+ device = device ,
129+ model_kwargs = {"dtype" : dtype },
130+ processor_kwargs = {"use_fast" : True },
131+ )
132+ except TypeError :
133+ self ._model = SentenceTransformer (
134+ self .model_name ,
135+ device = device ,
136+ model_kwargs = {"dtype" : dtype },
137+ tokenizer_kwargs = {"use_fast" : True },
138+ )
139+ return self ._model
140+
111141 async def embed (self , texts : list [str ], batch_size : int ) -> list [dict [str , Any ]]:
112142 model = await asyncio .to_thread (self ._ensure_loaded )
113- # BGE-M3 supports dense and sparse (lexical) weights
114- output = await asyncio .to_thread (
115- model .encode ,
116- texts ,
117- batch_size = batch_size ,
118- return_dense = True ,
119- return_sparse = True ,
120- return_colbert_vecs = False ,
121- convert_to_numpy = True ,
122- show_progress_bar = False ,
123- )
124-
125- dense_vecs = output ["dense_vecs" ]
126- lexical_weights = output ["lexical_weights" ]
127-
128- return [
129- {"dense" : d .tolist (), "sparse" : s }
130- for d , s in zip (dense_vecs , lexical_weights , strict = True )
131- ]
143+ if getattr (self , "_is_flag_model" , False ):
144+ # BGE-M3 supports dense and sparse (lexical) weights natively
145+ output = await asyncio .to_thread (
146+ model .encode ,
147+ texts ,
148+ batch_size = batch_size ,
149+ max_length = 8192 ,
150+ return_dense = True ,
151+ return_sparse = True ,
152+ return_colbert_vecs = False ,
153+ )
154+ dense_vecs = output ["dense_vecs" ]
155+ lexical_weights = output ["lexical_weights" ]
156+
157+ return [
158+ {"dense" : d .tolist (), "sparse" : s }
159+ for d , s in zip (dense_vecs , lexical_weights , strict = True )
160+ ]
161+ else :
162+ import numpy as np
163+ embeddings = await asyncio .to_thread (
164+ model .encode ,
165+ texts ,
166+ batch_size = batch_size ,
167+ convert_to_numpy = True ,
168+ show_progress_bar = False ,
169+ normalize_embeddings = True ,
170+ )
171+ return [
172+ {"dense" : emb .tolist (), "sparse" : None } for emb in np .asarray (embeddings )
173+ ]
132174
133175
134176class BGESlimStrategy (EmbeddingStrategy ):
0 commit comments