"""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)