+ tests

Author: noooop

tests/models/language/pooling_mteb_test/mteb_utils.py

@@ -212,10 +212,13 @@ def mteb_test_embed_models(hf_runner,
         vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
         head_dtype = model_config.head_dtype
 
-        # Test embed_dims, isnan and whether to use normalize
+        # Test embedding_size, isnan and whether to use normalize
         vllm_outputs = vllm_model.embed(example_prompts,
                                         truncate_prompt_tokens=-1)
-        assert not torch.any(torch.isnan(torch.tensor(vllm_outputs)))
+        outputs_tensor = torch.tensor(vllm_outputs)
+        assert not torch.any(torch.isnan(outputs_tensor))
+        embedding_size = model_config.embedding_size
+        assert torch.tensor(vllm_outputs).shape[-1] == embedding_size
 
     # Accelerate mteb test by setting
     # SentenceTransformers mteb score to a constant
@@ -233,7 +236,7 @@ def mteb_test_embed_models(hf_runner,
             st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
             st_dtype = next(hf_model.model.parameters()).dtype
 
-            # Test embed_dims and whether to use normalize
+            # Check embeddings close to hf outputs
             hf_outputs = hf_model.encode(example_prompts)
             check_embeddings_close(
                 embeddings_0_lst=hf_outputs,

vllm/transformers_utils/config.py

@@ -1003,12 +1003,11 @@ def try_get_tokenizer_config(
 
 @cache
 def try_get_dense_modules(
-    pretrained_model_name_or_path: Union[str, os.PathLike],
+    model: Union[str, Path],
     revision: Optional[str] = None,
 ) -> Optional[list[dict[str, Any]]]:
     try:
-        modules = get_hf_file_to_dict("modules.json",
-                                      pretrained_model_name_or_path, revision)
+        modules = get_hf_file_to_dict("modules.json", model, revision)
         if not modules:
             return None
 
@@ -1027,9 +1026,7 @@ def try_get_dense_modules(
             folder = module.get("path", "")
 
             config_path = f"{folder}/config.json" if folder else "config.json"
-            layer_config = get_hf_file_to_dict(config_path,
-                                               pretrained_model_name_or_path,
-                                               revision)
+            layer_config = get_hf_file_to_dict(config_path, model, revision)
             if not layer_config:
                 continue
             layer_config["folder"] = folder