| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm.config.pooler import PoolerConfig |
| 5 | |
| 6 | |
| 7 | def test_idefics_multimodal( |
| 8 | vllm_runner, |
| 9 | ) -> None: |
| 10 | prompts = [ |
| 11 | "Hello, my name is", |
| 12 | "The president of the United States is", |
| 13 | "The capital of France is", |
| 14 | "The future of AI is", |
| 15 | ] |
| 16 | |
| 17 | with vllm_runner( |
| 18 | model_name="HuggingFaceM4/Idefics3-8B-Llama3", |
| 19 | runner="pooling", |
| 20 | convert="classify", |
| 21 | load_format="dummy", |
| 22 | max_model_len=512, |
| 23 | enforce_eager=True, |
| 24 | tensor_parallel_size=1, |
| 25 | disable_log_stats=True, |
| 26 | dtype="bfloat16", |
| 27 | ) as vllm_model: |
| 28 | llm = vllm_model.get_llm() |
| 29 | outputs = llm.classify(prompts) |
| 30 | for output in outputs: |
| 31 | assert len(output.outputs.probs) == 2 |
| 32 | |
| 33 | |
| 34 | def update_config(config): |
| 35 | text_config = config.get_text_config() |
| 36 | text_config.update( |
| 37 | { |
| 38 | "architectures": ["Gemma3ForSequenceClassification"], |
| 39 | "classifier_from_token": ["A", "B", "C", "D", "E"], |
| 40 | "method": "no_post_processing", |
| 41 | "id2label": { |
| 42 | "A": "Chair", |
| 43 | "B": "Couch", |
| 44 | "C": "Table", |
| 45 | "D": "Bed", |
| 46 | "E": "Cupboard", |
| 47 | }, |
| 48 | } |
| 49 | ) |
| 50 | return config |
| 51 | |
| 52 | |
| 53 | def test_gemma_multimodal( |
| 54 | vllm_runner, |
| 55 | ) -> None: |
| 56 | messages = [ |
| 57 | { |
| 58 | "role": "system", |
| 59 | "content": """ |
| 60 | You are a helpful assistant. You will be given a product description |
| 61 | which may also include an image. Classify the following product into |
| 62 | one of the categories: |
| 63 | |
| 64 | A = chair |
| 65 | B = couch |
| 66 | C = table |
| 67 | D = bed |
| 68 | E = cupboard |
| 69 | |
| 70 | You'll answer with exactly one letter (A, B, C, D, or E).""", |
| 71 | }, |
| 72 | { |
| 73 | "role": "user", |
| 74 | "content": [ |
| 75 | { |
| 76 | "type": "image_url", |
| 77 | "image_url": { |
| 78 | "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/red_chair.jpg" |
| 79 | }, |
| 80 | }, |
| 81 | {"type": "text", "text": "A fine 19th century piece of furniture."}, |
| 82 | ], |
| 83 | }, |
| 84 | ] |
| 85 | |
| 86 | with vllm_runner( |
| 87 | model_name="google/gemma-3-4b-it", |
| 88 | runner="pooling", |
| 89 | convert="classify", |
| 90 | load_format="auto", |
| 91 | hf_overrides=update_config, |
| 92 | pooler_config=PoolerConfig(seq_pooling_type="LAST"), |
| 93 | max_model_len=512, |
| 94 | enforce_eager=True, |
| 95 | tensor_parallel_size=1, |
| 96 | disable_log_stats=True, |
| 97 | dtype="bfloat16", |
| 98 | ) as vllm_model: |
| 99 | llm = vllm_model.get_llm() |
| 100 | prompts = llm._preprocess_chat([messages]) |
| 101 | |
| 102 | result = llm.classify(prompts) |
| 103 | assert result[0].outputs.probs[0] > 0.95 |
| 104 | assert all(c < 0.05 for c in result[0].outputs.probs[1:]) |