src/vlps/data.py

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