Spaces:
Running
Running
| """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) | |