346 lines
14 KiB
Python
346 lines
14 KiB
Python
# #region estimate_token_budget [C:3] [TYPE Module] [SEMANTICS translate, token, budget, estimation, llm]
|
|
# @BRIEF Calculate safe batch_size and max_tokens for LLM translation calls based on actual content length and model context window limits.
|
|
# @LAYER Domain
|
|
# @RELATION DEPENDS_ON -> [TranslationExecutor:Module]
|
|
# @RELATION DEPENDS_ON -> [TranslationExecutor:Module]
|
|
# @RATIONALE Added comment clarifying PROVIDER_DEFAULTS is a fallback — primary source should be LLMProvider API.
|
|
|
|
# DeepSeek v4 Flash supports up to 64K context window; output is limited by max_tokens.
|
|
# @REJECTED External tokenizer library — would introduce heavy dependency for estimation only.
|
|
# Fixed batch_size of 50 — causes truncation on long-content rows.
|
|
|
|
# #region DEFAULT_CONTEXT_WINDOW [TYPE Constant]
|
|
|
|
DEFAULT_CONTEXT_WINDOW = 64000
|
|
# #endregion DEFAULT_CONTEXT_WINDOW
|
|
# #region DEFAULT_MAX_OUTPUT_TOKENS [TYPE Constant]
|
|
|
|
DEFAULT_MAX_OUTPUT_TOKENS = 16384
|
|
# #endregion DEFAULT_MAX_OUTPUT_TOKENS
|
|
# @BRIEF CoT reasoning overhead tokens for DeepSeek models (~2000 tokens for chain-of-thought).
|
|
# #region REASONING_OVERHEAD [TYPE Constant]
|
|
|
|
REASONING_OVERHEAD = 2000
|
|
# #endregion REASONING_OVERHEAD
|
|
|
|
# #region PROVIDER_DEFAULTS [TYPE Constant]
|
|
|
|
# Maps model name (or "default" fallback) to capacity limits.
|
|
# @RATIONALE Different providers have drastically different context windows and
|
|
# output limits. Using a single default for all causes either wasted
|
|
# capacity (underestimation) or truncation (overestimation).
|
|
PROVIDER_DEFAULTS: dict[str, dict[str, int]] = {
|
|
"gpt-4o-mini": {"context_window": 128000, "max_output_tokens": 16384},
|
|
"gpt-4o": {"context_window": 128000, "max_output_tokens": 16384},
|
|
"o1-mini": {"context_window": 128000, "max_output_tokens": 65536},
|
|
"claude-3-5-sonnet": {"context_window": 200000, "max_output_tokens": 8192},
|
|
"deepseek-v4-flash": {"context_window": 64000, "max_output_tokens": 8192},
|
|
"default": {"context_window": 64000, "max_output_tokens": 16384},
|
|
}
|
|
# #endregion PROVIDER_DEFAULTS
|
|
# Increased from 60 to 120 because SQL/dashboard text and JSON structure need more.
|
|
# #region OUTPUT_PER_ROW_PER_LANG [TYPE Constant]
|
|
|
|
# Increased from 60 to 120 because SQL/dashboard text and JSON structure need more.
|
|
OUTPUT_PER_ROW_PER_LANG = 120
|
|
# #endregion OUTPUT_PER_ROW_PER_LANG
|
|
# #region JSON_OVERHEAD_PER_ROW [TYPE Constant]
|
|
|
|
JSON_OVERHEAD_PER_ROW = 50
|
|
# #endregion JSON_OVERHEAD_PER_ROW
|
|
# #region PROMPT_BASE_TOKENS [TYPE Constant]
|
|
|
|
# Increased from 300 to 600 to account for longer template, system msg, and dict preamble.
|
|
PROMPT_BASE_TOKENS = 600
|
|
# #endregion PROMPT_BASE_TOKENS
|
|
# @BRIEF Cap for dictionary tokens in estimation to avoid overestimating.
|
|
# #region DICT_TOKENS_PER_ENTRY [TYPE Constant]
|
|
|
|
DICT_TOKENS_PER_ENTRY = 20
|
|
# #endregion DICT_TOKENS_PER_ENTRY
|
|
|
|
# #region DICT_TOKENS_MAX [TYPE Constant]
|
|
|
|
DICT_TOKENS_MAX = 5000
|
|
# #endregion DICT_TOKENS_MAX
|
|
# #region CHARS_PER_TOKEN_MIXED [TYPE Constant]
|
|
|
|
CHARS_PER_TOKEN_MIXED = 2.2
|
|
# #endregion CHARS_PER_TOKEN_MIXED
|
|
# #region MIN_MAX_TOKENS [TYPE Constant]
|
|
|
|
MIN_MAX_TOKENS = 4096
|
|
# #endregion MIN_MAX_TOKENS
|
|
# #region MAX_OUTPUT_HEADROOM [TYPE Constant]
|
|
|
|
# Increased from 1000 to 3000 because SQL/dashboard text output varies significantly.
|
|
MAX_OUTPUT_HEADROOM = 3000
|
|
# #endregion MAX_OUTPUT_HEADROOM
|
|
|
|
|
|
# region _estimate_tokens_for_text [TYPE Function]
|
|
# @BRIEF Estimate token count for a text string with CJK-aware heuristics.
|
|
# CJK characters (~1.5 chars/token) vs other text (~2.2 chars/token).
|
|
# @PRE text is a string.
|
|
# @POST Returns estimated token count >= 1.
|
|
# @RATIONALE CJK characters are more token-dense than Latin/Cyrillic text.
|
|
# Using a single ratio undercounts CJK input and causes truncation.
|
|
# @REJECTED Using tiktoken — would introduce a heavy dependency for estimation only.
|
|
def _estimate_tokens_for_text(text: str) -> int:
|
|
"""Estimate token count with CJK-aware heuristics.
|
|
|
|
CJK characters (CJK Unified Ideographs) use ~1.5 chars/token.
|
|
All other characters use ~2.2 chars/token.
|
|
"""
|
|
if not text:
|
|
return 1
|
|
|
|
cjk_count = 0
|
|
other_count = 0
|
|
for ch in text:
|
|
if '\u4e00' <= ch <= '\u9fff' or '\u3000' <= ch <= '\u303f' or '\uff00' <= ch <= '\uffef':
|
|
cjk_count += 1
|
|
else:
|
|
other_count += 1
|
|
|
|
cjk_tokens = cjk_count / 1.5 if cjk_count else 0
|
|
other_tokens = other_count / 2.2 if other_count else 0
|
|
return max(1, int(cjk_tokens + other_tokens))
|
|
# endregion _estimate_tokens_for_text
|
|
|
|
|
|
# region _count_rows_that_fit [TYPE Function]
|
|
# @BRIEF Count how many rows fit within the available input budget.
|
|
# @PRE input_per_row is non-empty; available_budget > 0.
|
|
# @POST Returns (safe_count, total_input_tokens). When no rows fit, returns (0, 0).
|
|
def _count_rows_that_fit(
|
|
input_per_row: list[int],
|
|
available_budget: int,
|
|
) -> tuple[int, int]:
|
|
"""Count consecutive rows that fit within available_budget.
|
|
|
|
Returns:
|
|
(safe_count, total_input_tokens): Number of rows and their total tokens.
|
|
Returns (0, 0) when the first row alone does not fit (MEDIUM fix).
|
|
"""
|
|
running_total = 0
|
|
safe_size = 0
|
|
for tokens in input_per_row:
|
|
if running_total + tokens + REASONING_OVERHEAD < available_budget:
|
|
running_total += tokens
|
|
safe_size += 1
|
|
else:
|
|
break
|
|
# MEDIUM: Return (0, 0) when no rows fit — signal upstream that oversizing occurred.
|
|
# Prevents silent clamping to 1 which would produce truncated LLM calls.
|
|
return safe_size, running_total
|
|
# endregion _count_rows_that_fit
|
|
|
|
|
|
# region estimate_token_budget [C:3] [TYPE Function]
|
|
# @BRIEF Estimate token budget for a batch of source rows and return safe batch parameters.
|
|
# @PRE source_rows is a list of dicts (can be empty). target_languages is a non-empty list.
|
|
# @POST Returns dict with batch_size_adjusted, estimated_input_tokens, estimated_output_tokens, max_output_needed, warning.
|
|
# @RATIONALE Uses character-count heuristics (chars/2.2 for mixed text) since exact tokenization
|
|
# depends on the LLM model. Estimates are intentionally conservative to prevent truncation.
|
|
# @REJECTED Using tiktoken or similar tokenizer — would introduce a heavy dependency and still
|
|
# not match DeepSeek's tokenizer exactly.
|
|
def _calculate_output_tokens(
|
|
safe_size: int,
|
|
num_languages: int,
|
|
) -> int:
|
|
"""Calculate estimated output tokens for a batch."""
|
|
return (
|
|
safe_size * num_languages * OUTPUT_PER_ROW_PER_LANG
|
|
+ safe_size * JSON_OVERHEAD_PER_ROW
|
|
+ REASONING_OVERHEAD
|
|
)
|
|
|
|
|
|
def _apply_output_aware_batch_sizing(
|
|
safe_size: int,
|
|
num_languages: int,
|
|
max_output_tokens: int,
|
|
) -> int:
|
|
"""Reduce batch size until estimated output fits within max_output_tokens."""
|
|
while safe_size > 0:
|
|
needed_output = (
|
|
safe_size * num_languages * OUTPUT_PER_ROW_PER_LANG
|
|
+ safe_size * JSON_OVERHEAD_PER_ROW
|
|
+ REASONING_OVERHEAD
|
|
+ MAX_OUTPUT_HEADROOM
|
|
)
|
|
if needed_output <= max_output_tokens:
|
|
break
|
|
safe_size -= 1
|
|
return safe_size
|
|
|
|
|
|
def _build_warning(
|
|
batch_size: int | None,
|
|
safe_size: int,
|
|
total_rows: int,
|
|
context_window: int,
|
|
estimated_input: int,
|
|
max_output_needed: int,
|
|
dict_warning: str | None,
|
|
) -> str | None:
|
|
"""Build warning message when batch size is reduced."""
|
|
warning = None
|
|
if batch_size and safe_size < batch_size:
|
|
total_estimated = estimated_input + max_output_needed
|
|
warning = (
|
|
f"Reduced batch size from {batch_size} to {safe_size} "
|
|
f"(estimated {total_estimated} tokens vs {context_window} window)"
|
|
)
|
|
elif not batch_size and safe_size < total_rows:
|
|
warning = (
|
|
f"Auto-calculated batch size of {safe_size} from {total_rows} rows "
|
|
f"(output-limited)"
|
|
)
|
|
if dict_warning:
|
|
warning = f"{warning}; {dict_warning}" if warning else dict_warning
|
|
return warning
|
|
|
|
|
|
def estimate_token_budget(
|
|
source_rows: list[dict],
|
|
target_languages: list[str],
|
|
source_column: str = "source_text",
|
|
context_columns: list[str] | None = None,
|
|
dictionary_entries: list | None = None,
|
|
batch_size: int | None = None,
|
|
context_window: int | None = None,
|
|
max_output_tokens: int | None = None,
|
|
provider_info: str | None = None,
|
|
) -> dict:
|
|
"""Estimate token budget and return safe batch parameters.
|
|
|
|
Args:
|
|
source_rows: List of row dicts with source text.
|
|
target_languages: List of target language codes.
|
|
source_column: Key for the source text in each row dict.
|
|
context_columns: Optional list of keys for context columns.
|
|
dictionary_entries: Optional list of dictionary entries for glossary.
|
|
batch_size: Desired batch size. If None, auto-calculate max safe size.
|
|
context_window: Model context window. If None, resolved from provider_info.
|
|
max_output_tokens: Hard max output tokens limit. If None, resolved from provider_info.
|
|
provider_info: Optional provider model name (e.g. "gpt-4o-mini") for provider-aware defaults.
|
|
When provided without explicit context_window/max_output_tokens, uses
|
|
PROVIDER_DEFAULTS to set appropriate limits.
|
|
|
|
Returns:
|
|
dict with:
|
|
batch_size_adjusted: Safe batch size (may be less than requested).
|
|
estimated_input_tokens: Total estimated input tokens for the batch.
|
|
estimated_output_tokens: Total estimated output tokens.
|
|
max_output_needed: Recommended max_tokens for this batch.
|
|
warning: str | None if batch was reduced.
|
|
available_input_budget: Total input token budget per batch (context_window - max_output_tokens).
|
|
max_output_tokens: The effective max output tokens limit used for this batch.
|
|
"""
|
|
if not target_languages:
|
|
target_languages = ["en"]
|
|
|
|
# Resolve provider-aware defaults
|
|
if provider_info and context_window is None and max_output_tokens is None:
|
|
provider_settings = PROVIDER_DEFAULTS.get(provider_info)
|
|
if provider_settings:
|
|
context_window = provider_settings["context_window"]
|
|
max_output_tokens = provider_settings["max_output_tokens"]
|
|
|
|
# Fall back to module defaults if still None
|
|
if context_window is None:
|
|
context_window = DEFAULT_CONTEXT_WINDOW
|
|
if max_output_tokens is None:
|
|
max_output_tokens = DEFAULT_MAX_OUTPUT_TOKENS
|
|
|
|
num_languages = len(target_languages)
|
|
|
|
# 1. Estimate tokens per row using CJK-aware heuristics
|
|
input_per_row = []
|
|
limit = batch_size if batch_size else len(source_rows)
|
|
for i, row in enumerate(source_rows):
|
|
if i >= limit:
|
|
break
|
|
text = str(row.get(source_column, "") or "")
|
|
estimated_tokens = _estimate_tokens_for_text(text)
|
|
if context_columns:
|
|
for col in context_columns:
|
|
val = str(row.get(col, "") or "")
|
|
estimated_tokens += _estimate_tokens_for_text(val)
|
|
input_per_row.append(estimated_tokens)
|
|
|
|
# 2. Calculate dictionary tokens with warning if capped
|
|
dict_tokens = 0
|
|
dict_warning = None
|
|
if dictionary_entries:
|
|
raw_dict_tokens = len(dictionary_entries) * DICT_TOKENS_PER_ENTRY
|
|
dict_tokens = min(raw_dict_tokens, DICT_TOKENS_MAX)
|
|
if raw_dict_tokens > DICT_TOKENS_MAX:
|
|
dict_warning = (
|
|
f"Dictionary entries ({len(dictionary_entries)} entries "
|
|
f"≈ {raw_dict_tokens} tokens) exceed cap of {DICT_TOKENS_MAX} — "
|
|
f"truncated to {DICT_TOKENS_MAX} in estimation"
|
|
)
|
|
|
|
prompt_tokens = PROMPT_BASE_TOKENS + dict_tokens
|
|
available_input_budget = context_window - max_output_tokens
|
|
|
|
# 3. Calculate safe batch size using shared helper
|
|
safe_size, input_total = _count_rows_that_fit(input_per_row, available_input_budget)
|
|
estimated_input = prompt_tokens + input_total
|
|
|
|
# If batch_size was specified and we reduced it, recalculate
|
|
if batch_size and safe_size < batch_size:
|
|
_, truncated_total = _count_rows_that_fit(
|
|
input_per_row[:batch_size], available_input_budget,
|
|
)
|
|
estimated_input = prompt_tokens + truncated_total
|
|
|
|
# 4. Estimate output tokens (includes JSON overhead per row)
|
|
estimated_output = _calculate_output_tokens(safe_size, num_languages)
|
|
|
|
# 5. Output-aware batch sizing: reduce batch if output exceeds limit
|
|
safe_size = _apply_output_aware_batch_sizing(
|
|
safe_size, num_languages, max_output_tokens,
|
|
)
|
|
|
|
# Ensure at least 1 row per batch — prevents empty batch allocation
|
|
safe_size = max(safe_size, 1)
|
|
|
|
# Recalculate totals after output-aware reduction
|
|
if safe_size > 0:
|
|
new_input_total = sum(input_per_row[:safe_size])
|
|
estimated_input = prompt_tokens + new_input_total
|
|
estimated_output = _calculate_output_tokens(safe_size, num_languages)
|
|
|
|
# 6. Calculate recommended max_tokens
|
|
max_output_needed = min(
|
|
estimated_output + MAX_OUTPUT_HEADROOM,
|
|
context_window - estimated_input,
|
|
)
|
|
max_output_needed = max(max_output_needed, MIN_MAX_TOKENS)
|
|
max_output_needed = min(max_output_needed, max_output_tokens)
|
|
|
|
# 7. Generate warning if batch was reduced
|
|
warning = _build_warning(
|
|
batch_size, safe_size, len(source_rows),
|
|
context_window, estimated_input, max_output_needed,
|
|
dict_warning,
|
|
)
|
|
|
|
return {
|
|
"batch_size_adjusted": safe_size,
|
|
"estimated_input_tokens": estimated_input,
|
|
"estimated_output_tokens": estimated_output,
|
|
"max_output_needed": max_output_needed,
|
|
"warning": warning,
|
|
"available_input_budget": available_input_budget,
|
|
"max_output_tokens": max_output_tokens,
|
|
}
|
|
# endregion estimate_token_budget
|
|
# #endregion estimate_token_budget
|
|
# endregion estimate_token_budget
|
|
# #endregion estimate_token_budget
|