Spaces:

ORCMS
/

vlps-demo

Sleeping

Fix SigLIP text padding (max_length) — cross-modal was returning random images

bfe0960 5 months ago

15.8 kB

import base64
import io
import os
from typing import Optional
import pandas as pd
import numpy as np
from pathlib import Path
from PIL import Image


def load_coco_dataset(split="validation"):
    """
    Loads the MS-COCO 2017 dataset from HuggingFace.
    """
    # Imported lazily: only needed for COCO download, not for the retrieval/QA helpers
    # below (so the module imports cleanly in lean deployments without datasets/aiohttp).
    import aiohttp
    from datasets import load_dataset, DownloadConfig

    print(f"Downloading MS-COCO {split}...")
    dl_config = DownloadConfig(
        storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=7200, sock_read=3600)}}
    )
    ds = load_dataset(
        "HuggingFaceM4/COCO",
        split=split,
        trust_remote_code=True,
        download_config=dl_config,
        storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=7200, sock_read=3600)}}
    )
    return ds


def _get_image_id(item):
    return str(item['cocoid']) if 'cocoid' in item else str(item['image_id'])


def generate_coco_docs(dataset, index_name="coco_captions",
                       caption_embeddings_df=None, image_embeddings_df=None):
    """
    Generator that yields index-ready dicts for OpenSearch.
    Optionally attaches caption_vec and image_vec if DataFrames are provided.
    caption_embeddings_df: DataFrame with columns [image_id, caption, embedding]
    image_embeddings_df: DataFrame with columns [image_id, embedding] (one per image)
    """
    # Build lookup dicts using vectorised ops — much faster than iterrows()
    caption_vec_lookup = {}
    if caption_embeddings_df is not None:
        caption_vec_lookup = {
            (iid, cap): emb
            for iid, cap, emb in zip(
                caption_embeddings_df['image_id'],
                caption_embeddings_df['caption'],
                caption_embeddings_df['embedding']
            )
        }

    image_vec_lookup = {}
    if image_embeddings_df is not None:
        image_vec_lookup = dict(zip(
            image_embeddings_df['image_id'],
            image_embeddings_df['embedding']
        ))

    for item in dataset:
        image_id = _get_image_id(item)
        captions = [item['sentences']['raw']]  # one caption per dataset row

        # Use metadata fields instead of decoding the image for width/height
        width  = item.get('width',  item['image'].width  if 'image' in item else 0)
        height = item.get('height', item['image'].height if 'image' in item else 0)

        # Extract COCO category names if available
        categories = []
        num_objects = 0
        if 'objects' in item and 'categories' in item['objects']:
            categories = list(set(item['objects']['categories']))
            num_objects = len(item['objects']['categories'])

        img_vec = image_vec_lookup.get(image_id)

        for caption in captions:
            doc = {
                "_index": index_name,
                "image_id": image_id,
                "caption": caption,
                "caption_length": len(caption.split()),
                "width": width,
                "height": height,
                "categories": categories,
                "num_objects": num_objects,
            }
            cap_vec = caption_vec_lookup.get((image_id, caption))
            if cap_vec is not None:
                doc["caption_vec"] = cap_vec if isinstance(cap_vec, list) else cap_vec.tolist()
            if img_vec is not None:
                doc["image_vec"] = img_vec if isinstance(img_vec, list) else img_vec.tolist()
            yield doc


def compute_caption_embeddings(dataset, model_name, save_path, batch_size=128):
    """
    Computes caption embeddings using a sentence-transformers model and saves to parquet.
    Each row has: image_id, caption, embedding (as list).
    """
    from sentence_transformers import SentenceTransformer

    save_path = Path(save_path)
    if save_path.exists():
        print(f"Embeddings already exist at {save_path}, loading...")
        return pd.read_parquet(save_path)

    model = SentenceTransformer(model_name)

    rows = []
    for item in dataset:
        image_id = _get_image_id(item)
        caption = item['sentences']['raw']  # one caption per dataset row
        rows.append({"image_id": image_id, "caption": caption})

    df = pd.DataFrame(rows)
    print(f"Computing embeddings for {len(df)} captions with {model_name}...")
    embeddings = model.encode(df['caption'].tolist(), batch_size=batch_size, show_progress_bar=True, normalize_embeddings=True)
    df['embedding'] = [emb.tolist() for emb in embeddings]

    save_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(save_path)
    print(f"Saved embeddings to {save_path}")
    return df


