From a97409efef61a36baebe872bec06394b02a19d6b Mon Sep 17 00:00:00 2001
From: tzhouam <tzhouam@connect.ust.hk>
Date: Fri, 24 Oct 2025 11:40:43 +0800
Subject: [PATCH] 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.
---
 vllm_omni/engine/processor.py | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py
index 850da0b51..c49014487 100644
--- a/vllm_omni/engine/processor.py
+++ b/vllm_omni/engine/processor.py
@@ -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,