Spaces:
Running
Running
File size: 6,793 Bytes
d289270 c6afbdf d289270 c6afbdf d289270 c6afbdf 0458600 d289270 3adb6fd d289270 c6afbdf ccc33b7 c6afbdf d289270 c6afbdf d289270 c6afbdf d289270 ccc33b7 d289270 3fe3157 d289270 c6afbdf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """True token-budgeted batched inference with resident models.
The external service-wide inference gate is held by ModelRegistry for the full call.
No model initialization, truncation, or extra E5 normalization occurs here.
"""
from __future__ import annotations
import logging
from app.embedding import EmbeddingResult, TooManyTokens
from app.models import ColbertResult, validate_vector, INSTRUCTION
log = logging.getLogger(__name__)
def _split_by_padding_cost(counts, token_budget, max_size):
"""Partition ordered inputs, accounting for batch padding to longest text."""
groups = []
current = []
current_max = 0
for index, count in enumerate(counts):
proposed_max = max(current_max, count)
if current and (len(current) >= max_size or
proposed_max * (len(current) + 1) > token_budget):
groups.append(current)
current = []
current_max = 0
current.append(index)
current_max = max(current_max, count)
if current:
groups.append(current)
return groups
def _token_count(tokenizer, text):
encoded = tokenizer(text, add_special_tokens=True, truncation=False)
return len(encoded["input_ids"])
def _dense_batch(engine, texts, kind, is_qwen, raw=False, truncate_long=False):
if raw:
# OpenAI-compatible callers provide preformatted text; never add prompts.
strings = list(texts)
budget, max_size = (4096, 16) if is_qwen else (8192, 64)
elif is_qwen:
strings = [
f"Instruct: {INSTRUCTION}\nQuery:{text}" if kind == "query" else text
for text in texts
]
# Smaller padded batches limit CPU activation memory; single long
# documents are still allowed up to the model's advertised context.
budget, max_size = 4096, 16
else:
strings = [f"{kind}: {text}" for text in texts]
budget, max_size = 8192, 64
counts = [_token_count(engine.tokenizer, text) for text in strings]
for index, count in enumerate(counts):
if count > engine.max_tokens:
if not truncate_long:
raise TooManyTokens(count, engine.max_tokens, input_index=index)
log.warning(
"E5 embeddings truncated: input_index=%d "
"original_tokens=%d used_tokens=%d",
index, count, engine.max_tokens,
)
counts[index] = engine.max_tokens
# Bound total work in one HTTP request (not the model context limit).
if sum(counts) > 65536:
raise ValueError("Batch exceeds 65536 total input tokens")
results = [None] * len(texts)
for group in _split_by_padding_cost(counts, budget, max_size):
inputs = [strings[i] for i in group]
tokenizer_options = {
"return_tensors": "pt",
"padding": True,
"truncation": truncate_long,
"add_special_tokens": True,
}
if truncate_long:
# Hugging Face preserves special tokens inside max_length and
# truncates from the right, keeping the leading query:/passage: prefix.
tokenizer_options["max_length"] = engine.max_tokens
batch = engine.tokenizer(inputs, **tokenizer_options)
with engine.torch.inference_mode():
hidden = engine.model(**batch).last_hidden_state
if is_qwen:
# Qwen tokenizer is explicitly left padded. The final index
# therefore refers to the final real token in every row.
vectors = engine.torch.nn.functional.normalize(
hidden[:, -1, :], p=2, dim=1,
)
else:
mask = batch["attention_mask"].unsqueeze(-1)
vectors = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
vectors = vectors.float().cpu().tolist()
for index, vector in zip(group, vectors):
validate_vector(vector, engine.dimensions)
results[index] = EmbeddingResult(
tokens_count=counts[index], embedding=vector,
)
return results
def _colbert_batch(engine, texts, kind):
model = engine.model
query = kind == "query"
prefix = model.query_prefix if query else model.document_prefix
counts = [_token_count(model.tokenizer, prefix + text) for text in texts]
for count in counts:
if count > engine.max_tokens:
raise TooManyTokens(count, engine.max_tokens)
if sum(counts) > 32768:
raise ValueError("ColBERT request exceeds 32768 total input tokens")
results = [None] * len(texts)
for group in _split_by_padding_cost(counts, 4096, 8):
if query:
model.query_length = max(32, max(counts[i] for i in group))
outputs = model.encode(
[texts[i] for i in group],
is_query=query, batch_size=len(group), show_progress_bar=False,
)
if len(outputs) != len(group):
raise RuntimeError("Unexpected PyLate batch length")
for index, matrix in zip(group, outputs):
if hasattr(matrix, "detach"):
matrix = matrix.detach().float().cpu()
if hasattr(matrix, "tolist"):
matrix = matrix.tolist()
vectors = [[float(value) for value in row] for row in matrix]
if not vectors or len(vectors) > engine.max_tokens:
raise RuntimeError("Unexpected PyLate output token count")
for vector in vectors:
validate_vector(vector, engine.dimensions)
# Note: encoded tokens and emitted token vectors may differ
# because PyLate can filter special/padding tokens.
results[index] = ColbertResult(
tokens_count=counts[index], token_embeddings=vectors,
)
return results
def batch_encode(engine, texts, kind, is_qwen=False, is_colbert=False,
raw=False, truncate_long=False):
"""Native tensor batch for real models, single-item fallback for test doubles."""
if truncate_long and (is_qwen or is_colbert):
raise ValueError("Automatic truncation is only supported for E5 input")
if is_colbert:
if not hasattr(engine, "model") or not hasattr(engine.model, "tokenizer"):
return [engine.encode(text, kind) for text in texts]
return _colbert_batch(engine, texts, kind)
if not hasattr(engine, "tokenizer") or not hasattr(engine, "model"):
if raw:
raise RuntimeError("Raw embeddings require the loaded model and tokenizer")
return [engine.encode(text, kind) for text in texts]
return _dense_batch(engine, texts, kind, is_qwen, raw=raw,
truncate_long=truncate_long)
|