def compute_image_embeddings(dataset, model_name, save_path, batch_size=32):
    """
    Computes image embeddings using CLIP or SigLIP and saves to parquet.
    Each row has: image_id, embedding (as list).
    """
    import torch
    from transformers import AutoProcessor, AutoModel

    save_path = Path(save_path)
    if save_path.exists():
        print(f"Image embeddings already exist at {save_path}, loading...")
        return pd.read_parquet(save_path)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = AutoProcessor.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(device)
    model.eval()

    rows = []
    images_batch = []
    ids_batch = []

    for idx, item in enumerate(dataset):
        image_id = _get_image_id(item)
        img = item['image'].convert("RGB")
        images_batch.append(img)
        ids_batch.append(image_id)

        if len(images_batch) == batch_size or idx == len(dataset) - 1:
            inputs = processor(images=images_batch, return_tensors="pt", padding=True).to(device)
            with torch.no_grad():
                out = model.get_image_features(**inputs)
                image_features = out.pooler_output if hasattr(out, 'pooler_output') else out
                image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            embeddings = image_features.cpu().numpy()
            for i, emb in enumerate(embeddings):
                rows.append({"image_id": ids_batch[i], "embedding": emb.tolist()})
            images_batch = []
            ids_batch = []
            if (idx + 1) % 500 == 0:
                print(f"  Processed {idx + 1}/{len(dataset)} images...")

    df = pd.DataFrame(rows)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(save_path)
    print(f"Saved {len(df)} image embeddings to {save_path}")
    return df


def compute_text_embeddings_clip(texts, model_name, save_path=None, batch_size=128):
    """
    Computes text embeddings using a CLIP/SigLIP model (for cross-modal search).
    Returns numpy array of shape (len(texts), dim).
    Optionally saves to disk.
    """
    import torch
    from transformers import AutoProcessor, AutoModel

    if save_path:
        save_path = Path(save_path)
        if save_path.exists():
            print(f"Text embeddings already exist at {save_path}, loading...")
            return np.load(save_path)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = AutoProcessor.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(device)
    model.eval()

    all_embs = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        inputs = processor(text=batch, return_tensors="pt", padding=True, truncation=True).to(device)
        with torch.no_grad():
            out = model.get_text_features(**inputs)
            text_features = out.pooler_output if hasattr(out, 'pooler_output') else out
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        all_embs.append(text_features.cpu().numpy())

    embeddings = np.concatenate(all_embs, axis=0)
    if save_path:
        save_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(save_path, embeddings)
        print(f"Saved text embeddings to {save_path}")
    return embeddings


# ---------------------------------------------------------------------------
# Phase 2 helpers
# ---------------------------------------------------------------------------

def build_image_lookup(dataset) -> dict:
    """Build an image_id -> PIL Image dict from the COCO dataset."""
    lookup = {}
    for item in dataset:
        iid = _get_image_id(item)
        lookup[iid] = item['image'].convert('RGB')
    return lookup


def image_to_base64(img: Image.Image, fmt: str = 'JPEG') -> str:
    """Convert a PIL image to a base64 data-URL string."""
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f'data:image/jpeg;base64,{b64}'


def ask_lvlm(llm_client, image: Image.Image, question: str,
             model: Optional[str] = None, system_prompt: Optional[str] = None,
             context: Optional[str] = None, max_tokens: int = 512) -> str:
    """
    Send an image + question to the vLLM-hosted LVLM and return the answer.
    llm_client: openai.OpenAI instance pointed at the vLLM endpoint.
    model: model ID string (defaults to PRIMARY_MODEL env var or Qwen2.5-VL).
    context: optional retrieved caption to prepend as context.
    """
    model = model or os.getenv('VLLM_MODEL', 'google/gemma-4-31B-it')
    image_url = image_to_base64(image)

    user_content = []
    if context:
        user_content.append({'type': 'text', 'text': f'Context: {context}\n\n'})
    user_content.append({'type': 'image_url', 'image_url': {'url': image_url}})
    user_content.append({'type': 'text', 'text': question})

    messages = []
    if system_prompt:
        messages.append({'role': 'system', 'content': system_prompt})
    messages.append({'role': 'user', 'content': user_content})

    response = llm_client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=max_tokens,
    )
    return response.choices[0].message.content


def load_siglip_model(model_name: str = 'google/siglip-base-patch16-224'):
    """Load SigLIP processor and model. Returns (processor, model, device)."""
    import torch
    from transformers import AutoProcessor, AutoModel

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    processor = AutoProcessor.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(device)
    model.eval()
    return processor, model, device


def _as_tensor(out):
    """Coerce a model output to the feature tensor (newer transformers may wrap it)."""
    import torch
    if torch.is_tensor(out):
        return out
    if getattr(out, "pooler_output", None) is not None:
        return out.pooler_output
    if getattr(out, "last_hidden_state", None) is not None:
        return out.last_hidden_state.mean(dim=1)
    raise TypeError(f"Unexpected SigLIP feature output: {type(out)}")


def embed_text_siglip(text: str, processor, model, device) -> list:
    """Embed a text query with SigLIP for cross-modal search.
    SigLIP's text encoder is trained with fixed 64-token padding; using
    padding='max_length' (not dynamic) is required for correct embeddings.
    """
    import torch
    inputs = processor(text=[text], return_tensors='pt', padding='max_length',
                       truncation=True, max_length=64).to(device)
    with torch.no_grad():
        feat = _as_tensor(model.get_text_features(**inputs))
        feat = feat / feat.norm(dim=-1, keepdim=True)
    return feat.cpu().numpy()[0].tolist()


def embed_image_siglip(img: Image.Image, processor, model, device) -> list:
    """Embed a PIL image with SigLIP."""
    import torch
    inputs = processor(images=img, return_tensors='pt').to(device)
    with torch.no_grad():
        feat = _as_tensor(model.get_image_features(**inputs))
        feat = feat / feat.norm(dim=-1, keepdim=True)
    return feat.cpu().numpy()[0].tolist()


def retrieve_by_text_bm25(client, query: str, index_name: str, top_k: int = 3) -> list:
    """Retrieve image-caption pairs using BM25 keyword search."""
    resp = client.search(
        index=index_name,
        body={'size': top_k, 'query': {'match': {'caption': query}}}
    )
    return [h['_source'] for h in resp['hits']['hits']]


def retrieve_by_text_knn(client, query: str, index_name: str,
                         bge_model, top_k: int = 3) -> list:
    """Retrieve image-caption pairs using BGE semantic k-NN."""
    vec = bge_model.encode(query, normalize_embeddings=True).tolist()
    resp = client.search(
        index=index_name,
        body={'size': top_k, 'query': {'knn': {'caption_vec': {'vector': vec, 'k': top_k}}}}
    )
    return [h['_source'] for h in resp['hits']['hits']]


def retrieve_by_text_crossmodal(client, query: str, index_name: str,
                                siglip_processor, siglip_model, siglip_device,
                                top_k: int = 3) -> list:
    """Retrieve images using SigLIP cross-modal k-NN (text → image vectors)."""
    vec = embed_text_siglip(query, siglip_processor, siglip_model, siglip_device)
    resp = client.search(
        index=index_name,
        body={'size': top_k, 'query': {'knn': {'image_vec': {'vector': vec, 'k': top_k}}}}
    )
    return [h['_source'] for h in resp['hits']['hits']]


def retrieve_by_image_crossmodal(client, img: Image.Image, index_name: str,
                                 siglip_processor, siglip_model, siglip_device,
                                 exclude_id: Optional[str] = None, top_k: int = 3) -> list:
    """Retrieve similar images using SigLIP image embedding."""
    vec = embed_image_siglip(img, siglip_processor, siglip_model, siglip_device)
    resp = client.search(
        index=index_name,
        body={'size': top_k + 1, 'query': {'knn': {'image_vec': {'vector': vec, 'k': top_k + 1}}}}
    )
    results = [
        h['_source'] for h in resp['hits']['hits']
        if h['_source']['image_id'] != exclude_id
    ]
    return results[:top_k]


