gemma3.py 1.91 KB
Newer Older
1
2
import re
from typing import Dict, List, Union
3

Mick's avatar
Mick committed
4
5
from sglang.srt.managers.multimodal_processor import (
    BaseMultimodalProcessor as SGLangBaseProcessor,
6
7
)
from sglang.srt.models.gemma3_mm import Gemma3ForConditionalGeneration
8
from sglang.srt.multimodal.processors.base_processor import MultimodalSpecialTokens
9
10
11
12
13

# Copied from: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma3/image_processing_gemma3_fast.py
# will be removed in the future


Mick's avatar
Mick committed
14
class Gemma3SGLangImageProcessor(SGLangBaseProcessor):
15
16
    models = [Gemma3ForConditionalGeneration]

17
18
19
20
    def __init__(self, hf_config, server_args, _processor):
        super().__init__(hf_config, server_args, _processor)
        self.IM_START_TOKEN_ID = hf_config.boi_token_index
        self.IM_END_TOKEN_ID = hf_config.eoi_token_index
21
22
23
24
25
26
27
28
29
        self.mm_tokens = MultimodalSpecialTokens(
            # The single, pre-expanded image token.
            image_token="<start_of_image>",
            image_token_id=hf_config.image_token_index,
            # The regex that matches expanded image tokens.
            image_token_regex=re.compile(
                r"<start_of_image>(?:(?:<image_soft_token>)*<end_of_image>)?"
            ),
        ).build(_processor)
30

Mick's avatar
Mick committed
31
    async def process_mm_data_async(
32
        self,
33
        image_data: List[Union[str, bytes, Dict]],
34
        input_text,
35
36
37
38
        request_obj,
        *args,
        **kwargs,
    ):
Mick's avatar
Mick committed
39
        base_output = self.load_mm_data(
40
            prompt=input_text,
41
            image_data=image_data,
42
            multimodal_tokens=self.mm_tokens,
43
44
45
            discard_alpha_channel=True,
        )

46
47
48
        mm_items, input_ids, _ = self.process_and_combine_mm_data(
            base_output, self.mm_tokens
        )
49
        return {
50
            "input_ids": input_ids.tolist(),
51
            "mm_items": mm_items,
52
53
54
            "im_start_id": self.IM_START_TOKEN_ID,
            "im_end_id": self.IM_END_TOKEN_ID,
        }