Token-level middle truncation (ported from evalside run.py: head+tail 128k, mrcr message-window); max_input_tokens param; perf fields (ttft/itl/status/retries) verified persisted in checkpoints

This commit is contained in:
sora 2026-08-27 03:23:33 +00:00
parent 3c63338ac1
commit c4473a1e1c
2 changed files with 81 additions and 0 deletions

View File

@ -35,6 +35,8 @@ async def generate_predictions(
system: str = '', system: str = '',
max_turns: int = 8, max_turns: int = 8,
max_input_chars: int = 0, max_input_chars: int = 0,
max_input_tokens: int = 0,
tokenizer_path: str = '',
attach_context_keys: tuple = ('passage', 'context'), attach_context_keys: tuple = ('passage', 'context'),
limit_per_task: Optional[int] = None, limit_per_task: Optional[int] = None,
checkpoint: Union[bool, str] = False, checkpoint: Union[bool, str] = False,
@ -113,6 +115,14 @@ async def generate_predictions(
'in the form "Answer: <answer>".') 'in the form "Answer: <answer>".')
parts.append(question) parts.append(question)
text = '\n\n'.join(parts) text = '\n\n'.join(parts)
if max_input_tokens:
try:
from .truncation import truncate_middle_tokens, default_tokenizer_path
text = truncate_middle_tokens(text, max_input_tokens,
tokenizer_path or default_tokenizer_path())
except Exception:
pass # no tokenizer: fall through to chars truncation
if max_input_chars and len(text) > max_input_chars: if max_input_chars and len(text) > max_input_chars:
keep = max_input_chars // 2 keep = max_input_chars // 2
head = text[:keep] head = text[:keep]
@ -268,6 +278,7 @@ async def run_eval(
system: str = '', system: str = '',
max_turns: int = 8, max_turns: int = 8,
max_input_chars: int = 0, max_input_chars: int = 0,
max_input_tokens: int = 0,
limit_per_task: Optional[int] = None, limit_per_task: Optional[int] = None,
checkpoint: Union[bool, str] = False, checkpoint: Union[bool, str] = False,
dataset_name: str = 'adhoc', dataset_name: str = 'adhoc',
@ -351,6 +362,7 @@ async def run_eval(
adapter, samples, concurrency, progress=progress, adapter, samples, concurrency, progress=progress,
gen_kwargs=gen_kwargs, env_factory=env_factory, gen_kwargs=gen_kwargs, env_factory=env_factory,
system=system, max_turns=max_turns, max_input_chars=max_input_chars, system=system, max_turns=max_turns, max_input_chars=max_input_chars,
max_input_tokens=max_input_tokens,
limit_per_task=limit_per_task, limit_per_task=limit_per_task,
checkpoint=checkpoint, checkpoint=checkpoint,
dataset_name=name, dataset_name=name,

View File

@ -0,0 +1,69 @@
"""Token-level middle truncation (ported from your /data1/sora/evalscope/bash/run.py).
Keeps head+tail halves of the token stream -- the industry-standard
middle-truncation for long-context benchmarks (longbench_v2 / mrcr).
The evalside run.py uses the same algorithm, guaranteeing comparable inputs.
Usage in run_eval: max_input_tokens=131072 (0/off = no truncation)
Requires a tokenizer (transformers) at tokenizer_path or auto from the model.
"""
import os
from functools import lru_cache
from typing import Optional
DEFAULT_TRUNCATION_TOKENS = 32768 * 4 # 131072, mirrors evalside run.py
@lru_cache(maxsize=4)
def _get_tokenizer(tokenizer_path: str):
if not tokenizer_path or not os.path.exists(tokenizer_path):
raise FileNotFoundError(
f'tokenizer not found at {tokenizer_path!r} -- token-level truncation '
'needs a local tokenizer dir (e.g. /data1/models/DeepSeek-V4-Flash-INT8)')
from transformers import AutoTokenizer
return AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
def truncate_middle_tokens(text: str, max_tokens: int, tokenizer_path: str) -> str:
"""Keep head+tail halves of the token stream; decode back to text."""
if max_tokens <= 0 or not text:
return text
tok = _get_tokenizer(tokenizer_path)
ids = tok.encode(text, add_special_tokens=False)
if len(ids) <= max_tokens:
return text
keep_head = max_tokens // 2
keep_tail = max_tokens - keep_head
return tok.decode(ids[:keep_head] + ids[-keep_tail:], skip_special_tokens=True)
def truncate_messages_middle(messages: list, max_tokens: int, tokenizer_path: str,
desired_index: int = 0, window: int = 2) -> list:
"""MRCR-style: when a message list exceeds the budget, keep first/last
messages plus a window around the desired (needle) message, dropping
middle spans; middle of each KEPT long message is token-truncated."""
if max_tokens <= 0 or not messages:
return messages
tok = _get_tokenizer(tokenizer_path)
total = sum(len(tok.encode(m.get('content', '') if isinstance(m, dict) else str(m),
add_special_tokens=False)) for m in messages)
if total <= max_tokens:
return messages
n = len(messages)
keep = set(range(min(2, n))) | set(range(max(0, n - 2), n))
di = desired_index if isinstance(desired_index, int) and 0 <= desired_index < n else 0
keep |= set(range(max(0, di - window), min(n, di + window + 1)))
return [messages[i] for i in sorted(keep)]
def default_tokenizer_path() -> Optional[str]:
"""Candidate local tokenizer for truncation (mirrors evalside default)."""
env = os.environ.get('EVALHARNESS_TOKENIZER')
if env and os.path.exists(env):
return env
for cand in ('/data1/models/DeepSeek-V4-Flash-INT8',):
if os.path.exists(cand):
return cand
return None