Skip to content

Commit 8c98309

Browse files
committed
use modelscope
Signed-off-by: MengqingCao <cmq0113@163.com>
1 parent e3fa24f commit 8c98309

File tree

1 file changed

+28
-26
lines changed

1 file changed

+28
-26
lines changed

tests/singlecard/embedding/test_embedding.py

Lines changed: 28 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -60,34 +60,36 @@ def test_models(
6060
example_prompts,
6161
model,
6262
dtype: str,
63+
monkeypatch: pytest.MonkeyPatch,
6364
) -> None:
65+
with monkeypatch.context() as m:
66+
m.setenv("VLLM_USE_MODELSCOPE", "True")
67+
vllm_extra_kwargs: Dict[str, Any] = {}
6468

65-
vllm_extra_kwargs: Dict[str, Any] = {}
69+
# The example_prompts has ending "\n", for example:
70+
# "Write a short story about a robot that dreams for the first time.\n"
71+
# sentence_transformers will strip the input texts, see:
72+
# https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159
73+
# This makes the input_ids different between hf_model and vllm_model.
74+
# So we need to strip the input texts to avoid test failing.
75+
example_prompts = [str(s).strip() for s in example_prompts]
6676

67-
# The example_prompts has ending "\n", for example:
68-
# "Write a short story about a robot that dreams for the first time.\n"
69-
# sentence_transformers will strip the input texts, see:
70-
# https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159
71-
# This makes the input_ids different between hf_model and vllm_model.
72-
# So we need to strip the input texts to avoid test failing.
73-
example_prompts = [str(s).strip() for s in example_prompts]
77+
with vllm_runner(model,
78+
task="embed",
79+
dtype=dtype,
80+
max_model_len=None,
81+
**vllm_extra_kwargs) as vllm_model:
82+
vllm_outputs = vllm_model.encode(example_prompts)
7483

75-
with vllm_runner(model,
76-
task="embed",
77-
dtype=dtype,
78-
max_model_len=None,
79-
**vllm_extra_kwargs) as vllm_model:
80-
vllm_outputs = vllm_model.encode(example_prompts)
84+
with hf_runner(MODELSCOPE_CACHE + model,
85+
dtype=dtype,
86+
is_sentence_transformer=True) as hf_model:
87+
hf_outputs = hf_model.encode(example_prompts)
8188

82-
with hf_runner(MODELSCOPE_CACHE + model,
83-
dtype=dtype,
84-
is_sentence_transformer=True) as hf_model:
85-
hf_outputs = hf_model.encode(example_prompts)
86-
87-
check_embeddings_close(
88-
embeddings_0_lst=hf_outputs,
89-
embeddings_1_lst=vllm_outputs,
90-
name_0="hf",
91-
name_1="vllm",
92-
tol=1e-2,
93-
)
89+
check_embeddings_close(
90+
embeddings_0_lst=hf_outputs,
91+
embeddings_1_lst=vllm_outputs,
92+
name_0="hf",
93+
name_1="vllm",
94+
tol=1e-2,
95+
)

0 commit comments

Comments
 (0)