@@ -93,13 +93,15 @@ class ModelBase:
9393 # Mistral format specifics
9494 is_mistral_format : bool = False
9595 disable_mistral_community_chat_template : bool = False
96+ sentence_transformers_dense_modules : bool = False
9697
9798 def __init__ (self , dir_model : Path , ftype : gguf .LlamaFileType , fname_out : Path , * , is_big_endian : bool = False ,
9899 use_temp_file : bool = False , eager : bool = False ,
99100 metadata_override : Path | None = None , model_name : str | None = None ,
100101 split_max_tensors : int = 0 , split_max_size : int = 0 , dry_run : bool = False ,
101102 small_first_shard : bool = False , hparams : dict [str , Any ] | None = None , remote_hf_model_id : str | None = None ,
102- disable_mistral_community_chat_template : bool = False ):
103+ disable_mistral_community_chat_template : bool = False ,
104+ sentence_transformers_dense_modules : bool = False ):
103105 if type (self ) is ModelBase or \
104106 type (self ) is TextModel or \
105107 type (self ) is MmprojModel :
@@ -114,6 +116,7 @@ def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path,
114116 self .lazy = not eager or (remote_hf_model_id is not None )
115117 self .dry_run = dry_run
116118 self .remote_hf_model_id = remote_hf_model_id
119+ self .sentence_transformers_dense_modules = sentence_transformers_dense_modules
117120 if remote_hf_model_id is not None :
118121 self .is_safetensors = True
119122
@@ -5269,6 +5272,53 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
52695272@ModelBase .register ("Gemma3TextModel" )
52705273class EmbeddingGemma (Gemma3Model ):
52715274 model_arch = gguf .MODEL_ARCH .GEMMA_EMBEDDING
5275+ module_paths = []
5276+ dense_features_dims = {}
5277+
5278+ def __init__ (self , * args , ** kwargs ):
5279+ super ().__init__ (* args , ** kwargs )
5280+ if self .sentence_transformers_dense_modules :
5281+ # read modules.json to determine if model has Dense layers
5282+ modules_file = self .dir_model / "modules.json"
5283+ if modules_file .is_file ():
5284+ with open (modules_file , encoding = "utf-8" ) as modules_json_file :
5285+ mods = json .load (modules_json_file )
5286+ for mod in mods :
5287+ if mod ["type" ] == "sentence_transformers.models.Dense" :
5288+ mod_path = mod ["path" ]
5289+ # check if model.safetensors file for Dense layer exists
5290+ model_tensors_file = self .dir_model / mod_path / "model.safetensors"
5291+ if model_tensors_file .is_file ():
5292+ self .module_paths .append (mod_path )
5293+ # read config.json of the Dense layer to get in/out features
5294+ mod_conf_file = self .dir_model / mod_path / "config.json"
5295+ if mod_conf_file .is_file ():
5296+ with open (mod_conf_file , encoding = "utf-8" ) as mod_conf_json_file :
5297+ mod_conf = json .load (mod_conf_json_file )
5298+ # hparams dense_2_feat_out and dense_3_feat_in are required when loading model's dense weights
5299+ prefix = self ._get_dense_prefix (mod_path )
5300+ if mod_conf ["in_features" ] is not None and mod_conf ["out_features" ] is not None :
5301+ self .dense_features_dims [prefix ] = (mod_conf ["in_features" ], mod_conf ["out_features" ])
5302+
5303+ def generate_extra_tensors (self ) -> Iterable [tuple [str , Tensor ]]:
5304+ from safetensors .torch import load_file
5305+ module_paths = list (self .module_paths )
5306+ for i , module_path in enumerate (module_paths ):
5307+ tensors_file = self .dir_model / module_path / "model.safetensors"
5308+ local_tensors = load_file (tensors_file )
5309+ tensor_name = self ._get_dense_prefix (module_path )
5310+ for name , local_tensor in local_tensors .items ():
5311+ if not name .endswith (".weight" ):
5312+ continue
5313+ orig_name = name .replace ("linear" , tensor_name )
5314+ name = self .map_tensor_name (orig_name )
5315+ yield name , local_tensor .clone ()
5316+
5317+ @staticmethod
5318+ def _get_dense_prefix (module_path ) -> str :
5319+ """Get the tensor name prefix for the Dense layer from module path."""
5320+ tensor_name = "dense_2" if module_path == "2_Dense" else "dense_3"
5321+ return tensor_name
52725322
52735323 def set_gguf_parameters (self ):
52745324 super ().set_gguf_parameters ()
@@ -5285,6 +5335,10 @@ def set_gguf_parameters(self):
52855335 logger .info (f"Using original sliding_window from config: { orig_sliding_window }
52865336 f"instead of { self .hparams ['sliding_window' ]} )
52875337 self .gguf_writer .add_sliding_window (orig_sliding_window )
5338+ if self .sentence_transformers_dense_modules :
5339+ for dense , dims in self .dense_features_dims .items ():
5340+ logger .info (f"Setting dense layer { dense } { dims } )
5341+ self .gguf_writer .add_dense_features_dims (dense , dims [0 ], dims [1 ])
52885342
52895343 self ._try_set_pooling_type ()
52905344
@@ -9335,6 +9389,13 @@ def parse_args() -> argparse.Namespace:
93359389 )
93369390 )
93379391
9392+ parser .add_argument (
9393+ "--sentence-transformers-dense-modules" , action = "store_true" ,
9394+ help = ("Whether to include sentence-transformers dense modules."
9395+ "It can be used for sentence-transformers models, like google/embeddinggemma-300m"
9396+ "Default these modules are not included." )
9397+ )
9398+
93389399 args = parser .parse_args ()
93399400 if not args .print_supported_models and args .model is None :
93409401 parser .error ("the following arguments are required: model" )
@@ -9397,9 +9458,13 @@ def main() -> None:
93979458 if args .remote :
93989459 hf_repo_id = args .model
93999460 from huggingface_hub import snapshot_download
9461+ allowed_patterns = ["LICENSE" , "*.json" , "*.md" , "*.txt" , "tokenizer.model" ]
9462+ if args .sentence_transformers_dense_modules :
9463+ # include sentence-transformers dense modules safetensors files
9464+ allowed_patterns .append ("*.safetensors" )
94009465 local_dir = snapshot_download (
94019466 repo_id = hf_repo_id ,
9402- allow_patterns = [ "LICENSE" , "*.json" , "*.md" , "*.txt" , "tokenizer.model" ] )
9467+ allow_patterns = allowed_patterns )
94039468 dir_model = Path (local_dir )
94049469 logger .info (f"Downloaded config and tokenizer to { local_dir } )
94059470 else :
@@ -9467,7 +9532,8 @@ def main() -> None:
94679532 split_max_tensors = args .split_max_tensors ,
94689533 split_max_size = split_str_to_n_bytes (args .split_max_size ), dry_run = args .dry_run ,
94699534 small_first_shard = args .no_tensor_first_split ,
9470- remote_hf_model_id = hf_repo_id , disable_mistral_community_chat_template = disable_mistral_community_chat_template
9535+ remote_hf_model_id = hf_repo_id , disable_mistral_community_chat_template = disable_mistral_community_chat_template ,
9536+ sentence_transformers_dense_modules = args .sentence_transformers_dense_modules
94719537 )
94729538
94739539 if args .vocab_only :
0 commit comments