使用 bash ./scripts/eval_mteb.sh 测试v2在CMTEB数据集上的效果，会比榜单上的分数低

[sanshi9523]

使用 bash ./scripts/eval_mteb.sh 测试v2在CMTEB数据集上的效果，会比榜单上的分数低
这里的脚本支持测试 tencent/KaLM-Embedding-Gemma3-12B-2511 模型吗？

[YanshekWoo]

mteb的版本已迭代多次，该项目下的旧版已不推荐使用。

tencent/KaLM-Embedding-Gemma3-12B-2511 评测发布使用的是 mteb==2.3.9 。
评测脚本示例：

from __future__ import annotations

import os
import logging
import argparse
import functools
import json
import time
from dataclasses import dataclass, field
from collections.abc import Sequence
from typing import Tuple, List, Optional, Union, Dict, Any

import numpy as np
import torch
import torch.multiprocessing as mp
import mteb
import sentence_transformers
# from sentence_transformers import SentenceTransformer, SentenceTransformerModelCardData
from transformers import HfArgumentParser, is_torch_npu_available
from mteb.encoder_interface import PromptType
from mteb.models.wrapper import Wrapper


class SmartPathFormatter(logging.Formatter):
    def format(self, record):
        # 智能显示路径：如果路径在当前目录下显示相对路径，否则显示文件名
        if hasattr(record, 'pathname'):
            cwd = os.getcwd()
            try:
                relative_path = os.path.relpath(record.pathname, cwd)
                # 如果相对路径不以..开头，说明在当前目录或子目录下
                if not relative_path.startswith('..'):
                    record.smartpath = relative_path
                else:
                    # 否则只显示文件名
                    record.smartpath = os.path.basename(record.pathname)
            except ValueError:
                record.smartpath = record.pathname
        else:
            record.smartpath = record.filename
            
        return super().format(record)

def setup_logging():
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    
    # 使用智能路径formatter
    formatter = SmartPathFormatter(
        '%(asctime)s | %(levelname)-8s | %(smartpath)s:%(lineno)d | %(funcName)s() | %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    
    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    
    # 文件处理器
    file_handler = logging.FileHandler('debug.log', encoding='utf-8')
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(formatter)
    
    # 添加处理器
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)



setup_logging()
logger = logging.getLogger("MTEB Benchmark")


@dataclass
class EvalArguments:
    model_name_or_path: str = field(
        default=None,
        metadata={"help": "Model identifier or path to pretrained model"}
    )
    model_name_or_path_doc: Optional[str] = field(
        default=None,
        metadata={"help": "Model identifier or path to pretrained model for document encoder"}
    )
    batch_size: int = field(
        default=16,
        metadata={"help": "Batch size for evaluation"}
    )
    max_length: int = field(
        default=512,
        metadata={"help": "Maximum sequence length for processing"}
    )
    truncate_dim: Optional[int] = field(
        default=None,
        metadata={"help": "Dimension to truncate embeddings (optional)"}
    )
    benchmark: str = field(
        default="MTEB(Multilingual, v2)",
        metadata={"help": "MTEB benchmark specification"}
    )
    tasks: List[str] = field(
        default=None,
        metadata={"help": "MTEB tasks specification"}
    )
    task_prompts_type: str = field(
        default="asymmetric",
        metadata={"help": "Type of task prompts. Choice: 'none', 'asymmetric', 'all'"},
    )
    task_prompts_file: str = field(
        default=None,
        metadata={"help": "Path to task prompts JSON file"}
    )
    output_folder: Optional[str] = field(
        default=None,
        metadata={"help": "Output directory for evaluation results"}
    )


