diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index 2e79ae5d22..cc34e7c9a7 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -1,17 +1,25 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Figure captioning via the user's currently-loaded chat VLM. +"""Figure captioning for RAG ingestion. -No separate vision model is loaded. If the chat model is vision-capable -(detected at ingestion-enqueue time and passed in as ``vlm_url`` / -``vlm_model``), we call its OpenAI-compatible ``/v1/chat/completions`` -with the figure as a base64 ``image_url``. If no vision-capable chat -model is loaded, captioning is skipped entirely — the ingestion falls -back to the parser's page-text ``nearest_caption``. +Two captioning sources, tried in order: -Defensive: any per-image request failure returns an empty string so -ingestion stays resilient. +1. The user's currently-loaded chat VLM — when ``vlm_url`` / ``vlm_model`` + are provided by the parent process (it probes ``llama_cpp.is_vision`` + at enqueue time). Captions go through that model's OpenAI-compatible + ``/v1/chat/completions`` endpoint as base64 ``image_url``. + +2. A helper llama-server fallback that loads the pre-cached + ``unsloth/gemma-4-E2B-it-GGUF`` (gemma-3n family, multimodal) with + its mmproj for vision. Spawned for the lifetime of a + ``caption_images`` call, unloaded before return so no llama-server + process leaks past ingestion. + +Defensive: any per-image failure returns an empty string; total +captioner unavailability (no chat VLM + helper load failure) returns +empty strings for every image. The caller (``_stream_image_chunks``) +falls back to the parser's page-text caption in that case. """ from __future__ import annotations @@ -19,7 +27,7 @@ from __future__ import annotations import base64 import logging from io import BytesIO -from typing import Optional +from typing import Any, Optional logger = logging.getLogger(__name__) @@ -36,6 +44,13 @@ _MAX_NEW_TOKENS = 120 _MAX_IMAGE_SIZE = 1600 _REQUEST_TIMEOUT_SECONDS = 120.0 +# Helper VLM (used when no vision-capable chat model is loaded). +# Matches the model pre-cached by precache_helper_gguf() at studio +# startup so the captioner doesn't have to wait on a fresh download. +_HELPER_REPO = "unsloth/gemma-4-E2B-it-GGUF" +_HELPER_VARIANT = "UD-Q4_K_XL" +_HELPER_MODEL_NAME = "helper" + def _image_to_data_url(blob: bytes) -> str: from PIL import Image @@ -49,61 +64,119 @@ def _image_to_data_url(blob: bytes) -> str: return f"data:image/jpeg;base64,{encoded}" +def _load_helper_vlm() -> Optional[tuple[Any, str, str]]: + """Spawn a private LlamaCppBackend with the helper VLM + mmproj. + + Returns ``(backend, base_url, model_name)`` on success, ``None`` on + failure. The caller is responsible for unloading the backend when + done (so the helper doesn't outlive the ingestion subprocess). + """ + try: + from core.inference.llama_cpp import LlamaCppBackend + + backend = LlamaCppBackend() + logger.info( + "RAG captioner: loading helper VLM %s (%s) as fallback", + _HELPER_REPO, + _HELPER_VARIANT, + ) + ok = backend.load_model( + hf_repo = _HELPER_REPO, + hf_variant = _HELPER_VARIANT, + model_identifier = f"rag-captioner:{_HELPER_REPO}:{_HELPER_VARIANT}", + is_vision = True, + n_ctx = 4096, + n_gpu_layers = -1, + ) + if not ok: + logger.warning("RAG captioner: helper VLM failed to start") + return None + return backend, backend.base_url, _HELPER_MODEL_NAME + except Exception as exc: # noqa: BLE001 + logger.warning("RAG captioner: helper VLM load raised: %s", exc) + return None + + +def _post_one(client: Any, endpoint: str, model: str, blob: bytes) -> str: + """POST one image to the OpenAI-compatible endpoint, return caption.""" + data_url = _image_to_data_url(blob) + payload = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": _PROMPT}, + { + "type": "image_url", + "image_url": {"url": data_url}, + }, + ], + } + ], + "max_tokens": _MAX_NEW_TOKENS, + "temperature": 0.0, + } + response = client.post(endpoint, json = payload) + response.raise_for_status() + data = response.json() + content = data.get("choices", [{}])[0].get("message", {}).get("content", "") + return content.strip() if isinstance(content, str) else "" + + def caption_images( image_bytes_list: list[bytes], *, vlm_url: Optional[str] = None, vlm_model: Optional[str] = None, ) -> list[str]: - """Caption each image via the loaded chat VLM. + """Generate one short caption per image; same-length output. - ``vlm_url`` / ``vlm_model`` come from the parent's chat-backend probe. - When either is missing (no model loaded, or loaded model is text-only), - returns an empty string per image so the caller falls back to its - parser-provided caption. Never raises. + Tries the loaded chat VLM first (``vlm_url`` + ``vlm_model``). If + those are missing, spawns the helper VLM, captions, and unloads it + before returning. On any failure returns ``""`` for the affected + image. Never raises. """ if not image_bytes_list: return [] - if not vlm_url or not vlm_model: - return ["" for _ in image_bytes_list] import httpx - endpoint = f"{vlm_url.rstrip('/')}/v1/chat/completions" - out: list[str] = [] - with httpx.Client(timeout = _REQUEST_TIMEOUT_SECONDS) as client: - for blob in image_bytes_list: + helper_backend: Optional[Any] = None + try: + # Resolve endpoint + model: chat VLM if available, else helper. + if vlm_url and vlm_model: + endpoint = f"{vlm_url.rstrip('/')}/v1/chat/completions" + model_name = vlm_model + else: + loaded = _load_helper_vlm() + if loaded is None: + # No chat VLM and helper failed → all empty strings; + # caller falls back to page-text captions. + return ["" for _ in image_bytes_list] + helper_backend, helper_base_url, helper_model_name = loaded + endpoint = f"{helper_base_url.rstrip('/')}/v1/chat/completions" + model_name = helper_model_name + + out: list[str] = [] + with httpx.Client(timeout = _REQUEST_TIMEOUT_SECONDS) as client: + for blob in image_bytes_list: + try: + out.append(_post_one(client, endpoint, model_name, blob)) + except Exception as exc: # noqa: BLE001 + logger.warning( + "caption_images: per-image request to %s failed: %s", + endpoint, + exc, + ) + out.append("") + return out + finally: + # Always tear down the helper if we spawned one. Chat VLM (when + # provided by the parent) is left alone — it's not ours to manage. + if helper_backend is not None: try: - data_url = _image_to_data_url(blob) - payload = { - "model": vlm_model, - "messages": [ - { - "role": "user", - "content": [ - {"type": "text", "text": _PROMPT}, - { - "type": "image_url", - "image_url": {"url": data_url}, - }, - ], - } - ], - "max_tokens": _MAX_NEW_TOKENS, - "temperature": 0.0, - } - response = client.post(endpoint, json = payload) - response.raise_for_status() - data = response.json() - content = ( - data.get("choices", [{}])[0].get("message", {}).get("content", "") - ) - out.append(content.strip() if isinstance(content, str) else "") + helper_backend.unload_model() + logger.info("RAG captioner: helper VLM unloaded") except Exception as exc: # noqa: BLE001 - logger.warning( - "caption_images: per-image request to %s failed: %s", - endpoint, - exc, - ) - out.append("") - return out + logger.warning("RAG captioner: helper unload failed: %s", exc) diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index df33c488ae..3f460617ab 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -64,11 +64,15 @@ def _subprocess_worker( vlm_model: str | None = None, ) -> None: try: + from core.rag.captioner import caption_images from core.rag.chunking import chunk_pages - from core.rag.parsers import parse + from core.rag.parsers import inline_image_captions, parse out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) - parsed = parse(Path(stored_path), want_images = (mode == "multimodal")) + # Always extract images so we can caption + splice for both + # modes. Text mode uses the captions inline in markdown; multimodal + # additionally embeds the raw images as image-kind chunks. + parsed = parse(Path(stored_path), want_images = True) pages = parsed.pages if not pages and not parsed.images: out_queue.put( @@ -76,6 +80,23 @@ def _subprocess_worker( ) return + # Caption figures once (chat VLM if available, else helper VLM + # fallback), then splice captions into the page markdown so the + # chunker indexes them like any other text. Multimodal mode also + # passes these same captions through to _stream_image_chunks + # below — no duplicate VLM calls per image. + captions: list[str] = [] + if parsed.images: + out_queue.put( + {"type": "progress", "stage": "caption_images", "progress": 0.08} + ) + captions = caption_images( + [img.image_bytes for img in parsed.images], + vlm_url = vlm_url, + vlm_model = vlm_model, + ) + pages = inline_image_captions(pages, parsed.images, captions) + out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1}) from core.rag.embeddings import ( get_embedder, @@ -119,8 +140,7 @@ def _subprocess_worker( model_name = model_name, out_queue = out_queue, first_index = text_count, - vlm_url = vlm_url, - vlm_model = vlm_model, + precomputed_captions = captions, ) out_queue.put({"type": "complete", "num_chunks": text_count + image_count}) except Exception as exc: # noqa: BLE001 @@ -194,11 +214,15 @@ def _stream_image_chunks( model_name: str, out_queue, first_index: int, - vlm_url: str | None = None, - vlm_model: str | None = None, + precomputed_captions: list[str] | None = None, ) -> int: - """Persist images, emit image+caption chunks; pairs share pair_group.""" - from core.rag.captioner import caption_images + """Persist images, emit image+caption chunks; pairs share pair_group. + + ``precomputed_captions`` come from the parent's earlier + caption_images call (used so we don't VLM-caption the same images + twice — once for markdown splicing, once for the caption-kind chunk). + If absent we fall back to each image's nearest_caption (page text). + """ from core.rag.embeddings import encode, encode_images from utils.paths.storage_roots import ensure_dir, rag_uploads_root @@ -211,8 +235,9 @@ def _stream_image_chunks( paths: list[str] = [] bytes_for_encoding: list[bytes] = [] - fallback_captions: list[str] = [] + captions: list[str] = [] pages: list[int | None] = [] + pre_caps = precomputed_captions or [] for idx, img in enumerate(images): ext = _MIME_TO_EXT.get(img.mime_type, ".bin") path = img_dir / f"img-{idx:04d}{ext}" @@ -223,30 +248,13 @@ def _stream_image_chunks( continue paths.append(str(path)) bytes_for_encoding.append(img.image_bytes) - fallback_captions.append(img.nearest_caption or "") + vlm_cap = pre_caps[idx].strip() if idx < len(pre_caps) and pre_caps[idx] else "" + captions.append(vlm_cap or (img.nearest_caption or "")) pages.append(img.page_number) if not paths: return 0 - # VLM-generated captions via the loaded chat VLM (when one is loaded - # and vision-capable). Falls back to the parser's nearest_caption - # (page-text blob) when no VLM is available or a request fails. - out_queue.put({"type": "progress", "stage": "caption_images", "progress": 0.87}) - vlm_captions = caption_images( - bytes_for_encoding, - vlm_url = vlm_url, - vlm_model = vlm_model, - ) - captions: list[str] = [ - ( - vlm_captions[i].strip() - if i < len(vlm_captions) and vlm_captions[i].strip() - else fallback_captions[i] - ) - for i in range(len(bytes_for_encoding)) - ] - image_vectors = encode_images(bytes_for_encoding, model_name = model_name) caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()] @@ -696,22 +704,24 @@ def enqueue_ingestion( or resolve_embedder(mode, chunking_strategy) or RAG_EMBEDDING_MODEL ) - # Probe before forking the subprocess so the loaded-model info is - # captured in the parent's process state, then passed to the child. - vlm_url, vlm_model = (None, None) - if mode == "multimodal": - vlm_url, vlm_model = _probe_loaded_vlm() - if vlm_url: - logger.info( - "RAG ingest: will caption figures via loaded chat VLM %s at %s", - vlm_model, - vlm_url, - ) - else: - logger.info( - "RAG ingest: no vision-capable chat model loaded; " - "skipping figure captioning (fallback to page-text)." - ) + # Probe the loaded chat backend so the subprocess can route figure + # captioning to the user's own vision model (no extra VRAM). Runs + # for both modes — text mode splices captions into markdown, and + # multimodal mode additionally feeds them to the image-vector + # encoder. If no vision chat model is loaded, the subprocess falls + # back to the helper VLM (pre-cached at studio startup). + vlm_url, vlm_model = _probe_loaded_vlm() + if vlm_url: + logger.info( + "RAG ingest: will caption figures via loaded chat VLM %s at %s", + vlm_model, + vlm_url, + ) + else: + logger.info( + "RAG ingest: no vision-capable chat model loaded; " + "subprocess will use the helper gemma-3n VLM fallback." + ) job_id = str(uuid4()) with get_connection() as conn: conn.execute( diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py index 040e2b0861..bc0b613b1e 100644 --- a/studio/backend/core/rag/parsers/__init__.py +++ b/studio/backend/core/rag/parsers/__init__.py @@ -44,6 +44,62 @@ class UnsupportedFormatError(ValueError): pass +def inline_image_captions( + pages: list[ParsedPage], + images: list[ParsedImage], + captions: list[str], +) -> list[ParsedPage]: + """Splice per-image captions into the markdown of the pages they came from. + + Mirrors PR #5351's chat-composer pattern: figure captions become + inline text in the page markdown so the chunker indexes them like + any other content. Captions appear at the end of the page's text + block as ``**Figure**: …`` lines. + + ``captions`` is parallel to ``images`` (same length, same order). + Empty or whitespace-only captions are skipped. Images without a + page_number are bucketed onto the single-page documents (DOCX/HTML/ + TXT all collapse to one page). + """ + if not images or not captions: + return list(pages) + if len(captions) != len(images): + # Defensive: caller mismatch shouldn't happen but we don't want + # to lose pages over it. + return list(pages) + + # Bucket captions per page_number (None bucket → single-page docs). + per_page: dict[int | None, list[str]] = {} + for img, cap in zip(images, captions): + cleaned = (cap or "").strip() + if not cleaned: + continue + per_page.setdefault(img.page_number, []).append(cleaned) + + if not per_page: + return list(pages) + + out: list[ParsedPage] = [] + null_bucket = per_page.get(None, []) + for page in pages: + captions_for_this = per_page.get(page.page_number, []) + # If this is the single-page case (no page_number) also flush + # the null-bucket so DOCX/HTML/TXT pick up captions correctly. + if page.page_number is None and null_bucket: + captions_for_this = captions_for_this + null_bucket + if not captions_for_this: + out.append(page) + continue + appendix = "\n\n".join(f"**Figure**: {cap}" for cap in captions_for_this) + out.append( + ParsedPage( + text = f"{page.text}\n\n{appendix}", + page_number = page.page_number, + ) + ) + return out + + def parse(path: Path, *, want_images: bool = False) -> ParseResult: suffix = path.suffix.lower() if suffix == ".pdf": diff --git a/studio/backend/core/rag/scope.py b/studio/backend/core/rag/scope.py index 0049d3be1c..054915ce24 100644 --- a/studio/backend/core/rag/scope.py +++ b/studio/backend/core/rag/scope.py @@ -38,7 +38,7 @@ def resolve_scope_embedder(scope: str) -> str | None: explicit = per_thread.get("embedding_model") or defaults.get("embedding_model") if explicit: return explicit - mode = per_thread.get("mode") or defaults.get("mode") or "multimodal" + mode = per_thread.get("mode") or defaults.get("mode") or "text" chunking_strategy = ( per_thread.get("chunking_strategy") or defaults.get("chunking_strategy") diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 7ad38c3b64..beefc0b2f3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -66,7 +66,7 @@ class CreateKBRequest(BaseModel): description: str | None = None embedding_model: str | None = None chunking_strategy: ChunkingStrategy = "standard" - mode: KBMode = "multimodal" + mode: KBMode = "text" class KBResponse(BaseModel): @@ -401,7 +401,7 @@ def list_knowledge_bases( class RagDefaults(BaseModel): chunking_strategy: ChunkingStrategy = "standard" - mode: KBMode = "multimodal" + mode: KBMode = "text" embedding_model: str | None = None @@ -423,7 +423,7 @@ def _load_rag_defaults() -> RagDefaults: raw = {} return RagDefaults( chunking_strategy = raw.get("chunking_strategy") or "standard", - mode = raw.get("mode") or "multimodal", + mode = raw.get("mode") or "text", embedding_model = raw.get("embedding_model"), ) @@ -495,7 +495,7 @@ def set_rag_defaults( class ThreadRagSettings(BaseModel): chunking_strategy: ChunkingStrategy = "standard" - mode: KBMode = "multimodal" + mode: KBMode = "text" embedding_model: str | None = None diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 4c66d2ebf6..a36a121c5d 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -100,6 +100,34 @@ def precache_helper_gguf(): logger.info(f"Helper GGUF cached: {len(matching)} file(s)") else: logger.warning(f"No GGUF matching variant '{variant}' in {repo}") + + # If the repo also ships an mmproj (vision projection), grab it + # so the helper can be used as a vision-language model by the + # RAG captioner path. Preference order: F16 → BF16 → F32 → any. + # Best-effort — the LLM-assist path doesn't need vision, so a + # missing mmproj is fine and only logged. + mmproj_files = [ + f for f in files if "mmproj" in f.lower() and f.endswith(".gguf") + ] + if mmproj_files: + mmproj_target: Optional[str] = None + for pref in ("mmproj-f16.gguf", "mmproj-bf16.gguf", "mmproj-f32.gguf"): + for cand in mmproj_files: + if cand.lower() == pref: + mmproj_target = cand + break + if mmproj_target: + break + if mmproj_target is None: + mmproj_target = mmproj_files[0] + try: + logger.info(f"Pre-caching helper mmproj: {repo}/{mmproj_target}") + hf_hub_download(repo_id = repo, filename = mmproj_target) + except Exception as mmproj_exc: + logger.warning( + f"Helper mmproj download failed (vision fallback " + f"will be unavailable until cached): {mmproj_exc}" + ) except Exception as e: logger.warning(f"Failed to pre-cache helper GGUF: {e}") finally: diff --git a/studio/backend/utils/rag/config.py b/studio/backend/utils/rag/config.py index b265b90b51..c346ba5624 100644 --- a/studio/backend/utils/rag/config.py +++ b/studio/backend/utils/rag/config.py @@ -34,21 +34,21 @@ RAG_EMBEDDING_MODEL: str = ( # Default embedder per (mode, chunking). (multimodal, late) is unsupported # and rejected at KB-create time in routes/rag.py. # -# Multimodal default is BAAI/BGE-VL-large (~400 M params, 768-d, ~800 MB -# bf16) — small, fast, shared text/image space. Loaded via the -# `_BGEVLAdapter` in core/rag/embeddings.py which bypasses BGE-VL's -# fragile sentence-transformers shim and pre-truncates text to CLIP's -# 77-token cap. +# Text mode is the default; figures from PDFs are captioned at ingest by +# the loaded chat VLM (or a helper gemma-3n fallback) and spliced into +# the page markdown before chunking, so a single 384-d text embedder +# handles all retrieval. Multimodal mode adds image-vector rows on top, +# embedded by Qwen3-VL-Embedding-2B (2 B params, 2048-d, no CLIP text +# cap — full 512-token chunks embed losslessly). # -# To switch back to Qwen3-VL-Embedding-2B (2 B params, 2048-d, no CLIP -# text cap; ~4 GB bf16 / ~1.5 GB 4-bit via FastSentenceTransformer), -# change the ("multimodal", "standard") entry below — the in-process -# loader supports both via `model_name.startswith("BAAI/BGE-VL")` -# routing. +# Alternative multimodal embedders left in tree for manual override: +# - "BAAI/BGE-VL-large" — smaller (~400 M / 768-d) but CLIP-family +# with a 77-token text cap; routed via `_BGEVLAdapter` in +# core/rag/embeddings.py. RAG_EMBEDDER_MATRIX: dict[tuple[str, str], str] = { ("text", "standard"): "BAAI/bge-small-en-v1.5", ("text", "late"): "nomic-ai/nomic-embed-text-v1.5", - ("multimodal", "standard"): "BAAI/BGE-VL-large", + ("multimodal", "standard"): "Qwen/Qwen3-VL-Embedding-2B", } diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 51eb079312..175c878740 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -484,7 +484,7 @@ export function ChatSettingsPanel({ ragDefaults?.chunking_strategy ?? "standard"; const effectiveThreadMode: KBMode = - threadSettings?.mode ?? ragDefaults?.mode ?? "multimodal"; + threadSettings?.mode ?? ragDefaults?.mode ?? "text"; const aui = useAui(); // Brand-new chat has no backend thread yet — initialize the local diff --git a/studio/frontend/src/features/rag/components/kb-create-dialog.tsx b/studio/frontend/src/features/rag/components/kb-create-dialog.tsx index c82e665a7f..c314e59753 100644 --- a/studio/frontend/src/features/rag/components/kb-create-dialog.tsx +++ b/studio/frontend/src/features/rag/components/kb-create-dialog.tsx @@ -43,7 +43,7 @@ export function KBCreateDialog({ const initialStrategy: ChunkingStrategy = defaults?.chunking_strategy ?? "standard"; - const initialMode: KBMode = defaults?.mode ?? "multimodal"; + const initialMode: KBMode = defaults?.mode ?? "text"; const initialEmbedder = defaults?.embedding_model ?? ""; const [name, setName] = useState(""); @@ -76,7 +76,7 @@ export function KBCreateDialog({ setDescription(""); setEmbeddingModel(defaults?.embedding_model ?? ""); setChunkingStrategy(defaults?.chunking_strategy ?? "standard"); - setMode(defaults?.mode ?? "multimodal"); + setMode(defaults?.mode ?? "text"); setError(null); setSubmitting(false); }; diff --git a/studio/frontend/src/features/rag/components/rag-defaults-section.tsx b/studio/frontend/src/features/rag/components/rag-defaults-section.tsx index a8f5dafd0a..e9536840f6 100644 --- a/studio/frontend/src/features/rag/components/rag-defaults-section.tsx +++ b/studio/frontend/src/features/rag/components/rag-defaults-section.tsx @@ -22,7 +22,7 @@ export function RagDefaultsSection() { const [chunkingStrategy, setChunkingStrategy] = useState("standard"); - const [mode, setMode] = useState("multimodal"); + const [mode, setMode] = useState("text"); const [embeddingModel, setEmbeddingModel] = useState(""); const [error, setError] = useState(null);