1from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
2from PIL import Image
3import torch
4import math
5
6# more pixels -> better embeddings -> more VRAM -> slower inference
7# From my experience, 768 image patches is the right spot for compute efficient embeddings.
8max_pixels = 768 * 28 * 28
9min_pixels = 1 * 28 * 28
10
11# Load the embedding model and processor
12model = Qwen2VLForConditionalGeneration.from_pretrained(
13 'llamaindex/vdr-2b-multi-v1',
14 # These are the recommended kwargs for the model, but change them as needed
15 attn_implementation="flash_attention_2",
16 torch_dtype=torch.bfloat16,
17 device_map="cuda:0"
18).eval()
19
20processor = AutoProcessor.from_pretrained(
21 'llamaindex/vdr-2b-multi-v1',
22 min_pixels=min_pixels,
23 max_pixels=max_pixels
24)
25
26model.padding_side = "left"
27processor.tokenizer.padding_side = "left"
28
29document_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is shown in this image?<|im_end|>\n<|endoftext|>"
30
31query_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Query: %s<|im_end|>\n<|endoftext|>"
1def encode_queries(queries: list[str], dimension: int) -> torch.Tensor:
2 """
3 Encode a list of queries into a tensor of embeddings.
4
5 Args:
6 queries: A list of strings, each representing a query.
7 dimension: The desired dimension of the output embeddings.
8
9 Returns:
10 A tensor of shape (num_queries, dimension) containing the encoded queries.
11 """
12
13 dummy_image = Image.new('RGB', (56, 56))
14 inputs = processor(
15 text=[query_prompt % x for x in queries],
16 images=[dummy_image for _ in queries],
17 videos=None,
18 padding='longest',
19 return_tensors='pt'
20 ).to('cuda:0')
21
22 cache_position = torch.arange(0, len(queries))
23 inputs = model.prepare_inputs_for_generation(
24 **inputs, cache_position=cache_position, use_cache=False)
25
26 with torch.no_grad():
27 output = self.model(
28 **inputs,
29 return_dict=True,
30 output_hidden_states=True
31 )
32
33 embeddings = output.hidden_states[-1][:, -1]
34 return torch.nn.functional.normalize(embeddings[:, :dimension], p=2, dim=-1)
1def round_by_factor(number: float, factor: int) -> int:
2 return round(number / factor) * factor
3
4def ceil_by_factor(number: float, factor: int) -> int:
5 return math.ceil(number / factor) * factor
6
7def floor_by_factor(number: float, factor: int) -> int:
8 return math.floor(number / factor) * factor
9
10def smart_resize(height: int, width: int) -> tuple[int, int]:
11 h_bar = max(28, round_by_factor(height, 28))
12 w_bar = max(28, round_by_factor(width, 28))
13 if h_bar * w_bar > max_pixels:
14 beta = math.sqrt((height * width) / max_pixels)
15 h_bar = floor_by_factor(height / beta, 28)
16 w_bar = floor_by_factor(width / beta, 28)
17 elif h_bar * w_bar < min_pixels:
18 beta = math.sqrt(min_pixels / (height * width))
19 h_bar = ceil_by_factor(height * beta, 28)
20 w_bar = ceil_by_factor(width * beta, 28)
21 return w_bar, h_bar
22
23def resize(image: Image.Image):
24 new_size = smart_resize(image.height, image.width)
25 return image.resize(new_size)
26
27def encode_documents(documents: list[Image.Image], dimension: int):
28 """
29 Encode a list of images into a tensor of embeddings.
30
31 Args:
32 documents: A list of PIL Image objects.
33 dimension: The desired dimension of the output embeddings.
34
35 Returns:
36 A tensor of shape (num_documents, dimension) containing the encoded images.
37 """
38
39 inputs = processor(
40 text=[document_prompt] * len(documents),
41 images=[resize(x) for x in documents],
42 videos=None,
43 padding='longest',
44 return_tensors='pt'
45 ).to('cuda:0')
46
47 cache_position = torch.arange(0, len(queries))
48 inputs = model.prepare_inputs_for_generation(
49 **inputs, cache_position=cache_position, use_cache=False)
50
51 with torch.no_grad():
52 output = self.model(
53 **inputs,
54 return_dict=True,
55 output_hidden_states=True
56 )
57
58 embeddings = output.hidden_states[-1][:, -1]
59 return torch.nn.functional.normalize(embeddings[:, :dimension], p=2, dim=-1)
