custom_generate/vllm_renderer.py

"""vLLM renderer that routes every ALTAY prompt through its model memory."""

from __future__ import annotations

from functools import cached_property
from typing import Any

from vllm.inputs import MultiModalInput, TokensInput
from vllm.renderers import TokenizeParams
from vllm.renderers.hf import HfRenderer


LOGICAL_CONTEXT_TOKENS = 1_010_000


class LomonosovZenitAltayRenderer(HfRenderer):
    """Separate the public logical context from the physical engine window."""

    @cached_property
    def default_cmpl_tok_params(self) -> TokenizeParams:
        return TokenizeParams(
            max_total_tokens=LOGICAL_CONTEXT_TOKENS,
            do_lower_case=False,
            add_special_tokens=True,
        )

    @cached_property
    def default_chat_tok_params(self) -> TokenizeParams:
        return TokenizeParams(
            max_total_tokens=LOGICAL_CONTEXT_TOKENS,
            do_lower_case=False,
            add_special_tokens=False,
        )

    def _process_tokens(
        self,
        prompt: Any,
        *,
        skip_mm_cache: bool = False,
    ) -> TokensInput | MultiModalInput:
        if prompt.get("multi_modal_data"):
            return super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)

        engine_input = self._process_multimodal(
            prompt["prompt_token_ids"],
            {},
            mm_processor_kwargs=prompt.get("mm_processor_kwargs"),
            tokenization_kwargs=None,
            mm_uuids=None,
            skip_mm_cache=skip_mm_cache,
        )
        if prompt_text := prompt.get("prompt"):
            engine_input["prompt"] = prompt_text
        if cache_salt := prompt.get("cache_salt"):
            engine_input["cache_salt"] = cache_salt
        return engine_input

    async def _process_tokens_async(
        self,
        prompt: Any,
        *,
        skip_mm_cache: bool = False,
    ) -> TokensInput | MultiModalInput:
        if prompt.get("multi_modal_data"):
            return await super()._process_tokens_async(
                prompt,
                skip_mm_cache=skip_mm_cache,
            )

        engine_input = await self._process_multimodal_async(
            prompt["prompt_token_ids"],
            {},
            mm_processor_kwargs=prompt.get("mm_processor_kwargs"),
            tokenization_kwargs=None,
            mm_uuids=None,
            skip_mm_cache=skip_mm_cache,
        )
        if prompt_text := prompt.get("prompt"):
            engine_input["prompt"] = prompt_text
        if cache_salt := prompt.get("cache_salt"):
            engine_input["cache_salt"] = cache_salt
        return engine_input


__all__ = ["LOGICAL_CONTEXT_TOKENS", "LomonosovZenitAltayRenderer"]