1import fitz # PyMuPDF
2import torch
3from PIL import Image
4from transformers import AutoProcessor
5from transformers.models.qwen3_vl import Qwen3VLForConditionalGeneration
6
7MODEL_ID = "mtri-admin/ZipRerank"
8
9processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
10model = Qwen3VLForConditionalGeneration.from_pretrained(
11 MODEL_ID,
12 torch_dtype=torch.bfloat16,
13 device_map="auto",
14 attn_implementation="flash_attention_2", # or "sdpa" if flash-attn is unavailable
15 trust_remote_code=True,
16).eval()
17tokenizer = processor.tokenizer
18
19
20def pdf_to_images(pdf_path: str, max_size: int = 1024):
21 """Render every page so the longest edge is at most ``max_size`` pixels."""
22 doc = fitz.open(pdf_path)
23 images = []
24 for page in doc:
25 scale = max_size / max(page.rect.width, page.rect.height)
26 pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale))
27 images.append(Image.frombytes("RGB", (pix.width, pix.height), pix.samples))
28 doc.close()
29 return images
30
31
32def create_ranking_prompt(query: str, num_passages: int) -> str:
33 lines = [
34 "You are RankGPT, an intelligent assistant that can rank passages "
35 "based on their relevancy to the query.",
36 "",
37 f"I will provide you with {num_passages} passages as images.",
38 "Rank the passages based on their relevance to the search query.",
39 "",
40 "The images are provided in order: "
41 + ", ".join(
42 f"Picture {i + 1} is passage [{chr(ord('A') + i)}]"
43 for i in range(num_passages)
44 )
45 + ".",
46 "",
47 f"Search Query: {query}",
48 "",
49 "Rank the passages above based on their relevance to the search query.",
50 "The passages should be listed in descending order using identifiers.",
51 "The most relevant passages should be listed first.",
52 "The output format should be [A] > [B], etc.",
53 "Only output the ranking results, do not say anything else.",
54 ]
55 return "
56".join(lines)
57
58
59@torch.no_grad()
60def rerank_window(query: str, images):
61 """Rank up to 20 page images in a single forward pass.
62
63 Returns a list of 0-based indices into ``images``, ordered best-first.
64 """
65 assert 1 <= len(images) <= 20, "Window size must be between 1 and 20."
66 messages = [{
67 "role": "user",
68 "content": [{"type": "text", "text": create_ranking_prompt(query, len(images))}]
69 + [{"type": "image", "image": img} for img in images],
70 }]
71 inputs = processor.apply_chat_template(
72 [messages],
73 tokenize=True,
74 add_generation_prompt=True,
75 return_dict=True,
76 return_tensors="pt",
77 )
78 # Force the first predicted token to be a letter by appending "["
79 prompt_ids = inputs["input_ids"][0].tolist()
80 prompt_ids.append(tokenizer.encode("[", add_special_tokens=False)[0])
81 input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=model.device)
82
83 logits = model(
84 input_ids=input_ids,
85 attention_mask=torch.ones_like(input_ids),
86 pixel_values=inputs["pixel_values"].to(model.device),
87 image_grid_thw=inputs["image_grid_thw"].to(model.device),
88 ).logits[0, -1, :]
89
90 letter_ids = [
91 tokenizer.encode(chr(ord("A") + i), add_special_tokens=False)[0]
92 for i in range(len(images))
93 ]
94 scores = [logits[tid].item() for tid in letter_ids]
95 return sorted(range(len(images)), key=lambda i: scores[i], reverse=True)
96
97
98pages = pdf_to_images("report.pdf", max_size=1024)
99ranking = rerank_window("What is the company revenue?", pages[:20])
100print("Best-first page indices:", ranking)