Spaces:
Running
Running
| """Resident E5, Qwen3, PyLate mLateOn and BGE-M3. No per-request model loading.""" | |
| import gc | |
| import math | |
| import os | |
| import resource | |
| import threading | |
| from pathlib import Path | |
| from dataclasses import dataclass | |
| from app.embedding import EmbeddingEngine, EmbeddingResult, MODEL_ID as E5_ID, TooManyTokens | |
| QWEN_ID = "Qwen/Qwen3-Embedding-0.6B" | |
| COLBERT_ID = "lightonai/mLateOn" | |
| BGE_M3_ID = "BAAI/bge-m3" | |
| QWEN_MAX = 32768 | |
| COLBERT_MAX = 8192 | |
| INSTRUCTION = "Given a question, retrieve relevant passages that answer the question" | |
| class ServiceBusy(Exception): | |
| pass | |
| class ModelUnavailable(Exception): | |
| """Requested model was intentionally not loaded at startup.""" | |
| MODEL_NAMES = ("e5", "qwen", "colbert", "bge-m3") | |
| def enabled_models_from_env(): | |
| """Parse the runtime load allowlist; omitted means all four models.""" | |
| raw = os.getenv("ENABLED_MODELS") | |
| if raw is None: | |
| enabled = set(MODEL_NAMES) | |
| # Retain the earlier BGE-only toggle for existing deployments. | |
| if os.getenv("BGE_M3_ENABLED", "1") != "1": | |
| enabled.remove("bge-m3") | |
| return frozenset(enabled) | |
| aliases = { | |
| "e5": "e5", E5_ID: "e5", | |
| "qwen": "qwen", QWEN_ID: "qwen", | |
| "colbert": "colbert", "mlateon": "colbert", COLBERT_ID: "colbert", | |
| "bge-m3": "bge-m3", "bge_m3": "bge-m3", BGE_M3_ID: "bge-m3", | |
| } | |
| parts = [part.strip() for part in raw.split(",")] | |
| if not parts or any(not part for part in parts): | |
| raise ValueError("ENABLED_MODELS must list at least one model: " | |
| + ", ".join(MODEL_NAMES)) | |
| unknown = [part for part in parts if part not in aliases] | |
| if unknown: | |
| raise ValueError("Unknown ENABLED_MODELS entries: " + ", ".join(unknown)) | |
| mapped = [aliases[part] for part in parts] | |
| if len(set(mapped)) != len(mapped): | |
| raise ValueError("Duplicate models in ENABLED_MODELS") | |
| return frozenset(mapped) | |
| def memory_mb(): | |
| current = None | |
| try: | |
| with open("/proc/self/status", encoding="utf-8") as f: | |
| for line in f: | |
| if line.startswith("VmRSS:"): | |
| current = round(int(line.split()[1]) / 1024, 1) | |
| break | |
| except OSError: | |
| pass | |
| peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss | |
| return {"rss_mb": current, "peak_rss_mb": round(peak / 1024, 1), | |
| "budget_mb": int(os.getenv("MEMORY_BUDGET_MB", "15360"))} | |
| def validate_vector(vector, dimensions): | |
| if len(vector) != dimensions or any(not math.isfinite(v) for v in vector): | |
| raise RuntimeError("Invalid embedding vector") | |
| class QwenEngine: | |
| dimensions = 1024 | |
| max_tokens = QWEN_MAX | |
| def __init__(self, tokenizer, model, torch): | |
| self.tokenizer, self.model, self.torch = tokenizer, model, torch | |
| def load(cls): | |
| import torch | |
| from transformers import AutoModel, AutoTokenizer | |
| path = os.getenv("QWEN_MODEL_DIR", "/opt/models/qwen3-embedding-0.6b") | |
| tokenizer = AutoTokenizer.from_pretrained(path, padding_side="left", | |
| local_files_only=True) | |
| model = AutoModel.from_pretrained(path, torch_dtype=torch.float32, | |
| attn_implementation="sdpa", | |
| local_files_only=True).to("cpu").eval() | |
| if int(model.config.hidden_size) != 1024: | |
| raise RuntimeError("Incorrect Qwen output size") | |
| return cls(tokenizer, model, torch) | |
| def encode(self, text, kind): | |
| inp = (f"Instruct: {INSTRUCTION}\nQuery:{text}" | |
| if kind == "query" else text) | |
| tokens = self.tokenizer(inp, add_special_tokens=True, truncation=False, | |
| return_tensors="pt") | |
| count = int(tokens["attention_mask"].sum().item()) | |
| if count > self.max_tokens: | |
| raise TooManyTokens(count, self.max_tokens) | |
| with self.torch.inference_mode(): | |
| last = self.model(**tokens).last_hidden_state[:, -1, :] | |
| normalized = self.torch.nn.functional.normalize(last, p=2, dim=1) | |
| values = normalized[0].float().cpu().tolist() | |
| validate_vector(values, self.dimensions) | |
| return EmbeddingResult(tokens_count=count, embedding=values) | |
| class ColbertResult: | |
| tokens_count: int | |
| token_embeddings: list[list[float]] | |
| class ColbertEngine: | |
| dimensions = 128 | |
| max_tokens = COLBERT_MAX | |
| def __init__(self, model): | |
| self.model = model | |
| def load(cls): | |
| from pylate import models | |
| model = models.ColBERT( | |
| model_name_or_path=os.getenv("COLBERT_MODEL_DIR", "/opt/models/mlateon"), | |
| device="cpu", query_length=32, document_length=COLBERT_MAX, | |
| ) | |
| model.eval() | |
| model.document_length = COLBERT_MAX | |
| return cls(model) | |
| def encode(self, text, kind): | |
| is_query = kind == "query" | |
| prefix = self.model.query_prefix if is_query else self.model.document_prefix | |
| tokens = self.model.tokenizer(f"{prefix}{text}", add_special_tokens=True, | |
| truncation=False) | |
| count = len(tokens["input_ids"]) | |
| if count > self.max_tokens: | |
| raise TooManyTokens(count, self.max_tokens) | |
| if is_query: | |
| # PyLate pads queries to query_length. Never pad short queries to 8192. | |
| self.model.query_length = max(32, count) | |
| output = self.model.encode([text], is_query=is_query, batch_size=1, | |
| show_progress_bar=False) | |
| if len(output) != 1: | |
| raise RuntimeError("Unexpected PyLate batch size") | |
| matrix = output[0] | |
| if hasattr(matrix, "detach"): | |
| matrix = matrix.detach().float().cpu() | |
| if hasattr(matrix, "tolist"): | |
| matrix = matrix.tolist() | |
| vectors = [[float(x) for x in row] for row in matrix] | |
| if not vectors or len(vectors) > COLBERT_MAX: | |
| raise RuntimeError("Unexpected PyLate output length") | |
| for row in vectors: | |
| validate_vector(row, self.dimensions) | |
| return ColbertResult(tokens_count=len(vectors), token_embeddings=vectors) | |
| class ModelRegistry: | |
| def __init__(self, e5, qwen, colbert, startup_memory=None, bge_m3=None): | |
| self.e5, self.qwen, self.colbert = e5, qwen, colbert | |
| self.bge_m3 = bge_m3 | |
| self.startup_memory = startup_memory or {} | |
| # Only one CPU inference may execute at a time. Allow a bounded | |
| # number of other HTTP calls to wait instead of returning 429 instantly. | |
| # Dify's embeddings provider has a 60-second HTTP timeout. | |
| pending = max(0, int(os.getenv("INFERENCE_MAX_WAITERS", "16"))) | |
| self._waiting_slots = threading.BoundedSemaphore(pending + 1) | |
| self._gate = threading.BoundedSemaphore(1) | |
| def load(cls): | |
| import torch | |
| torch.set_num_threads(max(1, min(4, int(os.getenv("CPU_THREADS", "2"))))) | |
| enabled = enabled_models_from_env() | |
| readings = {} | |
| e5 = qwen = colbert = bge_m3 = None | |
| if "e5" in enabled: | |
| e5 = EmbeddingEngine.load() | |
| readings["e5"] = memory_mb() | |
| if "qwen" in enabled: | |
| qwen = QwenEngine.load() | |
| readings["qwen"] = memory_mb() | |
| if "colbert" in enabled: | |
| colbert = ColbertEngine.load() | |
| readings["colbert"] = memory_mb() | |
| if "bge-m3" in enabled: | |
| from app.bge_m3 import BgeM3Engine | |
| bge_m3 = BgeM3Engine.load() | |
| readings["bge_m3"] = memory_mb() | |
| gc.collect() | |
| ram = memory_mb() | |
| if ram["rss_mb"] is not None and ram["rss_mb"] > ram["budget_mb"]: | |
| raise MemoryError("All resident models exceed RAM budget") | |
| return cls(e5, qwen, colbert, readings, bge_m3=bge_m3) | |
| def loaded_models(self): | |
| return [name for name, engine in ( | |
| ("e5", self.e5), ("qwen", self.qwen), | |
| ("colbert", self.colbert), ("bge-m3", self.bge_m3), | |
| ) if engine is not None] | |
| def _run(self, callback): | |
| # Keep inference serialized to protect CPU RAM, but give parallel | |
| # clients time to queue. The queue is bounded to avoid thread exhaustion. | |
| if not self._waiting_slots.acquire(blocking=False): | |
| raise ServiceBusy("Inference queue is full; retry later") | |
| acquired = False | |
| try: | |
| timeout = max(0.0, float(os.getenv("INFERENCE_WAIT_SECONDS", "40"))) | |
| acquired = self._gate.acquire(timeout=timeout) | |
| if not acquired: | |
| raise ServiceBusy("Inference queue wait timed out; retry later") | |
| return callback() | |
| finally: | |
| if acquired: | |
| self._gate.release() | |
| self._waiting_slots.release() | |
| def embed(self, text, kind, selected): | |
| engine = self.qwen if selected in ("qwen", QWEN_ID) else self.e5 | |
| if engine is None: | |
| raise ModelUnavailable(f"Model {selected or 'e5'} is not enabled") | |
| result = self._run(lambda: engine.encode(text, kind)) | |
| return (QWEN_ID if engine is self.qwen else E5_ID), engine.dimensions, result | |
| def openai_encode(self, texts, canonical_model): | |
| """OpenAI-compatible output is always dense only.""" | |
| from app.batching import batch_encode | |
| from app.bge_m3 import BGE_M3_ID | |
| if canonical_model == BGE_M3_ID: | |
| if self.bge_m3 is None: | |
| raise ModelUnavailable("Model bge-m3 is not enabled") | |
| def encode_bge_dense(): | |
| engine = self.bge_m3 | |
| counts = [ | |
| len(engine.model.tokenizer(text, add_special_tokens=True, | |
| truncation=False)["input_ids"]) | |
| for text in texts | |
| ] | |
| for index, count in enumerate(counts): | |
| if count > engine.max_tokens: | |
| raise TooManyTokens(count, engine.max_tokens, | |
| input_index=index) | |
| if sum(counts) > 65536: | |
| raise ValueError("Batch exceeds 65536 total input tokens") | |
| results = [] | |
| group, group_tokens = [], 0 | |
| # Dify may send 16+ texts; internal native BGE max is 8. | |
| for text, count in zip(texts, counts): | |
| if group and (len(group) == 8 or group_tokens + count > 8192): | |
| results.extend(engine.encode( | |
| group, return_sparse=False, return_colbert=False)) | |
| group, group_tokens = [], 0 | |
| group.append(text) | |
| group_tokens += count | |
| if group: | |
| results.extend(engine.encode( | |
| group, return_sparse=False, return_colbert=False)) | |
| return [EmbeddingResult( | |
| tokens_count=item.tokens_count, embedding=item.dense) | |
| for item in results] | |
| return self._run(encode_bge_dense) | |
| if canonical_model == E5_ID: | |
| engine, is_qwen = self.e5, False | |
| elif canonical_model == QWEN_ID: | |
| engine, is_qwen = self.qwen, True | |
| else: | |
| raise ValueError("OpenAI embeddings only support dense E5 and Qwen models") | |
| if engine is None: | |
| raise ModelUnavailable(f"Model {canonical_model} is not enabled") | |
| return self._run(lambda: batch_encode( | |
| engine, texts, "passage", is_qwen=is_qwen, raw=True, | |
| truncate_long=(canonical_model == E5_ID), | |
| )) | |
| def bge_m3_encode(self, texts, return_colbert=False): | |
| if self.bge_m3 is None: | |
| raise ModelUnavailable("Model bge-m3 is not enabled") | |
| return self._run(lambda: self.bge_m3.encode( | |
| texts, return_colbert=return_colbert, | |
| )) | |
| def colbert_encode(self, text, kind): | |
| if self.colbert is None: | |
| raise ModelUnavailable("Model colbert is not enabled") | |
| return self._run(lambda: self.colbert.encode(text, kind)) | |
| def voyage_encode(self, texts, kind, canonical_model): | |
| """Atomic inference gate protects all microbatches in one HTTP request.""" | |
| from app.batching import batch_encode | |
| if canonical_model == E5_ID: | |
| # Same 512-token safety fallback as the OpenAI E5 endpoint. | |
| # _dense_batch adds query:/passage: before truncating the tail. | |
| engine, mode = self.e5, {"truncate_long": True} | |
| elif canonical_model == QWEN_ID: | |
| engine, mode = self.qwen, {"is_qwen": True} | |
| elif canonical_model == COLBERT_ID: | |
| engine, mode = self.colbert, {"is_colbert": True} | |
| else: | |
| raise ValueError("Unknown model") | |
| if engine is None: | |
| raise ModelUnavailable(f"Model {canonical_model} is not enabled") | |
| return self._run(lambda: batch_encode(engine, texts, kind, **mode)) | |