# class SentenceTransformer(sentence_transformers.SentenceTransformer):
class SentenceTransformer(Wrapper):
    # 该类只能命名为 SentenceTransformer，否则初始化校验会出错，无法加载为 SentenceTransformer 模型
    # https://github.com/huggingface/sentence-transformers/blob/67ff0fe316fde4aa6c4901dd577123d8fd249be8/sentence_transformers/SentenceTransformer.py#L316
    def __init__(
        self, 
        model_name_or_path: str, 
        model_name_or_path_doc: str, 
        device: str = None, 
        **kwargs
    ):
        # 1. 初始化 Query 模型 (父类)
        # 这会创建 modules 比如: '0': Transformer, '1': Pooling
        self.query_model = sentence_transformers.SentenceTransformer(model_name_or_path=model_name_or_path, device=device, **kwargs)
        
        # 2. 初始化 Doc 模型
        self.doc_model = sentence_transformers.SentenceTransformer(model_name_or_path=model_name_or_path_doc, device=device, **kwargs)

        model_prompts = self.validate_task_to_prompt_name(self.query_model.prompts)
        self.model_prompts = self.validate_task_to_prompt_name(model_prompts)

        self.similarity_fn_name = "cosine"
        self.model_card_data = kwargs.get("model_card_data", None)
    
    def encode(
        self,
        sentences: Sequence[str],
        *,
        task_name: str,
        prompt_type: PromptType | None = None,
        **kwargs: Any,
    ) -> np.ndarray:
        """Encodes the given sentences using the encoder.

        Args:
            sentences: The sentences to encode.
            task_name: The name of the task. Sentence-transformers uses this to
                determine which prompt to use from a specified dictionary.
            prompt_type: The name type of prompt. (query or passage)
            **kwargs: Additional arguments to pass to the encoder.

            The order of priorities for prompt selection are:
                1. Composed prompt of task name + prompt type (query or passage)
                2. Specific task prompt
                3. Composed prompt of task type + prompt type (query or passage)
                4. Specific task type prompt
                5. Specific prompt type (query or passage)


        Returns:
            The encoded sentences.
        """
        prompt = None
        prompt_name = None
        if self.model_prompts is not None:
            prompt_name = self.get_prompt_name(
                self.model_prompts, task_name, prompt_type
            )
            prompt = self.model_prompts.get(prompt_name, None)
        if prompt_name:
            logger.info(
                f"Using {prompt_name=} for task={task_name} {prompt_type=} with {prompt=}"
            )
        else:
            logger.info(f"No model prompts found for task={task_name} {prompt_type=}")
        logger.info(f"Encoding {len(sentences)} sentences.")

        print(f"### Encode with prompt_type= {prompt_type}")
        if prompt_type == PromptType.query or not prompt_type:
            return self.encode_query(sentences, prompt=prompt, **kwargs)
        elif prompt_type == PromptType.document:
            print("$$$ HET: encode_document")
            return self.encode_document(sentences, prompt=prompt, **kwargs)
        else:
            raise ValueError(f"prompt_type = {prompt_type} not implemneted")

    def encode_query(self, sentences: Union[str, List[str]], **kwargs) -> Union[List[float], np.ndarray]:
        """
        Query 编码
        直接调用父类的 encode，它会调用我们上面重写的 forward
        """
        # return super().encode(sentences, **kwargs)
        print("#" * 200)
        print("encode query")
        print(kwargs)
        return self.query_model.encode(sentences, **kwargs)

    def encode_document(self, sentences: Union[str, List[str]], **kwargs) -> Union[List[float], np.ndarray]:
        """
        Doc 编码
        显式调用内部的 doc_model
        """
        print("#" * 200)
        print("encode doc")
        print(kwargs)
        return self.doc_model.encode(sentences, **kwargs)


