[Model] Pooling models default to using chunked prefill & prefix caching if supported. (#20930)
Signed-off-by: wang.yuqi <noooop@126.com>
This commit is contained in:
@@ -4,57 +4,58 @@ from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from ...utils import check_transformers_version
|
||||
from .embed_utils import EmbedModelInfo, correctness_test_embed_models
|
||||
from ...utils import (CLSPoolingEmbedModelInfo, EmbedModelInfo,
|
||||
LASTPoolingEmbedModelInfo, check_transformers_version)
|
||||
from .embed_utils import correctness_test_embed_models
|
||||
from .mteb_utils import mteb_test_embed_models
|
||||
|
||||
MODELS = [
|
||||
########## BertModel
|
||||
EmbedModelInfo("thenlper/gte-large",
|
||||
architecture="BertModel",
|
||||
enable_test=True),
|
||||
EmbedModelInfo("thenlper/gte-base",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
EmbedModelInfo("thenlper/gte-small",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
EmbedModelInfo("thenlper/gte-large-zh",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
EmbedModelInfo("thenlper/gte-base-zh",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
EmbedModelInfo("thenlper/gte-small-zh",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
CLSPoolingEmbedModelInfo("thenlper/gte-large",
|
||||
architecture="BertModel",
|
||||
enable_test=True),
|
||||
CLSPoolingEmbedModelInfo("thenlper/gte-base",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
CLSPoolingEmbedModelInfo("thenlper/gte-small",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
CLSPoolingEmbedModelInfo("thenlper/gte-large-zh",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
CLSPoolingEmbedModelInfo("thenlper/gte-base-zh",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
CLSPoolingEmbedModelInfo("thenlper/gte-small-zh",
|
||||
architecture="BertModel",
|
||||
enable_test=False),
|
||||
########### NewModel
|
||||
EmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
|
||||
architecture="GteNewModel",
|
||||
enable_test=True),
|
||||
EmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
|
||||
architecture="GteNewModel",
|
||||
enable_test=True),
|
||||
EmbedModelInfo("Alibaba-NLP/gte-large-en-v1.5",
|
||||
architecture="GteNewModel",
|
||||
enable_test=True),
|
||||
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
|
||||
architecture="GteNewModel",
|
||||
enable_test=True),
|
||||
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
|
||||
architecture="GteNewModel",
|
||||
enable_test=True),
|
||||
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-large-en-v1.5",
|
||||
architecture="GteNewModel",
|
||||
enable_test=True),
|
||||
########### Qwen2ForCausalLM
|
||||
EmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
|
||||
architecture="Qwen2ForCausalLM",
|
||||
enable_test=True),
|
||||
LASTPoolingEmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
|
||||
architecture="Qwen2ForCausalLM",
|
||||
enable_test=True),
|
||||
########## ModernBertModel
|
||||
EmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
|
||||
architecture="ModernBertModel",
|
||||
enable_test=True),
|
||||
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
|
||||
architecture="ModernBertModel",
|
||||
enable_test=True),
|
||||
########## Qwen3ForCausalLM
|
||||
EmbedModelInfo("Qwen/Qwen3-Embedding-0.6B",
|
||||
architecture="Qwen3ForCausalLM",
|
||||
dtype="float32",
|
||||
enable_test=True),
|
||||
EmbedModelInfo("Qwen/Qwen3-Embedding-4B",
|
||||
architecture="Qwen3ForCausalLM",
|
||||
dtype="float32",
|
||||
enable_test=False),
|
||||
LASTPoolingEmbedModelInfo("Qwen/Qwen3-Embedding-0.6B",
|
||||
architecture="Qwen3ForCausalLM",
|
||||
dtype="float32",
|
||||
enable_test=True),
|
||||
LASTPoolingEmbedModelInfo("Qwen/Qwen3-Embedding-4B",
|
||||
architecture="Qwen3ForCausalLM",
|
||||
dtype="float32",
|
||||
enable_test=False),
|
||||
]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user