49 lines
2.0 KiB
Python
49 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Helpers for explicit thinker decode span metadata."""
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
TensorSpan = tuple[torch.Tensor, int, int]
|
|
|
|
|
|
def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None:
|
|
tensor = payload.get(tensor_key)
|
|
start = payload.get(start_key)
|
|
end = payload.get(end_key)
|
|
if not isinstance(tensor, torch.Tensor):
|
|
return None
|
|
if not isinstance(start, int) or not isinstance(end, int):
|
|
return None
|
|
if start < 0 or end < start or (end - start) != int(tensor.shape[0]):
|
|
return None
|
|
return tensor, start, end
|
|
|
|
|
|
def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None:
|
|
if existing_span is None or incoming_span is None:
|
|
return None
|
|
|
|
existing_tensor, existing_start, existing_end = existing_span
|
|
incoming_tensor, incoming_start, incoming_end = incoming_span
|
|
if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype:
|
|
incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype)
|
|
if incoming_start == existing_end:
|
|
if existing_start == existing_end:
|
|
return incoming_tensor, incoming_start, incoming_end
|
|
return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end
|
|
if incoming_start < existing_end:
|
|
overlap = existing_end - incoming_start
|
|
if overlap >= int(incoming_tensor.shape[0]):
|
|
return existing_tensor, existing_start, existing_end
|
|
trimmed_tensor = incoming_tensor[overlap:]
|
|
return (
|
|
torch.cat([existing_tensor, trimmed_tensor], dim=0),
|
|
existing_start,
|
|
existing_end + int(trimmed_tensor.shape[0]),
|
|
)
|
|
return None
|