def get_prompts(task_prompts_file: str, task_prompts_type: str) -> dict:
    if task_prompts_type == "none":
        return {}
    
    elif task_prompts_type == "asymmetric" or task_prompts_type == "all":
        asymmetric_prompts = {
            "Classification": "Instruct: classify the query into different classes. \n Query: ",
            "MultilabelClassification": "Instruct: classify the query into different classes. \n Query: ",
            "Clustering": "Instruct: classify the query into different classes. \n Query: ",
            "Reranking-query": "Instruct: Given a query, retrieve documents that answer the query. \n Query: ",
            "Retrieval-query": "Instruct: Given a query, retrieve documents that answer the query. \n Query: ",
            "InstructionRetrieval-query": "Instruct: Given a query, retrieve documents that answer the query. \n Query: ",
        }
        symmetric_prompts = {
            "BitextMining": "Instruct: Retrieve parallel sentences \n Query: ",
            "STS": "Instruct: Retrieve semantically similar text \n Query: ",
            "PairClassification": "Instruct: Retrieve semantically similar text \n Query: ",
        }

        prompts = asymmetric_prompts
        if task_prompts_type == "all":
            prompts.update(symmetric_prompts)

        ### Specific prompts
        if task_prompts_file is not None:
            with open(task_prompts_file) as f:
                instructs_dict = json.load(f)
            for task, instruct in instructs_dict.items():
                # 判断 instruct 是否为字符串
                if isinstance(instruct, str):
                    query_prompt = f"Instruct: {instruct} \n Query: "
                    if mteb.get_task(task).metadata.type in ["Reranking", "Retrieval", "InstructionRetrieval"]:
                        prompts[f"{task}-query"] = query_prompt
                    elif mteb.get_task(task).metadata.type in ["Classification", "MultilabelClassification", "Clustering"]:
                        prompts[f"{task}"] = query_prompt
                    elif mteb.get_task(task).metadata.type in ["STS", "BitextMining", "Summarization", "PairClassification"]:
                        if task_prompts_type == "all": 
                            prompts[f"{task}"] = query_prompt
                        else:
                            print(f"Parrallel task {task} instruction ignored.")
                    else:
                        print(f"""task type {mteb.get_task(task).metadata.type} not found""")
                        raise ValueError
                else:
                    query_prompt = f"""Instruct: {instruct["query"]} \n Query: """
                    document_prompt = f"""Instruct: {instruct["passage"]} \n Query: """
                    prompts[f"{task}-query"] = query_prompt
                    prompts[f"{task}-document"] = document_prompt
        
        return prompts
    else:
        raise ValueError(f"task_prompts_type={task_prompts_type} does not match any choices.")


def get_device(rank: int) -> str:
    """Detect available acceleration device."""
    if is_torch_npu_available():
        return f"npu:{rank}"
    if torch.cuda.is_available():
        return f"cuda:{rank}"
    raise RuntimeError("No available acceleration device (NPU/CUDA) found")


def setup_model(args, rank: int) -> sentence_transformers.SentenceTransformer:
    """Initialize sentence transformer model with proper configuration."""

    def extract_model_name(model_name_or_path):
        parent_dir = os.path.dirname(model_name_or_path)
        last_dir = os.path.basename(model_name_or_path)
        if not last_dir.startswith("checkpoint-") and not last_dir.startswith("embedding-checkpoint"):
            model_name = last_dir
        else:
            model_name = f"{os.path.basename(parent_dir)}-{last_dir}"
        return model_name

    device = get_device(rank)

    if not args.model_name_or_path_doc or args.model_name_or_path_doc.lower() == "none":
        model_name = extract_model_name(args.model_name_or_path)
    else:
        model_name = f"{extract_model_name(args.model_name_or_path)}-HET-{extract_model_name(args.model_name_or_path_doc)}"
    
    if args.task_prompts_type == "none":
        inst_suffic = ""
    elif args.task_prompts_type == "asymmetric":
        inst_suffic = "-asy-inst"
    elif args.task_prompts_type == "all":
        inst_suffic = "-all-inst"
    else:
        raise ValueError(f"task_prompts_type={args.task_prompts_type} does not match any choices.")

    model_name=f"KaLM/{model_name}{inst_suffic}"

    prompts = get_prompts(args.task_prompts_file, args.task_prompts_type)

    if not args.model_name_or_path_doc or args.model_name_or_path_doc.lower() == "none":
        model = sentence_transformers.SentenceTransformer(
            args.model_name_or_path,
            truncate_dim=args.truncate_dim,
            device=device,
            prompts=prompts,
            trust_remote_code=True,
            model_card_data=sentence_transformers.SentenceTransformerModelCardData(
                model_name=model_name
                ),
            model_kwargs={
                "torch_dtype": torch.bfloat16,
                "attn_implementation": "sdpa",
            },
        )
        model.max_seq_length = args.max_length
        model.tokenizer.model_max_length = args.max_length
    else:
        model = SentenceTransformer(
            model_name_or_path=args.model_name_or_path,
            model_name_or_path_doc=args.model_name_or_path_doc,
            truncate_dim=args.truncate_dim,
            device=device,
            prompts=prompts,
            trust_remote_code=True,
            model_card_data=sentence_transformers.SentenceTransformerModelCardData(
                model_name=model_name
                ),
            model_kwargs={
                "torch_dtype": torch.bfloat16,
                "attn_implementation": "sdpa",
            },
        )
        model.query_model.max_seq_length = args.max_length
        model.query_model.tokenizer.model_max_length = args.max_length
        model.doc_model.max_seq_length = args.max_length
        model.doc_model.tokenizer.model_max_length = args.max_length
    
    print()
    return model


