#!/usr/bin/env python
"""Compute UEmbed (Qwen3.5) embeddings with vLLM's pooling engine.

vLLM expects the backbone weights under a `model.` prefix, whereas the released
checkpoint stores them as `language_model.*` / `visual.*`. `prepare_vllm_model_dir`
remaps the keys once into a cached directory and returns the path to feed vLLM.
This script shows both dense (`last.normal`) and sparse (`splade.last`)
embeddings for text- and image-containing inputs.

Usage:
    python examples/vllm_example.py ./models/UEmbed-2B
"""

import argparse
import hashlib
import json
import os
import tempfile

import torch
import torch.nn.functional as F
from transformers import AutoProcessor
from vllm import LLM, PoolingParams
from vllm.config.pooler import PoolerConfig
from qwen_vl_utils.vision_process import process_vision_info

NUM_EOS_TOKENS = 16  # from sparse_info.json


def _remap_key(key: str) -> str:
    if key.startswith(("language_model.", "visual.")):
        return f"model.{key}"
    return key


def _needs_key_remap(model_dir: str) -> bool:
    index_path = os.path.join(model_dir, "model.safetensors.index.json")
    if os.path.exists(index_path):
        with open(index_path, encoding="utf-8") as handle:
            first_key = next(iter(json.load(handle).get("weight_map", {})), "")
    else:
        weight_path = os.path.join(model_dir, "model.safetensors")
        if not os.path.exists(weight_path):
            return False
        from safetensors import safe_open
        with safe_open(weight_path, framework="pt", device="cpu") as handle:
            first_key = next(iter(handle.keys()), "")
    return first_key.startswith(("language_model.", "visual."))


def prepare_vllm_model_dir(model_dir: str) -> str:
    """Remap `language_model.*` / `visual.*` keys under a `model.` prefix (once, cached)."""
    if not _needs_key_remap(model_dir):
        return model_dir

    from safetensors.torch import load_file, save_file

    fingerprint = hashlib.sha256(os.path.realpath(model_dir).encode()).hexdigest()[:20]
    cache_dir = os.path.join(
        os.environ.get("UEMBED_VLLM_MODEL_CACHE", os.path.join(tempfile.gettempdir(), "uembed_vllm_models")),
        fingerprint,
    )
    stamp_file = os.path.join(cache_dir, ".done")
    if os.path.exists(stamp_file):
        return cache_dir

    os.makedirs(cache_dir, exist_ok=True)
    index_path = os.path.join(model_dir, "model.safetensors.index.json")
    if os.path.exists(index_path):
        with open(index_path, encoding="utf-8") as handle:
            index = json.load(handle)
        new_weight_map = {_remap_key(k): v for k, v in index["weight_map"].items()}
        files_to_convert = set(index["weight_map"].values())
        with open(os.path.join(cache_dir, "model.safetensors.index.json"), "w", encoding="utf-8") as handle:
            json.dump({**index, "weight_map": new_weight_map}, handle, indent=2)
    else:
        files_to_convert = {"model.safetensors"}

    for filename in files_to_convert:
        tensors = load_file(os.path.join(model_dir, filename))
        save_file({_remap_key(k): v for k, v in tensors.items()}, os.path.join(cache_dir, filename))

    for name in os.listdir(model_dir):
        if name in files_to_convert or name == "model.safetensors.index.json":
            continue
        destination = os.path.join(cache_dir, name)
        if not os.path.lexists(destination):
            os.symlink(os.path.join(model_dir, name), destination)
    open(stamp_file, "w").close()
    return cache_dir

CONVERSATIONS = [
    [{"role": "user", "content": [
        {"type": "text", "text": "A woman playing with her dog on a beach at sunset."},
    ]}],
    [{"role": "user", "content": [
        {"type": "image", "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"},
        {"type": "text", "text": "A woman and her dog on the beach."},
    ]}],
]


def build_vllm_inputs(processor, conversations):
    texts = processor.apply_chat_template(conversations, add_generation_prompt=True, tokenize=False)
    vllm_inputs = []
    for text, conv in zip(texts, conversations):
        item = {"prompt": text}
        images, videos, _ = process_vision_info(
            [conv], image_patch_size=16, return_video_metadata=True, return_video_kwargs=True
        )
        mm = {}
        if images:
            mm["image"] = images
        if videos:
            mm["video"] = [video for video, _ in videos]
        if mm:
            item["multi_modal_data"] = mm
        vllm_inputs.append(item)
    return vllm_inputs


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("model_path", nargs="?", default="./models/UEmbed-2B")
    args = parser.parse_args()

    model_dir = prepare_vllm_model_dir(args.model_path)
    llm = LLM(
        model=model_dir,
        runner="pooling",
        pooler_config=PoolerConfig(task="token_embed", pooling_type="ALL"),
        dtype=torch.bfloat16,
        max_model_len=8192,
        gpu_memory_utilization=0.8,
        limit_mm_per_prompt={"image": 1, "video": 1},
    )
    processor = AutoProcessor.from_pretrained(model_dir, padding_side="right")

    outputs = llm.encode(
        build_vllm_inputs(processor, CONVERSATIONS),
        pooling_task="token_embed",
        pooling_params=PoolingParams(use_activation=False),
        use_tqdm=False,
    )

    device = outputs[0].outputs.data.device
    sw = torch.load(f"{model_dir}/sparse_weights.pt", map_location=device, weights_only=True)
    dense_list, sparse_list = [], []
    for out in outputs:
        hidden = out.outputs.data  # [seq_len, hidden], per-token states

        # Dense: hidden state of the EOS right before the appended sparse tokens.
        dense_list.append(F.normalize(hidden[-(NUM_EOS_TOKENS + 1)], p=2, dim=-1))

        # Sparse: project the N appended EOS hidden states with the per-cluster heads.
        logits = [
            F.linear(hidden[-(NUM_EOS_TOKENS - i)].to(head.dtype), head, bias)
            for i, (head, bias) in enumerate(zip(sw["sparse_lm_heads"], sw["sparse_bias"]))
        ]
        sparse_list.append(torch.log1p(F.relu(torch.cat(logits, dim=-1))))

    dense = torch.stack(dense_list)      # [B, 2048]
    sparse = torch.stack(sparse_list)    # [B, 184016]

    print(f"dense:  {tuple(dense.shape)}")
    print("dense cosine matrix:\n", (dense.float() @ dense.float().T).cpu().numpy().round(4))
    print(f"sparse: {tuple(sparse.shape)}  nnz={ (sparse > 0).sum(-1).tolist() }")


if __name__ == "__main__":
    main()