90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Multimodal data plumbing for Omni's custom, non-HF processors."""
|
|
|
|
from collections.abc import Mapping
|
|
from typing import TypeVar
|
|
|
|
from transformers import BatchFeature
|
|
from vllm.multimodal.parse import MultiModalDataItems
|
|
from vllm.multimodal.processing import (
|
|
BaseDummyInputsBuilder,
|
|
BaseMultiModalProcessor,
|
|
BaseProcessingInfo,
|
|
ProcessorInputs,
|
|
cached_encode,
|
|
)
|
|
|
|
_I = TypeVar("_I", bound=BaseProcessingInfo)
|
|
|
|
|
|
class OmniDummyInputsBuilder(BaseDummyInputsBuilder[_I]):
|
|
"""Restore Omni's processor-input hook on the vLLM 0.31 API."""
|
|
|
|
def get_dummy_processor_inputs(
|
|
self,
|
|
seq_len: int,
|
|
mm_counts: Mapping[str, int],
|
|
mm_options: Mapping[str, object],
|
|
) -> ProcessorInputs:
|
|
dummy_text = self.get_dummy_text(mm_counts)
|
|
dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
|
|
dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
|
|
|
|
tokenizer = self.info.ctx.tokenizer
|
|
dummy_prompt = (
|
|
[]
|
|
if tokenizer is None
|
|
else cached_encode(tokenizer, dummy_text, truncation=False)
|
|
)
|
|
return ProcessorInputs(prompt=dummy_prompt, mm_data_items=dummy_mm_items)
|
|
|
|
|
|
class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]):
|
|
"""Share data/passthrough handling, leaving model processing in its owner.
|
|
|
|
These processors implement ``_call_hf_processor`` locally, sometimes
|
|
without any HF processor object. Prompt tokens are prepared separately
|
|
in ``apply``; the upstream base still owns caching and prompt updates.
|
|
"""
|
|
|
|
_OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text"
|
|
|
|
def get_dummy_inputs(
|
|
self,
|
|
seq_len: int,
|
|
mm_counts: Mapping[str, int],
|
|
mm_options: Mapping[str, object],
|
|
) -> ProcessorInputs:
|
|
if isinstance(self.dummy_inputs, OmniDummyInputsBuilder):
|
|
return self.dummy_inputs.get_dummy_processor_inputs(
|
|
seq_len, mm_counts, mm_options
|
|
)
|
|
return super().get_dummy_inputs(seq_len, mm_counts, mm_options)
|
|
|
|
def _apply_hf_processor_main(
|
|
self,
|
|
mm_items: MultiModalDataItems,
|
|
hf_kwargs: Mapping[str, object],
|
|
) -> BatchFeature:
|
|
kwargs = dict(hf_kwargs)
|
|
prompt = kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, None)
|
|
if prompt is None:
|
|
prompt = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
|
|
valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
|
|
# Raw processor/passthrough data (upstream removed `_get_hf_mm_data` in
|
|
# favour of `_get_hf_mm_inputs`, which also renames `audios` -> `audio`
|
|
# and injects a dummy text key). Omni's custom `_call_hf_processor`
|
|
# subclasses still expect the un-renamed keys, so extract them here.
|
|
mm_data: dict[str, object] = {}
|
|
passthrough: dict[str, object] = {}
|
|
for items in valid_items.values():
|
|
if not items:
|
|
continue
|
|
mm_data.update(items.get_processor_data())
|
|
passthrough.update(items.get_passthrough_data())
|
|
result = self._call_hf_processor(str(prompt), dict(mm_data), kwargs, {})
|
|
result.update(passthrough)
|
|
return result
|