def main_worker(rank: int, args, task_queue: mp.Queue):
    """Worker process for task evaluation."""
    # Environment setup
    time.sleep(1 + rank)
    os.environ["RANK"] = str(rank)
    logger.info(f"Worker {rank} initialized on device {get_device(rank)}")

    # Model initialization
    model = setup_model(args, rank)
    
    # Task processing loop
    while not task_queue.empty():
        task_name = task_queue.get()
        logger.info(f"Worker {rank} processing task: {task_name}")

        # Task execution
        evaluation = mteb.MTEB(tasks=[task_name])
        evaluation.run(
                    model,
                    output_folder=args.output_folder,
                    encode_kwargs={"batch_size": args.batch_size},
                    verbosity=2,
                )


def main():
    parser = HfArgumentParser(EvalArguments)
    args, = parser.parse_args_into_dataclasses()
    print("-" * 100)
    print(args)
    print("-" * 100)

    args.tasks = [task for task in args.tasks if task and task.lower() != "none"]

    mp_manager = mp.Manager()
    task_queue = mp_manager.Queue()
    if args.tasks is not None and len(args.tasks):
        tasks = args.tasks
    else:
        tasks = mteb.get_benchmark(args.benchmark)
        
    for task in tasks:
        task_queue.put(task)

    if args.truncate_dim <= 0:
        args.truncate_dim = None

    # Device setup
    try:
        world_size = torch.npu.device_count() if is_torch_npu_available() else torch.cuda.device_count()
    except Exception as e:
        logger.error(e)
        logger.warning("Setting world_size to 1")
        world_size = 1

    # Distributed environment setup
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "8081"
    
    # Start worker processes
    mp.spawn(
        main_worker,
        args=(args, task_queue),
        nprocs=world_size,
        join=True,
    )


if __name__ == "__main__":
    main()

[sanshi9523]

我将mteb升级到2.3.9版本后这里会报错--from mteb.encoder_interface import PromptType
在降到1.39.7版本后可以进行评测，跑了几个数据集，有的分数和榜单分数一样，有的还是分数很低

[sanshi9523]

from mteb.encoder_interface import PromptType
from mteb.models.wrapper import Wrapper

这两个包并不能在2.9.3版本使用，将版本降到1.39.7可以使用，测试的数据还是有的比榜单上的分数低很多

[YanshekWoo]

@ItsukiFujii 有空跟进下这里V2评测的问题

[YanshekWoo]

from mteb.encoder_interface import PromptType
from mteb.models.wrapper import Wrapper

这两个包并不能在2.9.3版本使用，将版本降到1.39.7可以使用，测试的数据还是有的比榜单上的分数低很多

@sanshi9523 请问有评测的代码或者示例吗，尤其是prompt的指定

[ItsukiFujii]

from mteb.encoder_interface import PromptType
from mteb.models.wrapper import Wrapper

这两个包并不能在2.9.3版本使用，将版本降到1.39.7可以使用，测试的数据还是有的比榜单上的分数低很多

具体是那个任务跑出来的分数低呢 @sanshi9523