import os # ZeroGPU and library caches must be configured before importing spaces/torch. os.environ.setdefault("HF_HOME", os.path.expanduser("~/.cache/huggingface")) os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules") os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False") os.environ.setdefault("GRADIO_SSR_MODE", "false") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") import spaces import logging import threading import time from dataclasses import dataclass from pathlib import Path from typing import Any import gradio as gr import torch import torch.nn.functional as F from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles from gradio.data_classes import FileData from sentence_transformers import SentenceTransformer MODEL_ID = "tencent/WeMM-Embedding-9B" ROOT = Path(__file__).resolve().parent ASSET_DIR = ROOT / "assets" FRONTEND_DIR = ROOT / "frontend" MATRYOSHKA_DIMS = (64, 128, 256, 512, 1024, 2048, 4096) MAX_CUSTOM_TEXTS = 6 MAX_CUSTOM_MEDIA = 6 logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s") LOGGER = logging.getLogger("wemm-space") @dataclass(frozen=True) class Candidate: key: str title: str kind: str description: str payload: Any media_path: str | None = None SHOWCASE: tuple[Candidate, ...] = ( Candidate( "llama4", "Llama 4 model card", "visual document", "A dense model-card screenshot describing the Scout and Maverick variants.", str(ASSET_DIR / "llama4_hgf.png"), str(ASSET_DIR / "llama4_hgf.png"), ), Candidate( "qwen-omni", "Qwen2.5-Omni overview", "visual document", "A model page covering omni-modal perception, speech, and video capabilities.", str(ASSET_DIR / "qwen2.5omni_hgf.png"), str(ASSET_DIR / "qwen2.5omni_hgf.png"), ), Candidate( "likelihood-contour", "Scientific contour plot", "figure", "An orange likelihood contour plotted against gamma and log-scaled tau over mass.", str(ASSET_DIR / "doc1.jpg"), str(ASSET_DIR / "doc1.jpg"), ), Candidate( "budget-1971", "1971 budget infographic", "visual document", "A historical chart comparing US outlays, including natural resources spending.", str(ASSET_DIR / "doc2.jpg"), str(ASSET_DIR / "doc2.jpg"), ), Candidate( "scoring-rules", "Proper scoring rules paper", "visual document", "An academic page with lemmas, an algorithm, equations, and references.", str(ASSET_DIR / "doc3.jpg"), str(ASSET_DIR / "doc3.jpg"), ), Candidate( "road-safety", "Road-safety assessment", "visual document", "An environmental assessment page about driver training, signs, and road closures.", str(ASSET_DIR / "doc4.jpg"), str(ASSET_DIR / "doc4.jpg"), ), Candidate( "mapo-tofu", "Mapo tofu in motion", "video", "A short cooking clip showing the preparation of the Sichuan tofu dish.", str(ASSET_DIR / "mapo_tofu.mp4"), str(ASSET_DIR / "mapo_tofu.mp4"), ), Candidate( "zhajiang-noodles", "Zhajiang noodles in motion", "video", "A short cooking clip showing noodles with a savory fermented-bean sauce.", str(ASSET_DIR / "zhajiang_noodle.mp4"), str(ASSET_DIR / "zhajiang_noodle.mp4"), ), Candidate( "vector-search", "How vector search works", "text", "Dense retrieval maps queries and documents into one normalized vector space, then ranks candidates by cosine similarity.", "Dense retrieval maps queries and documents into one normalized vector space, then ranks candidates by cosine similarity.", ), Candidate( "night-train", "A quiet journey", "text", "夜行列车穿过雨中的城市，车窗映出霓虹灯和安静的乘客。", "夜行列车穿过雨中的城市，车窗映出霓虹灯和安静的乘客。", ), ) LOGGER.info("Loading %s on CPU", MODEL_ID) MODEL = SentenceTransformer( MODEL_ID, trust_remote_code=True, device="cpu", model_kwargs={"dtype": torch.bfloat16, "low_cpu_mem_usage": True}, ) MODEL.eval() torch.set_grad_enabled(False) LOGGER.info("Model loaded; waiting for a ZeroGPU allocation") _CACHE_LOCK = threading.Lock() _INFERENCE_LOCK = threading.Lock() _SHOWCASE_EMBEDDINGS: torch.Tensor | None = None def _file_path(value: Any) -> str | None: """Normalize Gradio file values across browser and programmatic clients.""" if value is None: return None if isinstance(value, (str, Path)): return str(value) if isinstance(value, dict): path = value.get("path") or value.get("name") return str(path) if path else None path = getattr(value, "path", None) or getattr(value, "name", None) return str(path) if path else None def _is_video(path: str) -> bool: return Path(path.split("?", 1)[0]).suffix.lower() in { ".mp4", ".webm", ".mov", ".mkv", ".avi", ".mpeg", ".mpg" } def _multimodal_payload(text: str | None, image: Any, video: Any) -> tuple[Any, str]: text = (text or "").strip() image_path = _file_path(image) video_path = _file_path(video) if image_path and video_path: raise gr.Error("Choose one visual query: an image or a video, not both.") if image_path: return ({"image": image_path, "text": text}, "image + text") if text else (image_path, "image") if video_path: return ({"video": video_path, "text": text}, "video + text") if text else (video_path, "video") if text: return text, "text" raise gr.Error("Add a text, image, or video query to begin.") def _parse_text_candidates(raw: str | None) -> list[Candidate]: candidates: list[Candidate] = [] lines = [line.strip() for line in (raw or "").splitlines() if line.strip()] if len(lines) > MAX_CUSTOM_TEXTS: raise gr.Error(f"Use at most {MAX_CUSTOM_TEXTS} custom text candidates.") for index, line in enumerate(lines, start=1): if "::" in line: title, body = (part.strip() for part in line.split("::", 1)) title = title or f"Custom text {index}" body = body or title else: title, body = f"Custom text {index}", line candidates.append(Candidate(f"custom-text-{index}", title[:80], "text", body[:240], body)) return candidates def _parse_media_candidates(raw: Any) -> list[Candidate]: candidates: list[Candidate] = [] items = raw or [] if len(items) > MAX_CUSTOM_MEDIA: raise gr.Error(f"Upload at most {MAX_CUSTOM_MEDIA} candidate media files.") for index, item in enumerate(items, start=1): path = _file_path(item) if not path: continue kind = "video" if _is_video(path) else "image" candidates.append( Candidate( f"custom-media-{index}", f"Uploaded {kind} {index}", kind, f"User-supplied {kind} candidate.", path, path, ) ) return candidates def _encode(items: list[Any], *, query: bool) -> torch.Tensor: method = MODEL.encode_query if query else MODEL.encode_document with torch.inference_mode(): embeddings = method( items, batch_size=1, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False, ) if embeddings.ndim == 1: embeddings = embeddings.unsqueeze(0) return embeddings.float().cpu() def _truncate_normalize(embeddings: torch.Tensor, dimension: int) -> torch.Tensor: return F.normalize(embeddings[..., :dimension], p=2, dim=-1) def _dimension_scores(query: torch.Tensor, documents: torch.Tensor) -> dict[int, list[float]]: scores: dict[int, list[float]] = {} for dimension in MATRYOSHKA_DIMS: query_d = _truncate_normalize(query, dimension) docs_d = _truncate_normalize(documents, dimension) scores[dimension] = (query_d @ docs_d.T).squeeze(0).tolist() return scores def _fingerprint(vector: torch.Tensor, bar_count: int = 96) -> list[float]: values = vector.detach().float().flatten() chunks = torch.tensor_split(values, min(bar_count, values.numel())) return [round(float(chunk.mean()), 7) for chunk in chunks] def _public_media(candidate: Candidate) -> str | None: if not candidate.media_path: return None try: relative = Path(candidate.media_path).resolve().relative_to(ASSET_DIR.resolve()) except ValueError: return None return f"/assets/{relative.as_posix()}" def _search_duration(*args: Any, **kwargs: Any) -> int: """Budget more time for the cold corpus pass and video queries.""" query_payload = args[0] if args else None custom_candidates = args[2] if len(args) > 2 else [] include_showcase = bool(args[3]) if len(args) > 3 else True query_has_video = ( isinstance(query_payload, dict) and bool(query_payload.get("video")) ) or (isinstance(query_payload, str) and _is_video(query_payload)) has_uploaded_video = any(item.kind == "video" for item in custom_candidates or []) if include_showcase and _SHOWCASE_EMBEDDINGS is None: return 30 return 30 if query_has_video or has_uploaded_video else 15 @spaces.GPU(duration=_search_duration) def _run_search( query_payload: Any, query_kind: str, custom_candidates: list[Candidate], include_showcase: bool, dimension: int, ) -> tuple[torch.Tensor, torch.Tensor, list[Candidate], float, bool]: global _SHOWCASE_EMBEDDINGS started = time.perf_counter() was_cold = include_showcase and _SHOWCASE_EMBEDDINGS is None MODEL.to("cuda") try: query_embedding = _encode([query_payload], query=True) document_blocks: list[torch.Tensor] = [] candidates: list[Candidate] = [] if include_showcase: with _CACHE_LOCK: cached = _SHOWCASE_EMBEDDINGS if cached is None: encoded = _encode([item.payload for item in SHOWCASE], query=False) with _CACHE_LOCK: if _SHOWCASE_EMBEDDINGS is None: _SHOWCASE_EMBEDDINGS = encoded cached = _SHOWCASE_EMBEDDINGS document_blocks.append(cached) candidates.extend(SHOWCASE) if custom_candidates: custom_embeddings = _encode([item.payload for item in custom_candidates], query=False) document_blocks.append(custom_embeddings) candidates.extend(custom_candidates) if not document_blocks: raise gr.Error("Include the showcase universe or add at least one candidate.") documents = torch.cat(document_blocks, dim=0) return query_embedding, documents, candidates, time.perf_counter() - started, was_cold finally: MODEL.to("cpu") if torch.cuda.is_available(): torch.cuda.empty_cache() def _pair_duration(*args: Any, **kwargs: Any) -> int: has_video = any( (isinstance(value, dict) and bool(value.get("video"))) or (isinstance(value, str) and _is_video(value)) for value in args[:2] ) return 30 if has_video else 15 @spaces.GPU(duration=_pair_duration) def _run_pair( query_payload: Any, candidate_payload: Any, ) -> tuple[torch.Tensor, torch.Tensor, float]: started = time.perf_counter() MODEL.to("cuda") try: query = _encode([query_payload], query=True) candidate = _encode([candidate_payload], query=False) return query, candidate, time.perf_counter() - started finally: MODEL.to("cpu") if torch.cuda.is_available(): torch.cuda.empty_cache() def _interpret_score(score: float) -> tuple[str, str]: if score >= 0.75: return "high alignment", "These inputs occupy a very similar region for this retrieval model." if score >= 0.50: return "meaningful alignment", "The model sees a substantial semantic relationship." if score >= 0.25: return "weak alignment", "There is some overlap, but stronger candidates may rank above it." return "low alignment", "The model places these inputs relatively far apart." app = gr.Server( title="WeMM Semantic Universe API", summary="Custom multimodal retrieval studio powered by tencent/WeMM-Embedding-9B", version="2.0.0", ) @app.api( name="search", description="Rank a mixed text/image/video collection using WeMM-Embedding-9B.", concurrency_limit=1, concurrency_id="wemm_gpu", time_limit=300, ) def search_api( query_text: str = "", query_image: FileData | None = None, query_video: FileData | None = None, custom_texts: str = "", candidate_media: list[FileData] | None = None, include_showcase: bool = True, dimension: int = 1024, ) -> dict[str, Any]: """Search a mixed media collection and return structured visualization data.""" dimension = int(dimension) if dimension not in MATRYOSHKA_DIMS: raise gr.Error("Choose one of the model's native Matryoshka dimensions.") query_payload, query_kind = _multimodal_payload(query_text, query_image, query_video) custom_candidates = _parse_text_candidates(custom_texts) + _parse_media_candidates(candidate_media) with _INFERENCE_LOCK: query, documents, candidates, elapsed, was_cold = _run_search( query_payload, query_kind, custom_candidates, bool(include_showcase), dimension, ) dimension_map = _dimension_scores(query, documents) chosen_scores = dimension_map[dimension] order = sorted(range(len(candidates)), key=lambda index: chosen_scores[index], reverse=True) ranked = [ { "rank": rank, "key": candidates[index].key, "title": candidates[index].title, "kind": candidates[index].kind, "description": candidates[index].description, "score": round(float(chosen_scores[index]), 7), "media_url": _public_media(candidates[index]), } for rank, index in enumerate(order, start=1) ] curve_series = [ { "key": candidates[index].key, "title": candidates[index].title, "values": [round(float(dimension_map[dim][index]), 7) for dim in MATRYOSHKA_DIMS], } for index in order[:4] ] return { "mode": "search", "model": MODEL_ID, "query_kind": query_kind, "dimension": dimension, "dimensions": list(MATRYOSHKA_DIMS), "candidate_count": len(candidates), "elapsed_seconds": round(elapsed, 3), "cache_state": "created" if was_cold else "reused" if include_showcase else "not requested", "top_match": ranked[0], "rankings": ranked, "dimension_series": curve_series, "query_fingerprint": _fingerprint(query[0, :dimension]), "full_embedding_dimension": int(query.shape[-1]), "query_vector_preview": [round(float(value), 7) for value in query[0, :12]], "l2_norm_after_truncation": round(float(_truncate_normalize(query, dimension).norm()), 7), "note": "Cosine similarity is a ranking signal, not a calibrated probability.", } @app.api( name="compare", description="Compare a multimodal query and candidate across every native Matryoshka dimension.", concurrency_limit=1, concurrency_id="wemm_gpu", time_limit=240, ) def compare_api( query_text: str = "", query_image: FileData | None = None, query_video: FileData | None = None, candidate_text: str = "", candidate_image: FileData | None = None, candidate_video: FileData | None = None, dimension: int = 1024, ) -> dict[str, Any]: """Compare two supported inputs and return structured visualization data.""" dimension = int(dimension) if dimension not in MATRYOSHKA_DIMS: raise gr.Error("Choose one of the model's native Matryoshka dimensions.") query_payload, query_kind = _multimodal_payload(query_text, query_image, query_video) candidate_payload, candidate_kind = _multimodal_payload( candidate_text, candidate_image, candidate_video ) with _INFERENCE_LOCK: query, candidate, elapsed = _run_pair(query_payload, candidate_payload) scores = { dim: float((_truncate_normalize(query, dim) @ _truncate_normalize(candidate, dim).T).item()) for dim in MATRYOSHKA_DIMS } selected_score = scores[dimension] label, explanation = _interpret_score(selected_score) return { "mode": "compare", "model": MODEL_ID, "query_kind": query_kind, "candidate_kind": candidate_kind, "dimension": dimension, "dimensions": list(MATRYOSHKA_DIMS), "selected_score": round(selected_score, 7), "label": label, "explanation": explanation, "scores": [round(scores[dim], 7) for dim in MATRYOSHKA_DIMS], "query_fingerprint": _fingerprint(query[0, :dimension]), "candidate_fingerprint": _fingerprint(candidate[0, :dimension]), "elapsed_seconds": round(elapsed, 3), "note": "Interpret thresholds relative to a task-specific candidate set.", } @app.get("/", response_class=HTMLResponse, include_in_schema=False) async def homepage() -> HTMLResponse: return HTMLResponse((FRONTEND_DIR / "index.html").read_text(encoding="utf-8")) @app.get("/health", include_in_schema=False) async def health() -> dict[str, str]: return {"status": "ok", "model": MODEL_ID, "interface": "gradio-server"} app.mount("/assets", StaticFiles(directory=str(ASSET_DIR)), name="assets") app.mount("/ui", StaticFiles(directory=str(FRONTEND_DIR)), name="ui") # Hugging Face Spaces expects the application object under this conventional name. demo = app if __name__ == "__main__": app.launch(show_error=True)