FlagEmbedding/inference/embedder/model_mapping.py

from enum import Enum
from typing import Type, List
from dataclasses import dataclass
from collections import OrderedDict

from FlagEmbedding.abc.inference import AbsEmbedder
from FlagEmbedding.inference.embedder import FlagModel, BGEM3FlagModel, FlagLLMModel, FlagICLModel, FlagPseudoMoEModel


class EmbedderModelClass(Enum):
    ENCODER_ONLY_BASE = "encoder-only-base"
    ENCODER_ONLY_M3 = "encoder-only-m3"
    DECODER_ONLY_BASE = "decoder-only-base"
    DECODER_ONLY_ICL = "decoder-only-icl"
    DECODER_ONLY_PSEUDO_MOE = "decoder-only-pseudo_moe"


EMBEDDER_CLASS_MAPPING = OrderedDict([
    (EmbedderModelClass.ENCODER_ONLY_BASE, FlagModel),
    (EmbedderModelClass.ENCODER_ONLY_M3, BGEM3FlagModel),
    (EmbedderModelClass.DECODER_ONLY_BASE, FlagLLMModel),
    (EmbedderModelClass.DECODER_ONLY_ICL, FlagICLModel),
    (EmbedderModelClass.DECODER_ONLY_PSEUDO_MOE, FlagPseudoMoEModel)
])


class PoolingMethod(Enum):
    LAST_TOKEN = "last_token"
    CLS = "cls"
    MEAN = "mean"


@dataclass
class EmbedderConfig:
    model_class: Type[AbsEmbedder]
    pooling_method: PoolingMethod
    trust_remote_code: bool = False
    query_instruction_format: str = "{}{}"


# BGE models mapping
BGE_MAPPING = OrderedDict([
    (
        "bge-reasoner-embed-qwen3-8b-0923",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        "bge-code-v1",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="<instruct>{}\n<query>{}")
    ),
    (
        "bge-en-icl", 
        EmbedderConfig(FlagICLModel, PoolingMethod.LAST_TOKEN, query_instruction_format="<instruct>{}\n<query>{}")
    ),
    (
        "bge-multilingual-gemma2",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="<instruct>{}\n<query>{}")
    ),
    (
        "bge-m3",
        EmbedderConfig(BGEM3FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-large-en-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-base-en-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-small-en-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-large-zh-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-base-zh-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-small-zh-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-large-en",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-base-en",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-small-en",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-large-zh",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-base-zh",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        "bge-small-zh",
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
])

# Qwen3-Embedding models mapping
QWEN3_EMBEDDING_MAPPING = OrderedDict([
    (
        "Qwen3-Embedding-0.6B",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
    ),
    (
        "Qwen3-Embedding-4B",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
    ),
    (
        "Qwen3-Embedding-8B",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
    ),
])


# E5 models mapping
E5_MAPPING = OrderedDict([
    (
        "e5-mistral-7b-instruct",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        "e5-large-v2",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "e5-base-v2",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "e5-small-v2",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "multilingual-e5-large-instruct",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        "multilingual-e5-large",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "multilingual-e5-base",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "multilingual-e5-small",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "e5-large",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "e5-base",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        "e5-small",
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
])

# GTE models mapping
GTE_MAPPING = OrderedDict([
    (
        "gte-Qwen2-7B-instruct",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        "gte-Qwen2-1.5B-instruct",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        "gte-Qwen1.5-7B-instruct",
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        "gte-multilingual-base",
        EmbedderConfig(FlagModel, PoolingMethod.CLS, trust_remote_code=True)
    ),
    (
        "gte-large-en-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS, trust_remote_code=True)
    ),
    (
        "gte-base-en-v1.5",
        EmbedderConfig(FlagModel, PoolingMethod.CLS, True)
    ),
    (
        'gte-large',
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        'gte-base',
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        'gte-small',
        EmbedderConfig(FlagModel, PoolingMethod.MEAN)
    ),
    (
        'gte-large-zh',
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        'gte-base-zh',
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
    (
        'gte-small-zh',
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
])

# SFR models mapping
SFR_MAPPING = OrderedDict([
    (
        'SFR-Embedding-2_R',
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
    (
        'SFR-Embedding-Mistral',
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
])

# Linq models mapping
LINQ_MAPPING = OrderedDict([
    (
        'Linq-Embed-Mistral',
        EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
    ),
])

# BCE models mapping
BCE_MAPPING = OrderedDict([
    (
        'bce-embedding-base_v1',
        EmbedderConfig(FlagModel, PoolingMethod.CLS)
    ),
])

# Combine all mappings
AUTO_EMBEDDER_MAPPING = OrderedDict()
AUTO_EMBEDDER_MAPPING.update(BGE_MAPPING)
AUTO_EMBEDDER_MAPPING.update(QWEN3_EMBEDDING_MAPPING)
AUTO_EMBEDDER_MAPPING.update(E5_MAPPING)
AUTO_EMBEDDER_MAPPING.update(GTE_MAPPING)
AUTO_EMBEDDER_MAPPING.update(SFR_MAPPING)
AUTO_EMBEDDER_MAPPING.update(LINQ_MAPPING)
AUTO_EMBEDDER_MAPPING.update(BCE_MAPPING)

# TODO: Add more models, such as Jina, Stella_v5, NV-Embed, etc.

def support_native_bge_model_list()->List[str]:
    return list(BGE_MAPPING.keys())

def support_model_list()->List[str]:
    return list(AUTO_EMBEDDER_MAPPING.keys())