wehub-resource-sync
223 lines
7.3 KiB
Python
223 lines
7.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Verify fine-tuned embeddings: positives should be closer than negatives.
|
|
|
|
Compares base model vs fine-tuned (LoRA) model on eval pairs.
|
|
For each query, computes cosine similarity to its positive image and to
|
|
all other images (negatives). Reports mean positive/negative similarity
|
|
and Recall@1/5/10.
|
|
|
|
Usage:
|
|
python training/verify_embeddings.py --adapter training/output_test
|
|
python training/verify_embeddings.py --adapter training/output_test --max-pairs 50
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
stream=sys.stdout,
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def load_model_and_processor(model_name, adapter_path=None, max_visual_tokens=256):
|
|
from models.biqwen3 import BiQwen3
|
|
from transformers import AutoProcessor
|
|
|
|
model = BiQwen3.from_pretrained(model_name, dtype=torch.bfloat16)
|
|
processor = AutoProcessor.from_pretrained(model_name)
|
|
patch_size = processor.image_processor.patch_size
|
|
merge_size = processor.image_processor.merge_size
|
|
tile = patch_size * merge_size
|
|
processor.image_processor.max_pixels = max_visual_tokens * tile * tile
|
|
processor.image_processor.size["longest_edge"] = (
|
|
processor.image_processor.max_pixels
|
|
)
|
|
processor.tokenizer.padding_side = "left"
|
|
|
|
if adapter_path:
|
|
from peft import PeftModel
|
|
|
|
model = PeftModel.from_pretrained(model, adapter_path)
|
|
logger.info(f"Loaded LoRA adapter from {adapter_path}")
|
|
|
|
model.eval()
|
|
model.cuda()
|
|
return model, processor
|
|
|
|
|
|
# Task-specific instructions (Qwen3-VL-Embedding style)
|
|
QUERY_INSTRUCTION = "Retrieve images or text relevant to the user's query."
|
|
DOC_INSTRUCTION = "Represent the user's input."
|
|
|
|
|
|
def _process_queries(processor, queries):
|
|
messages_batch = [
|
|
[
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": QUERY_INSTRUCTION}],
|
|
},
|
|
{"role": "user", "content": [{"type": "text", "text": q}]},
|
|
]
|
|
for q in queries
|
|
]
|
|
texts = [
|
|
processor.apply_chat_template(m, tokenize=False, add_generation_prompt=True)
|
|
for m in messages_batch
|
|
]
|
|
return processor(text=texts, return_tensors="pt", padding="longest")
|
|
|
|
|
|
def _process_doc_images(processor, images):
|
|
messages_batch = [
|
|
[
|
|
{"role": "system", "content": [{"type": "text", "text": DOC_INSTRUCTION}]},
|
|
{"role": "user", "content": [{"type": "image", "image": img}]},
|
|
]
|
|
for img in images
|
|
]
|
|
texts = [
|
|
processor.apply_chat_template(m, tokenize=False, add_generation_prompt=True)
|
|
for m in messages_batch
|
|
]
|
|
return processor(text=texts, images=images, return_tensors="pt", padding="longest")
|
|
|
|
|
|
def embed_queries(model, processor, queries, batch_size=16):
|
|
all_embs = []
|
|
for i in range(0, len(queries), batch_size):
|
|
batch = queries[i : i + batch_size]
|
|
inputs = _process_queries(processor, batch)
|
|
inputs = {k: v.cuda() for k, v in inputs.items()}
|
|
with torch.no_grad():
|
|
embs = model(**inputs)
|
|
all_embs.append(embs.cpu().float().numpy())
|
|
return np.concatenate(all_embs, axis=0)
|
|
|
|
|
|
def embed_images(model, processor, images, batch_size=8):
|
|
all_embs = []
|
|
for i in range(0, len(images), batch_size):
|
|
batch = images[i : i + batch_size]
|
|
inputs = _process_doc_images(processor, batch)
|
|
inputs = {k: v.cuda() for k, v in inputs.items()}
|
|
with torch.no_grad():
|
|
embs = model(**inputs)
|
|
all_embs.append(embs.cpu().float().numpy())
|
|
return np.concatenate(all_embs, axis=0)
|
|
|
|
|
|
def compute_metrics(q_embs, i_embs):
|
|
"""Compute retrieval metrics. Each query[i] matches image[i]."""
|
|
# Cosine similarity matrix (already L2 normalized by model)
|
|
sims = q_embs @ i_embs.T # (N, N)
|
|
n = sims.shape[0]
|
|
|
|
# Per-query: positive sim vs mean negative sim
|
|
pos_sims = np.diag(sims)
|
|
# Mask out diagonal for negative sims
|
|
mask = ~np.eye(n, dtype=bool)
|
|
neg_sims = sims[mask].reshape(n, n - 1)
|
|
mean_neg_sims = neg_sims.mean(axis=1)
|
|
max_neg_sims = neg_sims.max(axis=1)
|
|
|
|
# Recall@K
|
|
rankings = (-sims).argsort(axis=1)
|
|
correct = np.arange(n)
|
|
ranks = np.array([np.where(rankings[i] == correct[i])[0][0] for i in range(n)])
|
|
|
|
recall_1 = (ranks < 1).mean()
|
|
recall_5 = (ranks < 5).mean()
|
|
recall_10 = (ranks < 10).mean()
|
|
mrr = (1.0 / (ranks + 1)).mean()
|
|
|
|
return {
|
|
"mean_pos_sim": float(pos_sims.mean()),
|
|
"mean_neg_sim": float(mean_neg_sims.mean()),
|
|
"mean_max_neg_sim": float(max_neg_sims.mean()),
|
|
"margin": float((pos_sims - mean_neg_sims).mean()),
|
|
"recall@1": float(recall_1),
|
|
"recall@5": float(recall_5),
|
|
"recall@10": float(recall_10),
|
|
"mrr": float(mrr),
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", default="Qwen/Qwen3-VL-Embedding-2B")
|
|
parser.add_argument(
|
|
"--adapter", type=str, required=True, help="Path to LoRA adapter dir"
|
|
)
|
|
parser.add_argument("--eval-jsonl", default="training/data/eval.jsonl")
|
|
parser.add_argument("--max-pairs", type=int, default=100)
|
|
parser.add_argument("--max-visual-tokens", type=int, default=1024)
|
|
args = parser.parse_args()
|
|
|
|
# Load eval data
|
|
pairs = []
|
|
with open(args.eval_jsonl) as f:
|
|
for line in f:
|
|
item = json.loads(line)
|
|
try:
|
|
img = Image.open(item["chunk_path"]).convert("RGB")
|
|
pairs.append((item["query"], img))
|
|
except Exception:
|
|
continue
|
|
pairs = pairs[: args.max_pairs]
|
|
queries = [p[0] for p in pairs]
|
|
images = [p[1] for p in pairs]
|
|
logger.info(f"Loaded {len(pairs)} eval pairs")
|
|
|
|
results = {}
|
|
|
|
# 1. Base model (no LoRA)
|
|
logger.info("=== Base model (no fine-tuning) ===")
|
|
model, processor = load_model_and_processor(
|
|
args.model, adapter_path=None, max_visual_tokens=args.max_visual_tokens
|
|
)
|
|
q_embs = embed_queries(model, processor, queries)
|
|
i_embs = embed_images(model, processor, images)
|
|
base_metrics = compute_metrics(q_embs, i_embs)
|
|
results["base"] = base_metrics
|
|
for k, v in base_metrics.items():
|
|
logger.info(f" {k}: {v:.4f}")
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
|
|
# 2. Fine-tuned model (with LoRA)
|
|
logger.info(f"=== Fine-tuned model ({args.adapter}) ===")
|
|
model, processor = load_model_and_processor(
|
|
args.model, adapter_path=args.adapter, max_visual_tokens=args.max_visual_tokens
|
|
)
|
|
q_embs_ft = embed_queries(model, processor, queries)
|
|
i_embs_ft = embed_images(model, processor, images)
|
|
ft_metrics = compute_metrics(q_embs_ft, i_embs_ft)
|
|
results["finetuned"] = ft_metrics
|
|
for k, v in ft_metrics.items():
|
|
logger.info(f" {k}: {v:.4f}")
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
|
|
# 3. Comparison
|
|
logger.info("=== Comparison ===")
|
|
for k in base_metrics:
|
|
diff = ft_metrics[k] - base_metrics[k]
|
|
arrow = "↑" if diff > 0 else "↓" if diff < 0 else "="
|
|
logger.info(
|
|
f" {k}: {base_metrics[k]:.4f} → {ft_metrics[k]:.4f} ({arrow}{abs(diff):.4f})"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|