Views
No views yet
git clone https://github.com/cycraft-corp/RerAnchor-for-Visual-Document-Grounding.git1import torch
2from PIL import Image
3from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor
4from reranchor_lib import Qwen2_5_RerAnchor, ReranchorProcessor, denoise_screenshot
5
6# ----------------------
7# Load models
8# ----------------------
9model = ColQwen2_5.from_pretrained(
10 "vidore/colqwen2.5-v0.2",
11 device_map="cuda:0"
12).eval()
13
14processor = ColQwen2_5_Processor.from_pretrained(
15 "vidore/colqwen2.5-v0.2"
16)
17
18rerank_model = Qwen2_5_RerAnchor.from_pretrained(
19 "ricky42613/reranchor-qwen2.5-3b",
20 device_map="cuda:0",
21 torch_dtype=torch.bfloat16
22).eval()
23
24rerank_processor = ReranchorProcessor.from_pretrained(
25 "Qwen/Qwen2.5-VL-3B-Instruct"
26)
27
28# ----------------------
29# Input
30# ----------------------
31query = "What is the main finding?"
32image = Image.open("page.png")
33
34# ----------------------
35# Step 1: encode query
36# ----------------------
37query_inputs = processor.process_queries([query]).to(model.device)
38with torch.no_grad():
39 query_embed = model(**query_inputs)
40
41# ----------------------
42# Step 2: denoise image (RerAnchor)
43# ----------------------
44denoised_image = denoise_screenshot(
45 rerank_processor,
46 rerank_model,
47 query,
48 image,
49 k_tokens=200
50)
51
52# ----------------------
53# Step 3: encode image
54# ----------------------
55image_inputs = processor.process_images([denoised_image]).to(model.device)
56with torch.no_grad():
57 image_embed = model(**image_inputs)
58
59# ----------------------
60# Step 4: similarity score
61# ----------------------
62score = processor.score_multi_vector(query_embed, image_embed)
63
64print("Similarity score:", score.item())