Studio: VLM-caption figures at ingest + pass image hits to LLM + render in card
This commit is contained in:
Roland Tannous
parent
9d5719a475
commit
7c1a09b350
5 changed files with 261 additions and 10 deletions
145
studio/backend/core/rag/captioner.py
Normal file
145
studio/backend/core/rag/captioner.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
# 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 a small generative VLM (PaddleOCR-VL, 4-bit).
|
||||
|
||||
Defensive by design: any load or per-image failure returns an empty
|
||||
string so the ingestion pipeline can fall back to its prior page-text
|
||||
caption behaviour rather than crashing.
|
||||
|
||||
Lifecycle: lives inside the ingestion subprocess (`_subprocess_worker`
|
||||
in `ingestion.py`). The model loads on first `caption_images` call and
|
||||
is released when the subprocess exits — no impact on the parent chat
|
||||
model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_CAPTION_MODEL_NAME = "unsloth/PaddleOCR-VL"
|
||||
_PROMPT = (
|
||||
"Describe this figure in <=60 words. Focus on factual content "
|
||||
"(axes, labels, captions, visible text, main objects). "
|
||||
"Do not speculate beyond what is visible."
|
||||
)
|
||||
_MAX_NEW_TOKENS = 120
|
||||
|
||||
_lock = threading.Lock()
|
||||
_model: Any | None = None
|
||||
_processor: Any | None = None
|
||||
_load_failed: bool = False
|
||||
|
||||
|
||||
def _load() -> tuple[Any, Any] | None:
|
||||
"""Lazy-load PaddleOCR-VL. Returns (model, processor) or None on failure.
|
||||
|
||||
Sentinel-cached: once a load failure occurs in this subprocess we
|
||||
don't keep retrying for every batch of images.
|
||||
"""
|
||||
global _model, _processor, _load_failed
|
||||
if _load_failed:
|
||||
return None
|
||||
with _lock:
|
||||
if _model is not None and _processor is not None:
|
||||
return _model, _processor
|
||||
try:
|
||||
import torch
|
||||
from transformers import AutoModelForVision2Seq, AutoProcessor
|
||||
|
||||
logger.info("Loading RAG captioner: %s", _CAPTION_MODEL_NAME)
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
_CAPTION_MODEL_NAME,
|
||||
trust_remote_code = True,
|
||||
)
|
||||
model = AutoModelForVision2Seq.from_pretrained(
|
||||
_CAPTION_MODEL_NAME,
|
||||
trust_remote_code = True,
|
||||
load_in_4bit = True,
|
||||
device_map = "auto",
|
||||
torch_dtype = torch.float16,
|
||||
)
|
||||
model.eval()
|
||||
_model = model
|
||||
_processor = processor
|
||||
return _model, _processor
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"RAG captioner %s failed to load (%s). Falling back to "
|
||||
"page-text captions for this ingestion.",
|
||||
_CAPTION_MODEL_NAME,
|
||||
exc,
|
||||
)
|
||||
_load_failed = True
|
||||
return None
|
||||
|
||||
|
||||
def caption_images(image_bytes_list: list[bytes]) -> list[str]:
|
||||
"""Generate one short caption per image. Same-length output.
|
||||
|
||||
Returns ``""`` for any image whose captioning failed (or all images
|
||||
if the model couldn't load). Never raises — ingestion must stay
|
||||
resilient to VLM unavailability.
|
||||
"""
|
||||
if not image_bytes_list:
|
||||
return []
|
||||
loaded = _load()
|
||||
if loaded is None:
|
||||
return ["" for _ in image_bytes_list]
|
||||
|
||||
model, processor = loaded
|
||||
from io import BytesIO
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
out: list[str] = []
|
||||
for blob in image_bytes_list:
|
||||
try:
|
||||
image = Image.open(BytesIO(blob)).convert("RGB")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": _PROMPT},
|
||||
],
|
||||
}
|
||||
]
|
||||
prompt = processor.apply_chat_template(
|
||||
messages,
|
||||
add_generation_prompt = True,
|
||||
)
|
||||
inputs = processor(
|
||||
images = image,
|
||||
text = prompt,
|
||||
return_tensors = "pt",
|
||||
)
|
||||
input_ids_len = (
|
||||
int(inputs["input_ids"].shape[1]) if "input_ids" in inputs else 0
|
||||
)
|
||||
inputs = {
|
||||
k: (v.to(model.device) if hasattr(v, "to") else v)
|
||||
for k, v in inputs.items()
|
||||
}
|
||||
with torch.no_grad():
|
||||
output_ids = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens = _MAX_NEW_TOKENS,
|
||||
do_sample = False,
|
||||
)
|
||||
# Decode only the newly-generated tokens, not the prompt echo.
|
||||
new_tokens = output_ids[:, input_ids_len:]
|
||||
caption = processor.batch_decode(
|
||||
new_tokens,
|
||||
skip_special_tokens = True,
|
||||
)[0]
|
||||
out.append(caption.strip())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("caption_images: skipping one image: %s", exc)
|
||||
out.append("")
|
||||
return out
|
||||
|
|
@ -191,6 +191,7 @@ def _stream_image_chunks(
|
|||
first_index: int,
|
||||
) -> int:
|
||||
"""Persist images, emit image+caption chunks; pairs share pair_group."""
|
||||
from core.rag.captioner import caption_images
|
||||
from core.rag.embeddings import encode, encode_images
|
||||
from utils.paths.storage_roots import ensure_dir, rag_uploads_root
|
||||
|
||||
|
|
@ -203,7 +204,7 @@ def _stream_image_chunks(
|
|||
|
||||
paths: list[str] = []
|
||||
bytes_for_encoding: list[bytes] = []
|
||||
captions: list[str] = []
|
||||
fallback_captions: list[str] = []
|
||||
pages: list[int | None] = []
|
||||
for idx, img in enumerate(images):
|
||||
ext = _MIME_TO_EXT.get(img.mime_type, ".bin")
|
||||
|
|
@ -215,12 +216,25 @@ def _stream_image_chunks(
|
|||
continue
|
||||
paths.append(str(path))
|
||||
bytes_for_encoding.append(img.image_bytes)
|
||||
captions.append(img.nearest_caption or "")
|
||||
fallback_captions.append(img.nearest_caption or "")
|
||||
pages.append(img.page_number)
|
||||
|
||||
if not paths:
|
||||
return 0
|
||||
|
||||
# VLM-generated captions; falls back to the parser's nearest_caption
|
||||
# (page-text blob) when the VLM is unavailable or fails for an image.
|
||||
out_queue.put({"type": "progress", "stage": "caption_images", "progress": 0.87})
|
||||
vlm_captions = caption_images(bytes_for_encoding)
|
||||
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()]
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ so the model never sees KB UUIDs.
|
|||
from __future__ import annotations
|
||||
|
||||
from contextvars import ContextVar
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from loggers import get_logger
|
||||
|
|
@ -110,6 +111,13 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
|
|||
kind = hit.get("kind")
|
||||
if kind and kind != "text":
|
||||
attrs.append(f'kind="{_xml_attr(kind)}"')
|
||||
image_path = hit.get("image_path")
|
||||
document_id = hit.get("document_id")
|
||||
if kind == "image" and image_path and document_id:
|
||||
# Mirror routes/rag.py search-response shape so the frontend
|
||||
# tool card can render the image inline via the same route.
|
||||
image_url = f"/api/rag/images/{document_id}/{Path(image_path).name}"
|
||||
attrs.append(f'image_url="{_xml_attr(image_url)}"')
|
||||
text = (hit.get("text") or "").strip()
|
||||
blocks.append(f"<chunk {' '.join(attrs)}>\n{text}\n</chunk>")
|
||||
return "\n\n".join(blocks)
|
||||
|
|
@ -210,7 +218,8 @@ def search_knowledge_base(
|
|||
rows = conn.execute(
|
||||
f"""
|
||||
SELECT c.id AS chunk_id, c.text, c.page_number,
|
||||
c.token_count, c.kind, d.filename
|
||||
c.token_count, c.kind, c.image_path,
|
||||
c.document_id, d.filename
|
||||
FROM rag_chunks c
|
||||
JOIN rag_documents d ON d.id = c.document_id
|
||||
WHERE c.id IN ({placeholders})
|
||||
|
|
@ -241,13 +250,14 @@ def search_knowledge_base(
|
|||
else:
|
||||
hits = hits[:k]
|
||||
|
||||
# Skip image-kind hits; the paired caption surfaces separately.
|
||||
# Merge Hit-side metadata (score, dense_score, chunk_index) into the
|
||||
# sqlite-side row so the formatter sees one flat dict per chunk.
|
||||
# Image-kind hits flow through so the multimodal embedder's match
|
||||
# can reach the LLM; their image_url lets the UI render the picture.
|
||||
formatted: list[dict] = []
|
||||
for hit in hits:
|
||||
row = lookup.get(hit.chunk_id)
|
||||
if row is None or row.get("kind") == "image":
|
||||
if row is None:
|
||||
continue
|
||||
formatted.append(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ import {
|
|||
type ToolCallMessagePartComponent,
|
||||
useAuiState,
|
||||
} from "@assistant-ui/react";
|
||||
import { FileTextIcon, LoaderIcon } from "lucide-react";
|
||||
import { authFetch } from "@/features/auth";
|
||||
import { FileTextIcon, ImageIcon, LoaderIcon } from "lucide-react";
|
||||
import { memo, useEffect, useState } from "react";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
|
|
@ -25,6 +26,7 @@ export interface ParsedChunk {
|
|||
chunkIndex?: string;
|
||||
tokens?: string;
|
||||
kind?: string;
|
||||
imageUrl?: string;
|
||||
text: string;
|
||||
}
|
||||
|
||||
|
|
@ -62,6 +64,7 @@ export function parseChunks(raw: string): ParsedChunk[] {
|
|||
chunkIndex: attrs.chunk_index,
|
||||
tokens: attrs.tokens,
|
||||
kind: attrs.kind,
|
||||
imageUrl: attrs.image_url,
|
||||
text,
|
||||
});
|
||||
}
|
||||
|
|
@ -72,6 +75,57 @@ export function parseChunks(raw: string): ParsedChunk[] {
|
|||
return out;
|
||||
}
|
||||
|
||||
/** Fetch a backend image via the bearer-authed `authFetch`, expose it
|
||||
* as a blob URL for `<img src>`. Cleans up the object URL on unmount. */
|
||||
function useAuthedImageUrl(path: string | undefined): string | undefined {
|
||||
const [url, setUrl] = useState<string | undefined>(undefined);
|
||||
useEffect(() => {
|
||||
if (!path) {
|
||||
setUrl(undefined);
|
||||
return;
|
||||
}
|
||||
let cancelled = false;
|
||||
let objectUrl: string | undefined;
|
||||
void authFetch(path)
|
||||
.then(async (response) => {
|
||||
if (!response.ok) throw new Error(`image fetch ${response.status}`);
|
||||
return response.blob();
|
||||
})
|
||||
.then((blob) => {
|
||||
if (cancelled) return;
|
||||
objectUrl = URL.createObjectURL(blob);
|
||||
setUrl(objectUrl);
|
||||
})
|
||||
.catch(() => {
|
||||
if (!cancelled) setUrl(undefined);
|
||||
});
|
||||
return () => {
|
||||
cancelled = true;
|
||||
if (objectUrl) URL.revokeObjectURL(objectUrl);
|
||||
};
|
||||
}, [path]);
|
||||
return url;
|
||||
}
|
||||
|
||||
function ChunkImage({ url, alt }: { url: string; alt: string }) {
|
||||
const blobUrl = useAuthedImageUrl(url);
|
||||
if (!blobUrl) {
|
||||
return (
|
||||
<div className="mb-2 flex h-32 items-center justify-center rounded-md bg-muted/60 text-[10px] text-muted-foreground">
|
||||
<ImageIcon className="mr-1.5 size-3" />
|
||||
Loading image…
|
||||
</div>
|
||||
);
|
||||
}
|
||||
return (
|
||||
<img
|
||||
src={blobUrl}
|
||||
alt={alt}
|
||||
className="mb-2 max-h-64 w-full rounded-md object-contain"
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
|
||||
const meta: string[] = [];
|
||||
if (chunk.page) meta.push(`page ${chunk.page}`);
|
||||
|
|
@ -93,7 +147,11 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
|
|||
<span className="rounded bg-foreground/10 px-1.5 py-0.5 font-mono text-[10px] font-semibold">
|
||||
[{chunk.id}]
|
||||
</span>
|
||||
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
{chunk.kind === "image" ? (
|
||||
<ImageIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
) : (
|
||||
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
)}
|
||||
<span className="truncate font-medium" title={chunk.source}>
|
||||
{chunk.source}
|
||||
</span>
|
||||
|
|
@ -104,9 +162,14 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
|
|||
</span>
|
||||
) : null}
|
||||
</div>
|
||||
<pre className="max-h-48 overflow-auto whitespace-pre-wrap break-words text-[11px] leading-relaxed text-foreground/80">
|
||||
{chunk.text}
|
||||
</pre>
|
||||
{chunk.kind === "image" && chunk.imageUrl ? (
|
||||
<ChunkImage url={chunk.imageUrl} alt={chunk.source} />
|
||||
) : null}
|
||||
{chunk.text ? (
|
||||
<pre className="max-h-48 overflow-auto whitespace-pre-wrap break-words text-[11px] leading-relaxed text-foreground/80">
|
||||
{chunk.text}
|
||||
</pre>
|
||||
) : null}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -153,6 +153,25 @@ def test_format_hits_offsets_ids_by_start_id():
|
|||
assert '<chunk id="1"' not in result
|
||||
|
||||
|
||||
def test_format_hits_emits_image_url_for_image_kind():
|
||||
from core.rag.tool import _format_hits_for_llm
|
||||
|
||||
hits = [
|
||||
{
|
||||
"filename": "paper.pdf",
|
||||
"text": "Figure 1 shows a bar chart of X over time.",
|
||||
"kind": "image",
|
||||
"image_path": "/abs/path/images/doc-123/img-0007.png",
|
||||
"document_id": "doc-123",
|
||||
"page_number": 5,
|
||||
}
|
||||
]
|
||||
result = _format_hits_for_llm(hits)
|
||||
assert 'kind="image"' in result
|
||||
assert 'image_url="/api/rag/images/doc-123/img-0007.png"' in result
|
||||
assert "Figure 1 shows a bar chart" in result
|
||||
|
||||
|
||||
def test_format_hits_escapes_xml_in_source():
|
||||
from core.rag.tool import _format_hits_for_llm
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue