@@ -62,26 +62,28 @@ class TestSetting:
6262 TestSetting (
6363 model = "BAAI/bge-multilingual-gemma2" ,
6464 model_args = [
65- "--runner" , "pooling" , "--dtype" , "bfloat16" ,
66- "--max-model-len" , "2048"
65+ "--runner" ,
66+ "pooling" ,
67+ "--dtype" ,
68+ "bfloat16" ,
69+ "--max-model-len" ,
70+ "2048" ,
6771 ],
6872 pp_size = 1 ,
6973 tp_size = 1 ,
7074 attn_backend = "FLASH_ATTN" ,
7175 method = "encode" ,
7276 fullgraph = True ,
7377 ),
74- # TODO: bert models are not supported in V1 yet
75- # # encoder-based embedding model (BERT)
76- # TestSetting(
77- # model="BAAI/bge-base-en-v1.5",
78- # model_args=["--runner", "pooling"],
79- # pp_size=1,
80- # tp_size=1,
81- # attn_backend="XFORMERS",
82- # method="encode",
83- # fullgraph=True,
84- # ),
78+ TestSetting (
79+ model = "BAAI/bge-base-en-v1.5" ,
80+ model_args = ["--runner" , "pooling" ],
81+ pp_size = 1 ,
82+ tp_size = 1 ,
83+ attn_backend = "FLASH_ATTN" ,
84+ method = "encode" ,
85+ fullgraph = True ,
86+ ),
8587 # vision language model
8688 TestSetting (
8789 model = "microsoft/Phi-3.5-vision-instruct" ,
@@ -92,7 +94,8 @@ class TestSetting:
9294 method = "generate_with_image" ,
9395 fullgraph = False ,
9496 ),
95- ])
97+ ],
98+ )
9699def test_compile_correctness (
97100 monkeypatch : pytest .MonkeyPatch ,
98101 test_setting : TestSetting ,
0 commit comments