def rag_answer(llm_client, os_client, question: str,
               image_lookup: dict, bge_model,
               siglip_processor, siglip_model, siglip_device,
               index_bge: str, index_multi: str,
               retrieval_method: str = 'crossmodal',
               top_k: int = 1, use_caption_context: bool = True,
               model: Optional[str] = None) -> dict:
    """
    Full RAG pipeline: retrieve relevant image-caption, then generate answer with LVLM.
    retrieval_method: 'bm25' | 'knn' | 'crossmodal'
    """
    if retrieval_method == 'bm25':
        hits = retrieve_by_text_bm25(os_client, question, index_bge, top_k)
    elif retrieval_method == 'knn':
        hits = retrieve_by_text_knn(os_client, question, index_bge, bge_model, top_k)
    else:
        hits = retrieve_by_text_crossmodal(
            os_client, question, index_multi,
            siglip_processor, siglip_model, siglip_device, top_k
        )

    if not hits:
        return {'question': question, 'answer': 'No relevant images found.', 'hits': []}

    top_hit = hits[0]
    image_id = top_hit['image_id']
    caption: str = top_hit.get('caption', '')
    img = image_lookup.get(image_id)

    if img is None:
        return {'question': question, 'answer': f'Image {image_id} not found in dataset.', 'hits': hits}

    answer = ask_lvlm(
        llm_client, img, question,
        model=model,
        context=caption if use_caption_context else None,
    )
    return {
        'question': question,
        'retrieved_image_id': image_id,
        'retrieved_caption': caption,
        'answer': answer,
        'hits': hits,
    }

	import base64
	import io
	import os
	from typing import Optional
	import pandas as pd
	import numpy as np
	from pathlib import Path
	from PIL import Image


	def load_coco_dataset(split="validation"):
	"""
	Loads the MS-COCO 2017 dataset from HuggingFace.
	"""
	# Imported lazily: only needed for COCO download, not for the retrieval/QA helpers
	# below (so the module imports cleanly in lean deployments without datasets/aiohttp).
	import aiohttp
	from datasets import load_dataset, DownloadConfig

	print(f"Downloading MS-COCO {split}...")
	dl_config = DownloadConfig(
	storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=7200, sock_read=3600)}}
	)
	ds = load_dataset(
	"HuggingFaceM4/COCO",
	split=split,
	trust_remote_code=True,
	download_config=dl_config,
	storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=7200, sock_read=3600)}}
	)
	return ds


	def _get_image_id(item):
	return str(item['cocoid']) if 'cocoid' in item else str(item['image_id'])


	def generate_coco_docs(dataset, index_name="coco_captions",
	caption_embeddings_df=None, image_embeddings_df=None):
	"""
	Generator that yields index-ready dicts for OpenSearch.
	Optionally attaches caption_vec and image_vec if DataFrames are provided.
	caption_embeddings_df: DataFrame with columns [image_id, caption, embedding]
	image_embeddings_df: DataFrame with columns [image_id, embedding] (one per image)
	"""
	# Build lookup dicts using vectorised ops — much faster than iterrows()
	caption_vec_lookup = {}
	if caption_embeddings_df is not None:
	caption_vec_lookup = {
	(iid, cap): emb
	for iid, cap, emb in zip(
	caption_embeddings_df['image_id'],
	caption_embeddings_df['caption'],
	caption_embeddings_df['embedding']
	)
	}

	image_vec_lookup = {}
	if image_embeddings_df is not None:
	image_vec_lookup = dict(zip(
	image_embeddings_df['image_id'],
	image_embeddings_df['embedding']
	))

	for item in dataset:
	image_id = _get_image_id(item)
	captions = [item['sentences']['raw']] # one caption per dataset row

	# Use metadata fields instead of decoding the image for width/height
	width = item.get('width', item['image'].width if 'image' in item else 0)
	height = item.get('height', item['image'].height if 'image' in item else 0)

	# Extract COCO category names if available
	categories = []
	num_objects = 0
	if 'objects' in item and 'categories' in item['objects']:
	categories = list(set(item['objects']['categories']))
	num_objects = len(item['objects']['categories'])

	img_vec = image_vec_lookup.get(image_id)

	for caption in captions:
	doc = {
	"_index": index_name,
	"image_id": image_id,
	"caption": caption,
	"caption_length": len(caption.split()),
	"width": width,
	"height": height,
	"categories": categories,
	"num_objects": num_objects,
	}
	cap_vec = caption_vec_lookup.get((image_id, caption))
	if cap_vec is not None:
	doc["caption_vec"] = cap_vec if isinstance(cap_vec, list) else cap_vec.tolist()
	if img_vec is not None:
	doc["image_vec"] = img_vec if isinstance(img_vec, list) else img_vec.tolist()
	yield doc


	def compute_caption_embeddings(dataset, model_name, save_path, batch_size=128):
	"""
	Computes caption embeddings using a sentence-transformers model and saves to parquet.
	Each row has: image_id, caption, embedding (as list).
	"""
	from sentence_transformers import SentenceTransformer

	save_path = Path(save_path)
	if save_path.exists():
	print(f"Embeddings already exist at {save_path}, loading...")
	return pd.read_parquet(save_path)

	model = SentenceTransformer(model_name)

	rows = []
	for item in dataset:
	image_id = _get_image_id(item)
	caption = item['sentences']['raw'] # one caption per dataset row
	rows.append({"image_id": image_id, "caption": caption})

	df = pd.DataFrame(rows)
	print(f"Computing embeddings for {len(df)} captions with {model_name}...")
	embeddings = model.encode(df['caption'].tolist(), batch_size=batch_size, show_progress_bar=True, normalize_embeddings=True)
	df['embedding'] = [emb.tolist() for emb in embeddings]

	save_path.parent.mkdir(parents=True, exist_ok=True)
	df.to_parquet(save_path)
	print(f"Saved embeddings to {save_path}")
	return df


	def compute_image_embeddings(dataset, model_name, save_path, batch_size=32):
	"""
	Computes image embeddings using CLIP or SigLIP and saves to parquet.
	Each row has: image_id, embedding (as list).
	"""
	import torch
	from transformers import AutoProcessor, AutoModel

	save_path = Path(save_path)
	if save_path.exists():
	print(f"Image embeddings already exist at {save_path}, loading...")
	return pd.read_parquet(save_path)

	device = "cuda" if torch.cuda.is_available() else "cpu"
	processor = AutoProcessor.from_pretrained(model_name)
	model = AutoModel.from_pretrained(model_name).to(device)
	model.eval()

	rows = []
	images_batch = []
	ids_batch = []

	for idx, item in enumerate(dataset):
	image_id = _get_image_id(item)
	img = item['image'].convert("RGB")
	images_batch.append(img)
	ids_batch.append(image_id)

	if len(images_batch) == batch_size or idx == len(dataset) - 1:
	inputs = processor(images=images_batch, return_tensors="pt", padding=True).to(device)
	with torch.no_grad():
	out = model.get_image_features(**inputs)
	image_features = out.pooler_output if hasattr(out, 'pooler_output') else out
	image_features = image_features / image_features.norm(dim=-1, keepdim=True)
	embeddings = image_features.cpu().numpy()
	for i, emb in enumerate(embeddings):
	rows.append({"image_id": ids_batch[i], "embedding": emb.tolist()})
	images_batch = []
	ids_batch = []
	if (idx + 1) % 500 == 0:
	print(f" Processed {idx + 1}/{len(dataset)} images...")

	df = pd.DataFrame(rows)
	save_path.parent.mkdir(parents=True, exist_ok=True)
	df.to_parquet(save_path)
	print(f"Saved {len(df)} image embeddings to {save_path}")
	return df


	def compute_text_embeddings_clip(texts, model_name, save_path=None, batch_size=128):
	"""
	Computes text embeddings using a CLIP/SigLIP model (for cross-modal search).
	Returns numpy array of shape (len(texts), dim).
	Optionally saves to disk.
	"""
	import torch
	from transformers import AutoProcessor, AutoModel

	if save_path:
	save_path = Path(save_path)
	if save_path.exists():
	print(f"Text embeddings already exist at {save_path}, loading...")
	return np.load(save_path)

	device = "cuda" if torch.cuda.is_available() else "cpu"
	processor = AutoProcessor.from_pretrained(model_name)
	model = AutoModel.from_pretrained(model_name).to(device)
	model.eval()

	all_embs = []
	for i in range(0, len(texts), batch_size):
	batch = texts[i:i + batch_size]
	inputs = processor(text=batch, return_tensors="pt", padding=True, truncation=True).to(device)
	with torch.no_grad():
	out = model.get_text_features(**inputs)
	text_features = out.pooler_output if hasattr(out, 'pooler_output') else out
	text_features = text_features / text_features.norm(dim=-1, keepdim=True)
	all_embs.append(text_features.cpu().numpy())

	embeddings = np.concatenate(all_embs, axis=0)
	if save_path:
	save_path.parent.mkdir(parents=True, exist_ok=True)
	np.save(save_path, embeddings)
	print(f"Saved text embeddings to {save_path}")
	return embeddings


	# ---------------------------------------------------------------------------
	# Phase 2 helpers
	# ---------------------------------------------------------------------------

	def build_image_lookup(dataset) -> dict:
	"""Build an image_id -> PIL Image dict from the COCO dataset."""
	lookup = {}
	for item in dataset:
	iid = _get_image_id(item)
	lookup[iid] = item['image'].convert('RGB')
	return lookup


	def image_to_base64(img: Image.Image, fmt: str = 'JPEG') -> str:
	"""Convert a PIL image to a base64 data-URL string."""
	buf = io.BytesIO()
	img.save(buf, format=fmt)
	b64 = base64.b64encode(buf.getvalue()).decode()
	return f'data:image/jpeg;base64,{b64}'


	def ask_lvlm(llm_client, image: Image.Image, question: str,
	model: Optional[str] = None, system_prompt: Optional[str] = None,
	context: Optional[str] = None, max_tokens: int = 512) -> str:
	"""
	Send an image + question to the vLLM-hosted LVLM and return the answer.
	llm_client: openai.OpenAI instance pointed at the vLLM endpoint.
	model: model ID string (defaults to PRIMARY_MODEL env var or Qwen2.5-VL).
	context: optional retrieved caption to prepend as context.
	"""
	model = model or os.getenv('VLLM_MODEL', 'google/gemma-4-31B-it')
	image_url = image_to_base64(image)

	user_content = []
	if context:
	user_content.append({'type': 'text', 'text': f'Context: {context}\n\n'})
	user_content.append({'type': 'image_url', 'image_url': {'url': image_url}})
	user_content.append({'type': 'text', 'text': question})

	messages = []
	if system_prompt:
	messages.append({'role': 'system', 'content': system_prompt})
	messages.append({'role': 'user', 'content': user_content})

	response = llm_client.chat.completions.create(
	model=model,
	messages=messages,
	max_tokens=max_tokens,
	)
	return response.choices[0].message.content


	def load_siglip_model(model_name: str = 'google/siglip-base-patch16-224'):
	"""Load SigLIP processor and model. Returns (processor, model, device)."""
	import torch
	from transformers import AutoProcessor, AutoModel

	device = 'cuda' if torch.cuda.is_available() else 'cpu'
	processor = AutoProcessor.from_pretrained(model_name)
	model = AutoModel.from_pretrained(model_name).to(device)
	model.eval()
	return processor, model, device


	def _as_tensor(out):
	"""Coerce a model output to the feature tensor (newer transformers may wrap it)."""
	import torch
	if torch.is_tensor(out):
	return out
	if getattr(out, "pooler_output", None) is not None:
	return out.pooler_output
	if getattr(out, "last_hidden_state", None) is not None:
	return out.last_hidden_state.mean(dim=1)
	raise TypeError(f"Unexpected SigLIP feature output: {type(out)}")


	def embed_text_siglip(text: str, processor, model, device) -> list:
	"""Embed a text query with SigLIP for cross-modal search.

	SigLIP's text encoder is trained with fixed 64-token padding; using
	padding='max_length' (not dynamic) is required for correct embeddings.
	"""
	import torch
	inputs = processor(text=[text], return_tensors='pt', padding='max_length',
	truncation=True, max_length=64).to(device)
	with torch.no_grad():
	feat = _as_tensor(model.get_text_features(**inputs))
	feat = feat / feat.norm(dim=-1, keepdim=True)
	return feat.cpu().numpy()[0].tolist()


	def embed_image_siglip(img: Image.Image, processor, model, device) -> list:
	"""Embed a PIL image with SigLIP."""
	import torch
	inputs = processor(images=img, return_tensors='pt').to(device)
	with torch.no_grad():
	feat = _as_tensor(model.get_image_features(**inputs))
	feat = feat / feat.norm(dim=-1, keepdim=True)
	return feat.cpu().numpy()[0].tolist()


	def retrieve_by_text_bm25(client, query: str, index_name: str, top_k: int = 3) -> list:
	"""Retrieve image-caption pairs using BM25 keyword search."""
	resp = client.search(
	index=index_name,
	body={'size': top_k, 'query': {'match': {'caption': query}}}
	)
	return [h['_source'] for h in resp['hits']['hits']]


	def retrieve_by_text_knn(client, query: str, index_name: str,
	bge_model, top_k: int = 3) -> list:
	"""Retrieve image-caption pairs using BGE semantic k-NN."""
	vec = bge_model.encode(query, normalize_embeddings=True).tolist()
	resp = client.search(
	index=index_name,
	body={'size': top_k, 'query': {'knn': {'caption_vec': {'vector': vec, 'k': top_k}}}}
	)
	return [h['_source'] for h in resp['hits']['hits']]


	def retrieve_by_text_crossmodal(client, query: str, index_name: str,
	siglip_processor, siglip_model, siglip_device,
	top_k: int = 3) -> list:
	"""Retrieve images using SigLIP cross-modal k-NN (text → image vectors)."""
	vec = embed_text_siglip(query, siglip_processor, siglip_model, siglip_device)
	resp = client.search(
	index=index_name,
	body={'size': top_k, 'query': {'knn': {'image_vec': {'vector': vec, 'k': top_k}}}}
	)
	return [h['_source'] for h in resp['hits']['hits']]


	def retrieve_by_image_crossmodal(client, img: Image.Image, index_name: str,
	siglip_processor, siglip_model, siglip_device,
	exclude_id: Optional[str] = None, top_k: int = 3) -> list:
	"""Retrieve similar images using SigLIP image embedding."""
	vec = embed_image_siglip(img, siglip_processor, siglip_model, siglip_device)
	resp = client.search(
	index=index_name,
	body={'size': top_k + 1, 'query': {'knn': {'image_vec': {'vector': vec, 'k': top_k + 1}}}}
	)
	results = [
	h['_source'] for h in resp['hits']['hits']
	if h['_source']['image_id'] != exclude_id
	]
	return results[:top_k]


	def rag_answer(llm_client, os_client, question: str,
	image_lookup: dict, bge_model,
	siglip_processor, siglip_model, siglip_device,
	index_bge: str, index_multi: str,
	retrieval_method: str = 'crossmodal',
	top_k: int = 1, use_caption_context: bool = True,
	model: Optional[str] = None) -> dict:
	"""
	Full RAG pipeline: retrieve relevant image-caption, then generate answer with LVLM.
	retrieval_method: 'bm25' \| 'knn' \| 'crossmodal'
	"""
	if retrieval_method == 'bm25':
	hits = retrieve_by_text_bm25(os_client, question, index_bge, top_k)
	elif retrieval_method == 'knn':
	hits = retrieve_by_text_knn(os_client, question, index_bge, bge_model, top_k)
	else:
	hits = retrieve_by_text_crossmodal(
	os_client, question, index_multi,
	siglip_processor, siglip_model, siglip_device, top_k
	)

	if not hits:
	return {'question': question, 'answer': 'No relevant images found.', 'hits': []}

	top_hit = hits[0]
	image_id = top_hit['image_id']
	caption: str = top_hit.get('caption', '')
	img = image_lookup.get(image_id)

	if img is None:
	return {'question': question, 'answer': f'Image {image_id} not found in dataset.', 'hits': hits}

	answer = ask_lvlm(
	llm_client, img, question,
	model=model,
	context=caption if use_caption_context else None,
	)
	return {
	'question': question,
	'retrieved_image_id': image_id,
	'retrieved_caption': caption,
	'answer': answer,
	'hits': hits,
	}