[Bugfix] Fix HfExampleModels.find_hf_info (#12223)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
vllm-project · Jan 20, 2025 · 18572e3 · 18572e3
1 parent 86bfb6d
commit 18572e3
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/models/registry.py b/tests/models/registry.py
@@ -302,6 +302,11 @@ def find_hf_info(self, model_id: str) -> _HfExamplesInfo:
             if info.default == model_id:
                 return info
 
+        # Fallback to extras
+        for info in self.hf_models.values():
+            if any(extra == model_id for extra in info.extras.values()):
+                return info
+
         raise ValueError(f"No example model defined for {model_id}")