refactor: simply MultimodalTokens logic (#7924)

This commit is contained in:
Mick
2025-07-18 08:59:15 +08:00
committed by GitHub
parent 3586b4cef2
commit e1020dc588
13 changed files with 145 additions and 132 deletions

View File

@@ -21,7 +21,7 @@ class BaseMultiModalProcessorOutput:
# input_text, with each frame of video/image represented with a image_token
input_text: str
# frames loaded from image and video, in given order
# frames loaded from image, in given order
images: Optional[list[Union[Image.Image, dict]]] = None
# videos
@@ -44,14 +44,26 @@ class BaseMultiModalProcessorOutput:
@dataclasses.dataclass
class MultimodalSpecialTokens:
image_token: Optional[Union[int, str, List[str]]] = None
video_token: Optional[Union[int, str, List[str]]] = None
audio_token: Optional[Union[int, str, List[str]]] = None
image_token: Optional[Union[str, List[str]]] = None
video_token: Optional[Union[str, List[str]]] = None
audio_token: Optional[Union[str, List[str]]] = None
image_token_id: Optional[int] = None
video_token_id: Optional[int] = None
audio_token_id: Optional[int] = None
image_token_regex: Optional[re.Pattern] = None
video_token_regex: Optional[re.Pattern] = None
audio_token_regex: Optional[re.Pattern] = None
combined_regex: Optional[re.Pattern] = None
def build(self, processor):
self.convert_to_strs(processor)
self.parse_regex()
self.get_combined_regex()
return self
def convert_to_str(self, token: Union[str, int], processor) -> str:
if token is None:
return token
@@ -60,11 +72,14 @@ class MultimodalSpecialTokens:
return processor.tokenizer.convert_ids_to_tokens([token])[0]
def convert_to_strs(self, processor):
self.image_token = self.convert_to_str(self.image_token, processor)
self.video_token = self.convert_to_str(self.video_token, processor)
self.audio_token = self.convert_to_str(self.audio_token, processor)
if not self.image_token:
self.image_token = self.convert_to_str(self.image_token_id, processor)
if not self.video_token:
self.video_token = self.convert_to_str(self.video_token_id, processor)
if not self.audio_token:
self.audio_token = self.convert_to_str(self.audio_token_id, processor)
def get_modality_of_token(self, token) -> Optional[Modality]:
def get_modality_of_token(self, token: str) -> Optional[Modality]:
"""
:return: the modality associated with the given token, if the token is a special_token or matches with the multimodal token regex
"""
@@ -94,7 +109,12 @@ class MultimodalSpecialTokens:
if self.audio_token_regex is None and self.audio_token is not None:
self.audio_token_regex = re.compile(re.escape(self.audio_token))
def combine_regex(self) -> re.Pattern:
def get_combined_regex(self) -> re.Pattern:
"""
Builds and returns a regex, used to split input str into tokens (with mm special tokens)
"""
if self.combined_regex:
return self.combined_regex
tokens = [
self.image_token_regex,
self.video_token_regex,
@@ -107,7 +127,8 @@ class MultimodalSpecialTokens:
patterns.append(t.pattern)
flags |= t.flags
combined = "(" + "|".join(f"(?:{p})" for p in patterns) + ")"
return re.compile(combined, flags)
self.combined_regex = re.compile(combined, flags)
return self.combined_regex
class BaseMultimodalProcessor(ABC):
@@ -341,9 +362,8 @@ class BaseMultimodalProcessor(ABC):
discard_alpha_channel: if True, discards the alpha channel in the returned images
"""
multimodal_tokens.convert_to_strs(self._processor)
multimodal_tokens.parse_regex()
multimodal_tokens_pattern = multimodal_tokens.combine_regex()
multimodal_tokens_pattern = multimodal_tokens.get_combined_regex()
if isinstance(prompt, list) and return_text:
assert len(prompt) and isinstance(prompt[0], int)
prompt = self._processor.tokenizer.decode(prompt)
@@ -445,7 +465,6 @@ class BaseMultimodalProcessor(ABC):
return result = [(2,4),(6,7)]
"""
mask = input_ids == mm_token_id
start_positions = (mask & ~torch.roll(mask, 1)).nonzero(as_tuple=True)[0]
end_positions = (mask & ~torch.roll(mask, -1)).nonzero(as_tuple=True)[0]
@@ -554,7 +573,9 @@ class BaseMultimodalProcessor(ABC):
return collected_items, input_ids, ret
def process_and_combine_mm_data(
self, base_output: BaseMultiModalProcessorOutput
self,
base_output: BaseMultiModalProcessorOutput,
mm_tokens: MultimodalSpecialTokens,
) -> Tuple[List[MultimodalDataItem], torch.Tensor, dict]:
"""
Process multimodal data and return the combined multimodal items and input_ids.
@@ -618,22 +639,14 @@ class BaseMultimodalProcessor(ABC):
# Add offsets to all items
for mm_item in all_collected_items:
if mm_item.modality in [Modality.IMAGE, Modality.MULTI_IMAGES]:
mm_item.offsets = self.get_mm_items_offset(
input_ids=input_ids,
mm_token_id=self.IM_TOKEN_ID,
)
elif mm_item.modality == Modality.AUDIO:
mm_item.offsets = self.get_mm_items_offset(
input_ids=input_ids,
mm_token_id=self.AUDIO_TOKEN_ID,
)
elif mm_item.modality == Modality.VIDEO:
mm_item.offsets = self.get_mm_items_offset(
input_ids=input_ids,
mm_token_id=self.VIDEO_TOKEN_ID,
)
else:
raise ValueError(f"Unknown modality: {mm_item.modality}")
mm_item.offsets = self.get_mm_items_offset(
input_ids=input_ids,
mm_token_id={
Modality.IMAGE: mm_tokens.image_token_id,
Modality.MULTI_IMAGES: mm_tokens.image_token_id,
Modality.VIDEO: mm_tokens.video_token_id,
Modality.AUDIO: mm_tokens.audio_token_id,
}.get(mm_item.modality, None),
)
return all_collected_items, input_ids, ret