diff --git a/tests/models/decoder_only/language/test_hybrid.py b/tests/models/decoder_only/language/test_hybrid.py index 5931c25b8d80..e5e0c28ae2d5 100644 --- a/tests/models/decoder_only/language/test_hybrid.py +++ b/tests/models/decoder_only/language/test_hybrid.py @@ -28,7 +28,7 @@ HYBRID_MODELS = [ # not compatible with pip-compile. "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", - "ibm-ai-platform/Bamba-9B", + "hmellor/bamba-tiny-random", ] # Avoid OOM diff --git a/tests/models/registry.py b/tests/models/registry.py index 75832d83d725..f17f70189a9b 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -120,7 +120,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { trust_remote_code=True), "BaichuanForCausalLM": _HfExamplesInfo("baichuan-inc/Baichuan2-7B-chat", trust_remote_code=True), - "BambaForCausalLM": _HfExamplesInfo("ibm-ai-platform/Bamba-9B"), + "BambaForCausalLM": _HfExamplesInfo("ibm-ai-platform/Bamba-9B", + extras={"tiny": "hmellor/bamba-tiny-random"}), # 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 94c8ad7c94f6..c34c673e985e 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 - "ibm-ai-platform/Bamba-9B", # hybrid + "hmellor/bamba-tiny-random", # hybrid "BAAI/bge-m3", # embedding ]