Add dtype mapping method in OmniProcessor

- Introduced a static method `_dtype_to_name` to map PyTorch data types to their string representations.
- Updated the `dtype_str` assignments in the `OmniProcessor` class to utilize the new mapping method for improved readability and maintainability.
This commit is contained in:
tzhouam committed 2025-10-24 11:40:43 +08:00
1 parent 4cd3b36fa0
commit a97409efef
1 file changed
+23 -2
+23 -2
View File
@@ -21,6 +21,27 @@ from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload,
import torch
class OmniProcessor(Processor):
@staticmethod
def _dtype_to_name(dtype: torch.dtype) -> str:
mapping = {
torch.float32: "float32",
torch.float: "float32",
torch.float16: "float16",
torch.half: "float16",
torch.bfloat16: "bfloat16",
torch.float64: "float64",
torch.double: "float64",
torch.int64: "int64",
torch.long: "int64",
torch.int32: "int32",
torch.int: "int32",
torch.int16: "int16",
torch.short: "int16",
torch.int8: "int8",
torch.uint8: "uint8",
torch.bool: "bool",
}
return mapping.get(dtype, str(dtype).replace("torch.", ""))
def __init__(self,
vllm_config: VllmConfig,
tokenizer: TokenizerGroup,
@@ -165,7 +186,7 @@ class OmniProcessor(Processor):
# Move to CPU and ensure contiguous memory for stable serialization
pe_cpu = pe.detach().to("cpu").contiguous()
seq_len, hidden_size = pe_cpu.shape
dtype_str = str(pe_cpu.dtype).replace("torch.", "")
dtype_str = self._dtype_to_name(pe_cpu.dtype)
data_bytes = pe_cpu.numpy().tobytes()
prompt_embeds_payload = PromptEmbedsPayload(
data=data_bytes,
@@ -178,7 +199,7 @@ class OmniProcessor(Processor):
for key, value in raw_info.items():
if isinstance(value, torch.Tensor):
v_cpu = value.detach().to("cpu").contiguous()
dtype_str = str(v_cpu.dtype).replace("torch.", "")
dtype_str = self._dtype_to_name(v_cpu.dtype)
data_bytes = v_cpu.numpy().tobytes()
entry = AdditionalInformationEntry(
tensor_data=data_bytes,