[feature] support for roberta embedding models (#5730)

2025-04-27 09:47:06 +08:00
parent c5e1026f47
commit d8fbc7c096
3 changed files with 186 additions and 2 deletions
--- a/test/srt/models/test_encoder_embedding_models.py
+++ b/test/srt/models/test_encoder_embedding_models.py
@@ -25,10 +25,10 @@ from transformers import AutoConfig, AutoTokenizer
 from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
 from sglang.test.test_utils import CustomTestCase, get_similarities, is_in_ci

-MODELS = [("BAAI/bge-small-en", 1, 1e-5)]
+MODELS = [("BAAI/bge-small-en", 1, 1e-5), ("BAAI/bge-m3", 1, 1e-5)]

 ATTENTION_BACKEND = ["torch_native", "triton"]
-BATCH_SIZE = [30]
+BATCH_SIZE = [1, 2]
 TORCH_DTYPES = [torch.float32]
 sgl_to_st_ratio = []