mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
"""segments → tensor_dict 的多模态转换：每段路由到对应模态 processor，产物合并。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
|
||
from deepinfer.ve.processor.common import BaseProcessor, ImageInput, SegmentInput, TextLengthInput, merge_tensor_dicts
|
||
|
||
|
||
class MMProcessor:
|
||
def __init__(
|
||
self,
|
||
text_token_processor: BaseProcessor = None,
|
||
image_processor: BaseProcessor = None,
|
||
) -> None:
|
||
self.text_token_processor = text_token_processor
|
||
self.image_processor = image_processor
|
||
self.processors_map = {
|
||
TextLengthInput: text_token_processor,
|
||
ImageInput: image_processor,
|
||
}
|
||
|
||
def convert_segments_to_tensor_dict(self, seg: SegmentInput) -> dict:
|
||
"""每个分段调用该模态对应的 processor，同名 tensor 合并。"""
|
||
# 空段种子保证零段 / 全缓存 prompt 也有稳定的键集合（mm_token_type 恒产出）
|
||
tensor_dicts: list[dict] = [{"mm_token_type": torch.zeros(0, dtype=torch.int32)}]
|
||
for segment in seg.segments:
|
||
processor: BaseProcessor = self.processors_map[type(segment)]
|
||
tensor_dicts.append(processor.process(segment))
|
||
return merge_tensor_dicts(tensor_dicts)
|