Megha AgarwalandZhou Taichang
9eeac248f2
Signed-off-by: Megha Agarwal <agarwalmegha1308@gmail.com> Signed-off-by: Megha Agarwal <agarwalmegha@microsoft.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
"""Multimodal output data structures for vLLM-Omni.
|
|
|
|
This module defines structured types for multimodal outputs.
|
|
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.outputs import CompletionOutput
|
|
|
|
|
|
@dataclass
|
|
class MultimodalPayload:
|
|
"""Structured multimodal output payload.
|
|
|
|
Attributes:
|
|
tensors: Dictionary mapping modality/key names to their tensors.
|
|
metadata: Optional dictionary for non-tensor metadata
|
|
(e.g., sample rate for audio, image dimensions).
|
|
"""
|
|
|
|
tensors: dict[str, torch.Tensor] = field(default_factory=dict)
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
@property
|
|
def primary_tensor(self) -> torch.Tensor | None:
|
|
"""Return the first tensor in the payload, or None if empty."""
|
|
if self.tensors:
|
|
return next(iter(self.tensors.values()))
|
|
return None
|
|
|
|
@property
|
|
def is_empty(self) -> bool:
|
|
"""Return True if the payload has no tensors."""
|
|
return len(self.tensors) == 0
|
|
|
|
def get(self, key: str) -> torch.Tensor | None:
|
|
"""Get a tensor by key, returning None if not found."""
|
|
return self.tensors.get(key)
|
|
|
|
def __contains__(self, key: str) -> bool:
|
|
return key in self.tensors
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.tensors)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None:
|
|
"""Create a MultimodalPayload from a raw dictionary.
|
|
|
|
Separates torch.Tensor values into tensors and everything
|
|
else into metadata.
|
|
"""
|
|
if not data:
|
|
return None
|
|
tensors: dict[str, torch.Tensor] = {}
|
|
metadata: dict[str, Any] = {}
|
|
for k, v in data.items():
|
|
if isinstance(v, torch.Tensor):
|
|
tensors[k] = v
|
|
else:
|
|
metadata[k] = v
|
|
if not tensors and not metadata:
|
|
return None
|
|
return cls(tensors=tensors, metadata=metadata)
|
|
|
|
|
|
@dataclass
|
|
class MultimodalCompletionOutput(CompletionOutput):
|
|
"""CompletionOutput with multimodal support.
|
|
|
|
Inherits all CompletionOutput fields and adds multimodal_output.
|
|
As a CompletionOutput subclass, compatible with all existing vLLM consumers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
multimodal_output: MultimodalPayload | None = None,
|
|
**kwargs: Any,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.multimodal_output = multimodal_output
|
|
|
|
def __repr__(self) -> str:
|
|
base = super().__repr__()
|
|
return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})"
|