generation.py

"""Small generation helpers that do not require loading the model."""

from __future__ import annotations

import math
from collections.abc import Iterable


def gpu_duration_seconds(
    prompt_characters: int,
    output_tokens: int,
    max_context_tokens: int,
    max_duration_seconds: int = 120,
) -> int:
    """Estimate a ZeroGPU reservation for prompt prefill plus generation.

    Large coding prompts spend substantial GPU time on their 32k-token prefill
    even when the requested answer is short. Character count is available to
    the ZeroGPU duration callback before tokenization and is a conservative
    proxy for that cost.
    """
    characters = max(0, int(prompt_characters))
    output = max(1, int(output_tokens))
    context_limit = max(1, int(max_context_tokens))
    estimated_input_tokens = min(context_limit, math.ceil(characters / 3))
    estimate = (
        25
        + math.ceil(estimated_input_tokens * 0.005)
        + math.ceil(output * 0.08)
    )
    return min(max_duration_seconds, max(30, estimate))


def head_tail_token_counts(
    total_tokens: int,
    token_budget: int,
    preserved_prefix_tokens: int,
) -> tuple[int, int]:
    """Split an oversized prompt budget between its prefix and recent tail.

    Keeping only the tail can erase system instructions, tool definitions, and
    a task stated before a large code block. Keeping a bounded prefix plus the
    largest possible tail retains both the operating contract and the newest
    conversation state.
    """
    total = max(0, int(total_tokens))
    budget = max(1, int(token_budget))
    if total <= budget:
        return total, 0

    prefix = max(0, int(preserved_prefix_tokens))
    head = min(prefix, budget - 1)
    return head, budget - head


def ensure_bos_token(prompt: str, bos_token: str | None) -> str:
    """Prefix the model's BOS token when the chat template omits it."""
    if not bos_token or prompt.startswith(bos_token):
        return prompt
    return bos_token + prompt


def merge_eos_token_ids(
    model_ids: int | Iterable[int] | None,
    tokenizer_id: int | None,
) -> int | list[int] | None:
    """Keep every model stop token while preserving the tokenizer fallback."""
    if isinstance(model_ids, int):
        candidates = [model_ids]
    elif model_ids is None:
        candidates = []
    else:
        candidates = list(model_ids)

    if tokenizer_id is not None:
        candidates.append(tokenizer_id)

    unique: list[int] = []
    for candidate in candidates:
        if isinstance(candidate, int) and candidate not in unique:
            unique.append(candidate)

    if not unique:
        return None
    if len(unique) == 1:
        return unique[0]
    return unique