diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 9b7a42acece5..604cb854b32f 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -31,7 +31,7 @@ HYBRID_MODELS = [ # not compatible with pip-compile. "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", - "hmellor/bamba-tiny-random", + "hmellor/tiny-random-BambaForCausalLM", ] # Avoid OOM diff --git a/tests/models/registry.py b/tests/models/registry.py index a92dee3b642d..911a58e99d4c 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -124,7 +124,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "BaichuanForCausalLM": _HfExamplesInfo("baichuan-inc/Baichuan2-7B-chat", trust_remote_code=True), "BambaForCausalLM": _HfExamplesInfo("ibm-ai-platform/Bamba-9B", - extras={"tiny": "hmellor/bamba-tiny-random"}), # noqa: E501 + extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"}), # noqa: E501 "BloomForCausalLM": _HfExamplesInfo("bigscience/bloom-560m", {"1b": "bigscience/bloomz-1b1"}), "ChatGLMModel": _HfExamplesInfo("THUDM/chatglm3-6b", diff --git a/tests/v1/test_oracle.py b/tests/v1/test_oracle.py index c34c673e985e..1b77417a1bd3 100644 --- a/tests/v1/test_oracle.py +++ b/tests/v1/test_oracle.py @@ -12,7 +12,7 @@ UNSUPPORTED_MODELS_V1 = [ "openai/whisper-large-v3", # transcription "facebook/bart-large-cnn", # encoder decoder "mistralai/Mamba-Codestral-7B-v0.1", # mamba - "hmellor/bamba-tiny-random", # hybrid + "hmellor/tiny-random-BambaForCausalLM", # hybrid "BAAI/bge-m3", # embedding ]