chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,274 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user