embeddings: verify chunk sizes after retokenization + serialize llama ctx access

Two remaining GGUF crashers: (1) detokenize→retokenize is not
token-count-stable, so a chunk cut at exactly n_ctx could re-inflate
past it inside embed() and trip the GGML_ASSERT — chunks are now
shrunk until their re-tokenized length verifiably fits (n_ctx-64
margin); (2) llama.cpp contexts are not thread-safe and parallel
indexer requests segfaulted the engine (-11) — embed calls now hold a
per-model lock.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014S8VtAvG499SsCbeESRK7V
parent ffd6629c
...@@ -16,7 +16,7 @@ ...@@ -16,7 +16,7 @@
# Canonical product version for CoderAI — single source of truth. Both the API # Canonical product version for CoderAI — single source of truth. Both the API
# metadata and the admin web UI read from here. # metadata and the admin web UI read from here.
__version__ = "0.1.40" __version__ = "0.1.41"
# Configure the CUDA caching allocator BEFORE torch is imported anywhere. # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
# expandable_segments lets the allocator return freed pages to the driver even # expandable_segments lets the allocator return freed pages to the driver even
......
...@@ -58,11 +58,18 @@ class _EmbeddingModel: ...@@ -58,11 +58,18 @@ class _EmbeddingModel:
those branches and its VRAM is only reclaimed implicitly by gc. those branches and its VRAM is only reclaimed implicitly by gc.
""" """
__slots__ = ("backend", "model") __slots__ = ("backend", "model", "lock")
def __init__(self, backend, model): def __init__(self, backend, model):
self.backend = backend self.backend = backend
self.model = model self.model = model
# llama.cpp contexts are NOT thread-safe: two executor threads calling
# embed() on one ctx segfault the engine. Serialize per model instance.
if backend == 'llama':
import threading
self.lock = threading.Lock()
else:
self.lock = None
def __iter__(self): def __iter__(self):
yield self.backend yield self.backend
...@@ -442,10 +449,30 @@ def _embed_texts(model_obj, texts: List[str], dimensions=None) -> List[List[floa ...@@ -442,10 +449,30 @@ def _embed_texts(model_obj, texts: List[str], dimensions=None) -> List[List[floa
# those. Normalize either way (llama.cpp does not normalize). # those. Normalize either way (llama.cpp does not normalize).
import math import math
results = [] results = []
# Generous safety margin: embed() re-tokenizes the chunk text itself
# (with BOS/specials), and a detokenize→retokenize roundtrip is NOT
# token-count-stable — a window cut at exactly n_ctx can re-inflate
# past it and trip llama.cpp's GGML_ASSERT (SIGABRT).
try: try:
_tok_limit = max(16, model.n_ctx() - 8) _tok_limit = max(64, model.n_ctx() - 64)
except Exception: except Exception:
_tok_limit = 2040 _tok_limit = 1984
def _retok_len(txt):
return len(model.tokenize(txt.encode('utf-8', 'ignore'),
add_bos=True, special=True))
def _fit_chunk(toks):
"""Detokenize a token window, shrinking until its RE-tokenized
length verifiably fits the context (roundtrip can inflate)."""
end = len(toks)
while end > 0:
txt = model.detokenize(toks[:end]).decode('utf-8', 'ignore')
n = _retok_len(txt)
if n <= _tok_limit:
return txt, end
end -= max(8, n - _tok_limit)
return '', 0
def _embed_one(chunk_text): def _embed_one(chunk_text):
emb = model.embed(chunk_text) emb = model.embed(chunk_text)
...@@ -454,41 +481,52 @@ def _embed_texts(model_obj, texts: List[str], dimensions=None) -> List[List[floa ...@@ -454,41 +481,52 @@ def _embed_texts(model_obj, texts: List[str], dimensions=None) -> List[List[floa
emb = [sum(col) / n for col in zip(*emb)] emb = [sum(col) / n for col in zip(*emb)]
return emb return emb
for t in texts: # Serialize on the model's lock: llama.cpp ctxs are not thread-safe and
# An input longer than the context window would abort llama.cpp # concurrent embed() calls from parallel requests segfault the engine.
# (GGML_ASSERT) and take the whole engine down. Instead of cutting import contextlib
# the text, CHUNK it into context-sized token windows, embed each, _mlock = getattr(model_obj, 'lock', None)
# and combine with a token-count-weighted mean — no content is with (_mlock if _mlock is not None else contextlib.nullcontext()):
# dropped, and single-chunk inputs take the direct path. for t in texts:
_chunks = [(t, 1)] # An input longer than the context window would abort llama.cpp
try: # (GGML_ASSERT). Instead of cutting the text, CHUNK it into
_toks = model.tokenize(t.encode('utf-8', 'ignore'), # verified context-sized windows, embed each, and combine with a
add_bos=True, special=False) # token-count-weighted mean — no content is dropped;
if len(_toks) > _tok_limit: # single-chunk inputs take the direct path.
_chunks = [] _chunks = None
for _i in range(0, len(_toks), _tok_limit): try:
_w = _toks[_i:_i + _tok_limit] if _retok_len(t) <= _tok_limit:
_chunks.append( _chunks = [(t, 1)]
(model.detokenize(_w).decode('utf-8', 'ignore'),
len(_w)))
except Exception:
pass
if len(_chunks) == 1:
emb = _embed_one(_chunks[0][0])
else:
_acc = None
_tot = 0
for _ct, _cw in _chunks:
_e = _embed_one(_ct)
if _acc is None:
_acc = [x * _cw for x in _e]
else: else:
for _j, _x in enumerate(_e): _toks = model.tokenize(t.encode('utf-8', 'ignore'),
_acc[_j] += _x * _cw add_bos=True, special=False)
_tot += _cw _chunks = []
emb = [x / _tot for x in _acc] _i = 0
norm = math.sqrt(sum(x * x for x in emb)) or 1.0 while _i < len(_toks):
results.append([x / norm for x in emb]) _txt, _used = _fit_chunk(_toks[_i:_i + _tok_limit])
if _used <= 0:
break # pathological token — skip it
_chunks.append((_txt, _used))
_i += _used
except Exception:
_chunks = [(t, 1)]
if not _chunks:
_chunks = [(t[:2000], 1)]
if len(_chunks) == 1:
emb = _embed_one(_chunks[0][0])
else:
_acc = None
_tot = 0
for _ct, _cw in _chunks:
_e = _embed_one(_ct)
if _acc is None:
_acc = [x * _cw for x in _e]
else:
for _j, _x in enumerate(_e):
_acc[_j] += _x * _cw
_tot += _cw
emb = [x / _tot for x in _acc]
norm = math.sqrt(sum(x * x for x in emb)) or 1.0
results.append([x / norm for x in emb])
elif backend == 'vision': elif backend == 'vision':
raise ValueError( raise ValueError(
"this embedding model is image-only (no text tower) — send 'image' " "this embedding model is image-only (no text tower) — send 'image' "
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment