[Feature] Support reward model LxzGordon/URM-LLaMa-3.1-8B (#1525)

2024-09-27 23:32:11 -07:00
parent b1e330bcb0
commit 9aa6553d2a
13 changed files with 478 additions and 44 deletions
--- a/test/srt/models/test_embedding_models.py
+++ b/test/srt/models/test_embedding_models.py
@@ -39,7 +39,9 @@ class TestEmbeddingModels(unittest.TestCase):
        prefill_tolerance,
    ) -> None:
        with HFRunner(
-            model_path, torch_dtype=torch_dtype, is_generation=False
+            model_path,
+            torch_dtype=torch_dtype,
+            model_type="embedding",
        ) as hf_runner:
            hf_outputs = hf_runner.forward(prompts)

@@ -47,7 +49,7 @@ class TestEmbeddingModels(unittest.TestCase):
            model_path,
            tp_size=tp_size,
            torch_dtype=torch_dtype,
-            is_generation=False,
+            model_type="embedding",
        ) as srt_runner:
            srt_outputs = srt_runner.forward(prompts)