diff --git a/evalharness/model/runner.py b/evalharness/model/runner.py index 569715a..682cc0d 100644 --- a/evalharness/model/runner.py +++ b/evalharness/model/runner.py @@ -35,6 +35,8 @@ async def generate_predictions( system: str = '', max_turns: int = 8, max_input_chars: int = 0, + max_input_tokens: int = 0, + tokenizer_path: str = '', attach_context_keys: tuple = ('passage', 'context'), limit_per_task: Optional[int] = None, checkpoint: Union[bool, str] = False, @@ -113,6 +115,14 @@ async def generate_predictions( 'in the form "Answer: ".') parts.append(question) 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: keep = max_input_chars // 2 head = text[:keep] @@ -268,6 +278,7 @@ async def run_eval( system: str = '', max_turns: int = 8, max_input_chars: int = 0, + max_input_tokens: int = 0, limit_per_task: Optional[int] = None, checkpoint: Union[bool, str] = False, dataset_name: str = 'adhoc', @@ -351,6 +362,7 @@ async def run_eval( adapter, samples, concurrency, progress=progress, gen_kwargs=gen_kwargs, env_factory=env_factory, system=system, max_turns=max_turns, max_input_chars=max_input_chars, + max_input_tokens=max_input_tokens, limit_per_task=limit_per_task, checkpoint=checkpoint, dataset_name=name, diff --git a/evalharness/model/truncation.py b/evalharness/model/truncation.py new file mode 100644 index 0000000..ac7fac2 --- /dev/null +++ b/evalharness/model/truncation.py @@ -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