ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: text-mode default with VLM-captioned figure splicing; helper VLM fallback

Browse source
This commit is contained in:
Roland Tannous 2026-05-27 11:16:56 +04:00
commit aedede2f2e
10 changed files with 284 additions and 117 deletions
Download patch file Download diff file Expand all files Collapse all files

179
studio/backend/core/rag/captioner.py
View file

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

98
studio/backend/core/rag/ingestion.py
View file

@ -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(

56
studio/backend/core/rag/parsers/__init__.py
View file

@ -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":

2
studio/backend/core/rag/scope.py
View file

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

8
studio/backend/routes/rag.py
View file

@ -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

28
studio/backend/utils/datasets/llm_assist.py
View file

@ -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:

22
studio/backend/utils/rag/config.py
View file

@ -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",
}

2
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -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

4
studio/frontend/src/features/rag/components/kb-create-dialog.tsx
View file

@ -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);
};

2
studio/frontend/src/features/rag/components/rag-defaults-section.tsx
View file

@ -22,7 +22,7 @@ export function RagDefaultsSection() {
const [chunkingStrategy, setChunkingStrategy] =
useState<ChunkingStrategy>("standard");
const [mode, setMode] = useState<KBMode>("multimodal");
const [mode, setMode] = useState<KBMode>("text");
const [embeddingModel, setEmbeddingModel] = useState("");
const [error, setError] = useState<string | null>(null);