1import torch
2from transformers import AutoProcessor
3from qwen_vl_utils import process_vision_info
4
5from src.arguments import ModelArguments
6from src.encoder.select_encoder import AttentionSelectEncoder
7from src.models.qwen2_5_vl_embed.qwen2_5_vl_embed import Qwen2_5ForEmbedding
8from src.utils import get_appending_token_strings
9
10MODEL_ID = "hltcoe/AGC_qwen2.5-vl_msrvtt"
11VIDEO_PATH = "PLACEHOLDER"
12NUM_PROXY_TOKENS = 32
13APPENDING_SUFFIX = "".join(get_appending_token_strings(NUM_PROXY_TOKENS))
14
15# --- Setup ---
16model_args = ModelArguments(
17 model_name_or_path=MODEL_ID,
18 pooling="select",
19 normalize=True,
20 num_appending_token=NUM_PROXY_TOKENS,
21 use_cluster_pooling=True,
22 use_attn_weight_cluster_pooling=True,
23 attn_implementation="flash_attention_2",
24)
25
26processor = AutoProcessor.from_pretrained(MODEL_ID)
27model = AttentionSelectEncoder.load(
28 Qwen2_5ForEmbedding,
29 model_args,
30 attn_implementation=model_args.attn_implementation,
31 dtype=torch.bfloat16,
32)
33model = model.to("cuda").eval()
34
35# --- Encode a video document ---
36passage_messages = [
37 {
38 "role": "user",
39 "content": [
40 {"type": "text", "text": "Passage: "},
41 {"type": "video", "video": VIDEO_PATH, "nframes": 24, "max_pixels": 84672, "min_pixels": 75264},
42 ],
43 }
44]
45text = processor.apply_chat_template(passage_messages, tokenize=False, add_generation_prompt=False)
46text += APPENDING_SUFFIX
47image_inputs, video_inputs = process_vision_info(passage_messages)
48passage_inputs = processor(
49 text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt",
50).to("cuda")
51
52with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
53 with torch.inference_mode():
54 doc_embeddings, doc_mask = model.encode(passage_inputs, is_query=False)
55 print(doc_embeddings.shape)
56 # doc_embeddings: (1, 32, 2048) — 32 compressed AGC vectors
57
58# --- Encode a text query ---
59query_messages = [{"role": "user", "content": [{"type": "text", "text": "Query: a person is cooking"}]}]
60query_text = processor.apply_chat_template(query_messages, tokenize=False, add_generation_prompt=False)
61query_inputs = processor(text=[query_text], padding=True, return_tensors="pt").to("cuda")
62
63with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
64 with torch.inference_mode():
65 query_embeddings, query_mask = model.encode(query_inputs, is_query=True)
66 print(query_embeddings.shape)
67
68# --- ColBERT MaxSim scoring ---
69score = model.compute_similarity(query_embeddings, doc_embeddings, query_mask, doc_mask)
70print(f"Similarity score: {score.item():.4f}")