Serialize tokenizer first load; one-shot degradation warning
96 worker threads racing transformers 5.x lazy imports on the FIRST _get_tokenizer call raised ImportError and degraded that whole first batch to the char approximation (the old single-threaded path never raced). First load now holds a threading.Lock; the transformers '>model_max_length' logging is silenced inside truncation (counting a 2M-token doc before trimming it is the point), and the per-sample degradation print becomes a one-shot warning. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
dafd171d4d
commit
d83c2cc1df
@ -9,21 +9,34 @@ Requires a tokenizer (transformers) at tokenizer_path or auto from the model.
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
import threading
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
DEFAULT_TRUNCATION_TOKENS = 32768 * 4 # 131072, mirrors evalside run.py
|
DEFAULT_TRUNCATION_TOKENS = 32768 * 4 # 131072, mirrors evalside run.py
|
||||||
|
|
||||||
|
|
||||||
|
_TOK_LOCK = threading.Lock()
|
||||||
|
|
||||||
|
|
||||||
@lru_cache(maxsize=4)
|
@lru_cache(maxsize=4)
|
||||||
def _get_tokenizer(tokenizer_path: str):
|
def _get_tokenizer(tokenizer_path: str):
|
||||||
if not tokenizer_path or not os.path.exists(tokenizer_path):
|
if not tokenizer_path or not os.path.exists(tokenizer_path):
|
||||||
raise FileNotFoundError(
|
raise FileNotFoundError(
|
||||||
f'tokenizer not found at {tokenizer_path!r} -- token-level truncation '
|
f'tokenizer not found at {tokenizer_path!r} -- token-level truncation '
|
||||||
'needs a local tokenizer dir (e.g. /data1/models/DeepSeek-V4-Flash-INT8)')
|
'needs a local tokenizer dir (e.g. /data1/models/DeepSeek-V4-Flash-INT8)')
|
||||||
from transformers import AutoTokenizer
|
# serialize the FIRST load: 96 worker threads racing transformers 5.x's
|
||||||
|
# lazy imports raised ImportError and silently degraded batches to the
|
||||||
|
# char approximation; after one success lru_cache serves the rest
|
||||||
|
with _TOK_LOCK:
|
||||||
|
from transformers import AutoTokenizer
|
||||||
|
|
||||||
return AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
|
# the '> model_max_length' warnings are EXPECTED here -- counting a
|
||||||
|
# 2M-token doc before trimming it is the whole point of truncation
|
||||||
|
import logging
|
||||||
|
|
||||||
|
logging.getLogger('transformers').setLevel(logging.ERROR)
|
||||||
|
return AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
|
||||||
|
|
||||||
|
|
||||||
def truncate_middle_tokens(text: str, max_tokens: int, tokenizer_path: str) -> str:
|
def truncate_middle_tokens(text: str, max_tokens: int, tokenizer_path: str) -> str:
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user