modeling_unlimitedocr.py
52.2 KB · 1299 lines · python Raw
1 from .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM
2 from .configuration_deepseek_v2 import DeepseekV2Config
3 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
4 from typing import List, Optional, Tuple, Union
5 from transformers.cache_utils import Cache
6 import requests
7 from PIL import Image, ImageOps, ImageDraw, ImageFont
8 from io import BytesIO
9 import torch
10 import torch.nn as nn
11 from torch.nn import CrossEntropyLoss
12 from torchvision import transforms
13 from torchvision.transforms.functional import InterpolationMode
14 import os
15 from .deepencoder import build_sam_vit_b, build_clip_l, MlpProjector
16 from addict import Dict
17 from transformers import TextStreamer
18 from .conversation import get_conv_template
19 from abc import ABC
20 import math
21 import re
22 from tqdm import tqdm
23 import numpy as np
24 import time
25
26
27 def load_image(image_path):
28
29 try:
30 image = Image.open(image_path)
31
32 corrected_image = ImageOps.exif_transpose(image)
33
34 return corrected_image
35
36 except Exception as e:
37 print(f"error: {e}")
38 try:
39 return Image.open(image_path)
40 except:
41 return None
42
43
44 def re_match(text):
45 ref_pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'
46 matches = re.findall(ref_pattern, text, re.DOTALL)
47
48 det_pattern = r'(<\|det\|>\s*([A-Za-z_][\w-]*)\s*(\[[^\]]+\])\s*<\|/det\|>)'
49 for full_match, label, box in re.findall(det_pattern, text, re.DOTALL):
50 matches.append((full_match, label, box))
51
52 mathes_image = []
53 mathes_other = []
54 for a_match in matches:
55 if a_match[1].strip() == 'image' or '<|ref|>image<|/ref|>' in a_match[0]:
56 mathes_image.append(a_match[0])
57 else:
58 mathes_other.append(a_match[0])
59 return matches, mathes_image, mathes_other
60
61
62 def extract_coordinates_and_label(ref_text, image_width, image_height):
63
64 try:
65 label_type = ref_text[1]
66 cor_list = eval(ref_text[2])
67 if cor_list and isinstance(cor_list[0], (int, float)):
68 cor_list = [cor_list]
69 except Exception as e:
70 print(e)
71 return None
72
73 return (label_type, cor_list)
74
75
76 def draw_bounding_boxes(image, refs, ouput_path, image_prefix=''):
77
78 image_width, image_height = image.size
79
80 img_draw = image.copy()
81 draw = ImageDraw.Draw(img_draw)
82
83 overlay = Image.new('RGBA', img_draw.size, (0, 0, 0, 0))
84 draw2 = ImageDraw.Draw(overlay)
85
86 # try:
87 # except IOError:
88 # try:
89 # font = ImageFont.truetype("DejaVuSans.ttf", 20)
90 # except IOError:
91 font = ImageFont.load_default()
92
93 img_idx = 0
94
95 for i, ref in enumerate(refs):
96 try:
97 result = extract_coordinates_and_label(ref, image_width, image_height)
98 if result:
99 label_type, points_list = result
100
101 color = (np.random.randint(0, 200), np.random.randint(0, 200), np.random.randint(0, 255))
102
103 color_a = color + (20, )
104 for points in points_list:
105 x1, y1, x2, y2 = points
106
107 x1 = int(x1 / 999 * image_width)
108 y1 = int(y1 / 999 * image_height)
109
110 x2 = int(x2 / 999 * image_width)
111 y2 = int(y2 / 999 * image_height)
112
113 if label_type == 'image':
114 try:
115 cropped = image.crop((x1, y1, x2, y2))
116 cropped.save(f"{ouput_path}/images/{image_prefix}{img_idx}.jpg")
117 except Exception as e:
118 print(e)
119 pass
120 img_idx += 1
121
122 try:
123 if label_type == 'title':
124 draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
125 draw2.rectangle([x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1)
126 else:
127 draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
128 draw2.rectangle([x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1)
129 text_x = x1
130 text_y = max(0, y1 - 15)
131
132
133 text_bbox = draw.textbbox((0, 0), label_type, font=font)
134 text_width = text_bbox[2] - text_bbox[0]
135 text_height = text_bbox[3] - text_bbox[1]
136 draw.rectangle([text_x, text_y, text_x + text_width, text_y + text_height],
137 fill=(255, 255, 255, 30))
138
139 draw.text((text_x, text_y), label_type, font=font, fill=color)
140 except:
141 pass
142 except:
143 continue
144 img_draw.paste(overlay, (0, 0), overlay)
145 return img_draw
146
147
148 def process_image_with_refs(image, ref_texts, output_path, image_prefix=''):
149
150 result_image = draw_bounding_boxes(image, ref_texts, output_path, image_prefix=image_prefix)
151
152 return result_image
153
154
155
156
157
158 def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
159 best_ratio_diff = float('inf')
160 best_ratio = (1, 1)
161 area = width * height
162 for ratio in target_ratios:
163 target_aspect_ratio = ratio[0] / ratio[1]
164 ratio_diff = abs(aspect_ratio - target_aspect_ratio)
165 if ratio_diff < best_ratio_diff:
166 best_ratio_diff = ratio_diff
167 best_ratio = ratio
168 elif ratio_diff == best_ratio_diff:
169 if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
170 best_ratio = ratio
171 # print(f'width: {width}, height: {height}, best_ratio: {best_ratio}')
172 return best_ratio
173
174
175 def dynamic_preprocess(image, min_num=2, max_num=32, image_size=640, use_thumbnail=False):
176 orig_width, orig_height = image.size
177 aspect_ratio = orig_width / orig_height
178
179 # calculate the existing image aspect ratio
180 target_ratios = set(
181 (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
182 i * j <= max_num and i * j >= min_num)
183 # print(target_ratios)
184 target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
185
186 # find the closest aspect ratio to the target
187 target_aspect_ratio = find_closest_aspect_ratio(
188 aspect_ratio, target_ratios, orig_width, orig_height, image_size)
189
190 # print(target_aspect_ratio)
191 # calculate the target width and height
192 target_width = image_size * target_aspect_ratio[0]
193 target_height = image_size * target_aspect_ratio[1]
194 blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
195
196 # resize the image
197 resized_img = image.resize((target_width, target_height))
198 processed_images = []
199 for i in range(blocks):
200 box = (
201 (i % (target_width // image_size)) * image_size,
202 (i // (target_width // image_size)) * image_size,
203 ((i % (target_width // image_size)) + 1) * image_size,
204 ((i // (target_width // image_size)) + 1) * image_size
205 )
206 # split the image
207 split_img = resized_img.crop(box)
208 processed_images.append(split_img)
209 assert len(processed_images) == blocks
210 if use_thumbnail and len(processed_images) != 1:
211 thumbnail_img = image.resize((image_size, image_size))
212 processed_images.append(thumbnail_img)
213 return processed_images, target_aspect_ratio
214
215
216
217 def normalize_transform(mean, std):
218 if mean is None and std is None:
219 transform = None
220 elif mean is None and std is not None:
221 mean = [0.] * len(std)
222 transform = transforms.Normalize(mean=mean, std=std)
223 elif mean is not None and std is None:
224 std = [1.] * len(mean)
225 transform = transforms.Normalize(mean=mean, std=std)
226 else:
227 transform = transforms.Normalize(mean=mean, std=std)
228
229 return transform
230
231
232
233 def format_messages(
234 conversations: List[Dict[str, str]],
235 sft_format: str = "deepseek",
236 system_prompt: str = "",
237 ):
238 """
239 Applies the SFT template to conversation.
240
241 Args:
242 conversations (List[Dict]): A List of messages.
243 sft_format (str, optional): The format of the SFT template to use. Defaults to "deepseek".
244 system_prompt (str, optional): The system prompt to use in the SFT template. Defaults to "".
245
246 Returns:
247 sft_prompt (str): The formatted text.
248 """
249
250 conv = get_conv_template(sft_format)
251 conv.set_system_message(system_prompt)
252 for message in conversations:
253 conv.append_message(message["role"], message["content"].strip())
254 sft_prompt = conv.get_prompt().strip()
255
256 return sft_prompt
257
258
259 def text_encode(tokenizer, text: str, bos: bool = True, eos: bool = False):
260 t = tokenizer.encode(text, add_special_tokens=False)
261 bos_id = 0
262 eos_id = 1
263 if bos:
264 t = [bos_id] + t
265 if eos:
266 t = t + [eos_id]
267
268 return t
269
270 def load_pil_images(conversations: List[Dict[str, str]]) -> List[Image.Image]:
271 """
272
273 Args:
274 conversations (List[Dict[str, str]]): the conversations with a list of messages. An example is :
275 [
276 {
277 "role": "User",
278 "content": "<image_placeholder>\nExtract all information from this image and convert them into markdown format.",
279 "images": ["./examples/table_datasets.png"]
280 },
281 {"role": "Assistant", "content": ""},
282 ]
283
284 Returns:
285 pil_images (List[PIL.Image.Image]): the list of PIL images.
286
287 """
288
289 pil_images = []
290
291 for message in conversations:
292 if "images" not in message:
293 continue
294
295 for image_path in message["images"]:
296 # print('----------------')
297 # print(image_path)
298 # print('----------------')
299 # exit()
300
301 # pil_img = Image.open(image_path)
302 pil_img = load_image(image_path)
303 pil_img = pil_img.convert("RGB")
304 pil_images.append(pil_img)
305
306 return pil_images
307
308
309 class BaseTransform(ABC):
310
311 def set_rng(self, *args, **kwargs):
312 pass
313
314 def __call__(self, *args, **kwargs) -> torch.Tensor:
315 pass
316
317 @property
318 def default_shape(self):
319 raise NotImplementedError
320
321
322 class BasicImageTransform(BaseTransform):
323 def __init__(
324 self,
325 mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
326 std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
327 normalize: bool = True
328 ):
329 self.mean = mean
330 self.std = std
331
332 transform_pipelines = [
333 transforms.ToTensor()
334 ]
335
336 normalize = normalize_transform(mean, std) if normalize else nn.Identity()
337 if normalize is not None:
338 transform_pipelines.append(normalize)
339
340 self.transform = transforms.Compose(transform_pipelines)
341
342 def __call__(self, x):
343 x = self.transform(x)
344 return x
345
346 class NoEOSTextStreamer(TextStreamer):
347 def on_finalized_text(self, text: str, stream_end: bool = False):
348
349 eos_text = self.tokenizer.decode([self.tokenizer.eos_token_id], skip_special_tokens=False)
350 text = text.replace(eos_text, "\n")
351 print(text, flush=True, end="")
352
353
354 class SlidingWindowNoRepeatNgramProcessor:
355 """Block n-gram repetitions within a sliding window.
356 Aligned with SGLang DeepseekOCRNoRepeatNGramLogitProcessor."""
357 def __init__(self, ngram_size, window, whitelist_token_ids=None):
358 self.ngram_size = ngram_size
359 self.window = window
360 self.whitelist = set(whitelist_token_ids) if whitelist_token_ids else set()
361
362 def __call__(self, input_ids, scores):
363 for batch_idx in range(input_ids.shape[0]):
364 sequence = input_ids[batch_idx].tolist()
365 if len(sequence) < self.ngram_size:
366 continue
367 search_start = max(0, len(sequence) - self.window)
368 search_end = len(sequence) - self.ngram_size + 1
369 if search_end <= search_start:
370 continue
371 if self.ngram_size > 1:
372 current_prefix = tuple(sequence[-(self.ngram_size - 1):])
373 else:
374 current_prefix = tuple()
375 banned = set()
376 for idx in range(search_start, search_end):
377 ngram = sequence[idx:idx + self.ngram_size]
378 if self.ngram_size == 1 or tuple(ngram[:-1]) == current_prefix:
379 banned.add(ngram[-1])
380 banned.difference_update(self.whitelist)
381 for token_id in banned:
382 scores[batch_idx, token_id] = float('-inf')
383 return scores
384
385
386 class TPSTextStreamer(TextStreamer):
387 """Streamer that prints TPS every `interval` tokens. Set interval=0 to disable."""
388 def __init__(self, tokenizer, interval=100, **kwargs):
389 super().__init__(tokenizer, **kwargs)
390 self.interval = interval
391 self.token_count = 0
392 self.start_time = None
393 self.start_token_count = 0
394 self.last_report_count = 0
395 self.last_report_time = None
396
397 def put(self, value):
398 import time
399 if hasattr(value, 'numel'):
400 self.token_count += value.numel()
401 else:
402 self.token_count += 1
403 # 第一次 put 时开始计时（跳过 prefill）
404 if self.start_time is None:
405 self.start_time = time.time()
406 self.last_report_time = self.start_time
407 self.last_report_count = self.token_count
408 self.start_token_count = self.token_count
409 super().put(value)
410 return
411 if self.interval > 0 and self.token_count - self.last_report_count >= self.interval:
412 now = time.time()
413 delta_tokens = self.token_count - self.last_report_count
414 delta_time = now - self.last_report_time
415 recent_tps = delta_tokens / delta_time if delta_time > 0 else 0
416 avg_tps = (self.token_count - self.start_token_count) / (now - self.start_time) if (now - self.start_time) > 0 else 0
417 print(f"\n[TPS] tokens={self.token_count}, recent={recent_tps:.1f} t/s, avg={avg_tps:.1f} t/s", flush=True)
418 self.last_report_count = self.token_count
419 self.last_report_time = now
420 super().put(value)
421
422 def on_finalized_text(self, text: str, stream_end: bool = False):
423 eos_text = self.tokenizer.decode([self.tokenizer.eos_token_id], skip_special_tokens=False)
424 text = text.replace(eos_text, "\n")
425 print(text, flush=True, end="")
426
427
428 class UnlimitedOCRConfig(DeepseekV2Config):
429 model_type = "unlimited-ocr"
430
431 class UnlimitedOCRModel(DeepseekV2Model):
432 config_class = UnlimitedOCRConfig
433
434 def __init__(self, config: DeepseekV2Config):
435 super(UnlimitedOCRModel, self).__init__(config)
436
437 self.sam_model = build_sam_vit_b()
438 self.vision_model = build_clip_l()
439 # self.conv_2 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=2, stride=2)
440 n_embed = 1280
441 self.projector = MlpProjector(Dict(projector_type="linear", input_dim=2048, n_embed=n_embed))
442 embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
443 self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
444 self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
445
446
447
448
449 def forward(
450 self,
451 input_ids: torch.LongTensor = None,
452 attention_mask: Optional[torch.Tensor] = None,
453 position_ids: Optional[torch.LongTensor] = None,
454 past_key_values: Optional[List[torch.FloatTensor]] = None,
455 inputs_embeds: Optional[torch.FloatTensor] = None,
456 use_cache: Optional[bool] = None,
457 output_attentions: Optional[bool] = None,
458 output_hidden_states: Optional[bool] = None,
459 images: Optional[torch.FloatTensor] = None,
460 images_seq_mask: Optional[torch.FloatTensor] = None,
461 images_spatial_crop: Optional[torch.FloatTensor] = None,
462 return_dict: Optional[bool] = None,
463 ) -> Union[Tuple, BaseModelOutputWithPast]:
464
465
466
467
468 if inputs_embeds is None:
469 # inputs_embeds = self.embed_tokens(input_ids)
470 inputs_embeds = self.get_input_embeddings()(input_ids)
471
472
473
474 sam_model = getattr(self, 'sam_model', None)
475 # sam_model = self.sam_model
476 vision_model = getattr(self, 'vision_model', None)
477
478
479
480 if sam_model is not None and images is not None and (input_ids.shape[1] != 1 or self.training) and torch.sum(images[0][1]).item() != 0:
481
482 idx = 0
483
484 # sam_model = torch.jit.script(sam_model)
485
486 # start_time = time.time()
487 for image, crop_shape in zip(images, images_spatial_crop):
488 images_in_this_batch = []
489
490 patches = image[0]
491 image_ori = image[1]
492
493 with torch.no_grad():
494 # with torch.inference_mode():
495
496 if torch.sum(patches).item() != 0:
497 # P, C, H, W = patches.shape
498 crop_flag = 1
499 local_features_1 = sam_model(patches)
500
501 local_features_2 = vision_model(patches, local_features_1)
502 # vit_time = time.time()
503 local_features = torch.cat((local_features_2[:, 1:], local_features_1.flatten(2).permute(0, 2, 1)), dim=-1)
504 local_features = self.projector(local_features)
505
506
507 global_features_1 = sam_model(image_ori)
508 global_features_2 = vision_model(image_ori, global_features_1)
509 global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1)
510 global_features = self.projector(global_features)
511
512 # print('=====================')
513 # print('BASE: ', global_features.shape)
514 # print('PATCHES: ', local_features.shape)
515 # print('=====================')
516
517 _, hw, n_dim = global_features.shape
518 h = w = int(hw ** 0.5)
519
520 _2, hw2, n_dim2 = local_features.shape
521 h2 = w2 = int(hw2 ** 0.5)
522
523 width_crop_num, height_crop_num = crop_shape[0], crop_shape[1]
524
525 global_features = global_features.view(h, w, n_dim)
526
527 global_features = torch.cat(
528 [global_features, self.image_newline[None, None, :].expand(h, 1, n_dim)], dim=1
529 )
530
531 global_features = global_features.view(-1, n_dim)
532
533
534 local_features = local_features.view(height_crop_num, width_crop_num, h2, w2, n_dim2).permute(0, 2, 1, 3, 4).reshape(height_crop_num*h2, width_crop_num*w2, n_dim2)
535 local_features = torch.cat(
536 [local_features, self.image_newline[None, None, :].expand(height_crop_num * h2, 1, n_dim2)], dim=1
537 )
538 local_features = local_features.view(-1, n_dim2)
539
540 global_local_features = torch.cat([local_features, global_features, self.view_seperator[None, :]], dim=0)
541 images_in_this_batch.append(global_local_features)
542
543 # end_time = time.time()
544
545 # print('sam: ', sam_time - start_time)
546 # print('vit: ', vit_time - sam_time)
547 # print('all: ', end_time - start_time)
548
549 # exit()
550
551 else:
552 # Handle single or multiple images in image_ori
553 num_imgs = image_ori.shape[0]
554 for img_idx in range(num_imgs):
555 single_img = image_ori[img_idx:img_idx+1] # [1, 3, H, W]
556 global_features_1 = sam_model(single_img)
557 global_features_2 = vision_model(single_img, global_features_1)
558 global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1)
559 global_features = self.projector(global_features)
560
561 _, hw, n_dim = global_features.shape
562 h = w = int(hw ** 0.5)
563
564 global_features = global_features.view(h, w, n_dim)
565
566 global_features = torch.cat(
567 [global_features, self.image_newline[None, None, :].expand(h, 1, n_dim)], dim=1
568 )
569
570 global_features = global_features.view(-1, n_dim)
571
572 global_local_features = torch.cat([global_features, self.view_seperator[None, :]], dim=0)
573 images_in_this_batch.append(global_local_features)
574
575
576 # print(inputs_embeds.shape)
577
578 if images_in_this_batch:
579 images_in_this_batch = torch.cat(images_in_this_batch, dim=0)
580 # exit()
581
582 inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)
583
584 idx += 1
585
586
587 return super(UnlimitedOCRModel, self).forward(
588 input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values,
589 inputs_embeds=inputs_embeds, use_cache=use_cache, position_ids = position_ids,
590 output_attentions=output_attentions, output_hidden_states=output_hidden_states,
591 return_dict=return_dict
592 )
593
594
595 class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
596
597 config_class = UnlimitedOCRConfig
598 # supports_gradient_checkpointing = True
599
600 def __init__(self, config):
601 super(DeepseekV2ForCausalLM, self).__init__(config)
602 self.model = UnlimitedOCRModel(config)
603
604 self.vocab_size = config.vocab_size
605
606 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
607
608 # self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
609
610 # Initialize weights and apply final processing
611 self.post_init()
612
613 def get_model(self):
614 return self.model
615
616
617 def forward(
618 self,
619 input_ids: torch.LongTensor = None,
620 attention_mask: Optional[torch.Tensor] = None,
621 position_ids: Optional[torch.LongTensor] = None,
622 past_key_values: Optional[List[torch.FloatTensor]] = None,
623 inputs_embeds: Optional[torch.FloatTensor] = None,
624 labels: Optional[torch.LongTensor] = None,
625 use_cache: Optional[bool] = None,
626 output_attentions: Optional[bool] = None,
627 output_hidden_states: Optional[bool] = None,
628 images: Optional[torch.FloatTensor] = None,
629 images_seq_mask: Optional[torch.FloatTensor] = None,
630 images_spatial_crop: Optional[torch.FloatTensor] = None,
631 return_dict: Optional[bool] = None,
632
633 ) -> Union[Tuple, CausalLMOutputWithPast]:
634 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
635 output_hidden_states = (
636 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
637 )
638 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
639
640
641
642 outputs = self.model(
643 input_ids=input_ids,
644 past_key_values=past_key_values,
645 attention_mask=attention_mask,
646 position_ids=position_ids,
647 inputs_embeds=inputs_embeds,
648 use_cache=use_cache,
649 output_attentions=output_attentions,
650 output_hidden_states=output_hidden_states,
651 images=images,
652 images_seq_mask = images_seq_mask,
653 images_spatial_crop = images_spatial_crop,
654 return_dict=return_dict
655
656 )
657
658
659
660 # print(transformer_outputs)
661
662 hidden_states = outputs[0]
663 logits = self.lm_head(hidden_states)
664 logits = logits.float()
665
666 # logits
667
668 loss = None
669 if labels is not None:
670 # Shift so that tokens < n predict n
671 shift_logits = logits[..., :-1, :].contiguous()
672 shift_labels = labels[..., 1:].contiguous()
673 # Flatten the tokens
674 loss_fct = CrossEntropyLoss()
675 shift_logits = shift_logits.view(-1, self.config.vocab_size)
676 shift_labels = shift_labels.view(-1)
677 # Enable model parallelism
678 shift_labels = shift_labels.to(shift_logits.device)
679 loss = loss_fct(shift_logits, shift_labels)
680
681 if not return_dict:
682 output = (logits,) + outputs[1:]
683 return (loss,) + output if loss is not None else output
684
685 return CausalLMOutputWithPast(
686 loss=loss,
687 logits=logits,
688 past_key_values=outputs.past_key_values,
689 hidden_states=outputs.hidden_states,
690 attentions=outputs.attentions,
691 )
692
693
694 def prepare_inputs_for_generation(
695 self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
696 ):
697 # Omit tokens covered by past_key_values
698 past_length = 0
699 if past_key_values is not None:
700 if isinstance(past_key_values, Cache):
701 cache_length = past_key_values.get_seq_length()
702 past_length = past_key_values.get_seq_length()
703 max_cache_length = getattr(past_key_values, 'get_max_length', lambda: None)()
704 else:
705 cache_length = past_length = past_key_values[0][0].shape[2]
706 max_cache_length = None
707
708 # Ring buffer: cache size is fixed, but we've processed more tokens.
709 # Always just take the last token for decode.
710 if hasattr(past_key_values, '_prefill_length') and past_length > 0:
711 input_ids = input_ids[:, -1:]
712 # Keep only the unprocessed tokens:
713 # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
714 # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
715 # input)
716 elif attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
717 input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
718 # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
719 # input_ids based on the past_length.
720 elif past_length < input_ids.shape[1]:
721 input_ids = input_ids[:, past_length:]
722 # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
723
724 # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
725 if (
726 max_cache_length is not None
727 and attention_mask is not None
728 and cache_length + input_ids.shape[1] > max_cache_length
729 ):
730 attention_mask = attention_mask[:, -max_cache_length:]
731
732 position_ids = kwargs.get("position_ids", None)
733 if attention_mask is not None and position_ids is None:
734 # create position_ids on the fly for batch generation
735 position_ids = attention_mask.long().cumsum(-1) - 1
736 position_ids.masked_fill_(attention_mask == 0, 1)
737 if past_key_values:
738 position_ids = position_ids[:, -input_ids.shape[1] :]
739
740 # if self.generation_config.cache_implementation == "static":
741 # # generation with static cache
742 # cache_position = kwargs.get("cache_position", None)
743 # if cache_position is None:
744 # past_length = 0
745 # else:
746 # past_length = cache_position[-1] + 1
747 # input_ids = input_ids[:, past_length:]
748 # position_ids = position_ids[:, past_length:]
749
750 # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
751 # same goes for position ids. Could also help with continued generation.
752 cache_position = torch.arange(past_length, past_length + position_ids.shape[-1], device=position_ids.device)
753
754 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
755 if inputs_embeds is not None and past_key_values is None:
756 model_inputs = {"inputs_embeds": inputs_embeds}
757 else:
758 model_inputs = {"input_ids": input_ids}
759
760 # Pass images only on prefill (cache empty or None)
761 _is_prefill = (past_key_values is None or
762 (isinstance(past_key_values, Cache) and past_key_values.get_seq_length() == 0))
763 model_inputs.update(
764 {
765 "position_ids": position_ids,
766 "past_key_values": past_key_values,
767 "use_cache": kwargs.get("use_cache"),
768 "attention_mask": attention_mask,
769 "images": kwargs.get("images", None) if _is_prefill else None,
770 "images_seq_mask": kwargs.get("images_seq_mask", None) if _is_prefill else None,
771 "images_spatial_crop": kwargs.get("images_spatial_crop", None) if _is_prefill else None,
772 }
773 )
774 return model_inputs
775
776
777 def disable_torch_init(self):
778 """
779 Disable the redundant torch default initialization to accelerate model creation.
780 """
781 import torch
782 setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
783 setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
784
785
786
787 def infer(self, tokenizer, prompt='', image_file='', output_path = '', base_size=1024, image_size=640, crop_mode=True, test_compress=False, save_results=False, eval_mode=False, max_length=32768, tps_interval=0, no_repeat_ngram_size=0, ngram_window=0, temperature=0.0):
788 self.disable_torch_init()
789
790 os.makedirs(output_path, exist_ok=True)
791 os.makedirs(f'{output_path}/images', exist_ok=True)
792
793 if prompt and image_file:
794 conversation = [
795 {
796 "role": "<|User|>",
797 # "content": "<image>\n<|grounding|>Given the layout of the image. ",
798 "content": f'{prompt}',
799 # "content": "君不见黄河之水天上来的下一句是什么？",
800 # "content": "<image>\nFree OCR. ",
801 # "content": "<image>\nParse the figure. ",
802 # "content": "<image>\nExtract the text in the image. ",
803 "images": [f'{image_file}'],
804 },
805 {"role": "<|Assistant|>", "content": ""},
806 ]
807
808 elif prompt:
809 conversation = [
810 {
811 "role": "<|User|>",
812 # "content": "<image>\n<|grounding|>Given the layout of the image. ",
813 "content": f'{prompt}',
814 # "content": "君不见黄河之水天上来的下一句是什么？",
815 # "content": "<image>\nFree OCR. ",
816 # "content": "<image>\nParse the figure. ",
817 # "content": "<image>\nExtract the text in the image. ",
818 # "images": [f'{image_file}'],
819 },
820 {"role": "<|Assistant|>", "content": ""},
821 ]
822 else:
823 assert False, f'prompt is none!'
824
825 prompt = format_messages(conversations=conversation, sft_format='plain', system_prompt='')
826
827 patch_size = 16
828 downsample_ratio = 4
829 images = load_pil_images(conversation)
830
831 valid_img_tokens = 0
832 ratio = 1
833
834 image_draw = images[0].copy()
835
836 w,h = image_draw.size
837 # print(w, h)
838 ratio = 1 - ((max(w, h) - min(w, h)) / (max(w, h)))
839
840
841 image_transform=BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True)
842 images_seq_mask = []
843
844 image_token = '<image>'
845 image_token_id = 128815
846 text_splits = prompt.split(image_token)
847
848 images_list, images_crop_list, images_seq_mask = [], [], []
849 tokenized_str = []
850 images_spatial_crop = []
851 for text_sep, image in zip(text_splits, images):
852
853 tokenized_sep = text_encode(tokenizer, text_sep, bos=False, eos=False)
854 tokenized_str += tokenized_sep
855 images_seq_mask += [False] * len(tokenized_sep)
856
857 if crop_mode:
858
859 if image.size[0] <= 640 and image.size[1] <= 640:
860 crop_ratio = [1, 1]
861
862 else:
863 if crop_mode:
864 # best_width, best_height = select_best_resolution(image.size, self.candidate_resolutions)
865 images_crop_raw, crop_ratio = dynamic_preprocess(image)
866 else:
867 # best_width, best_height = self.image_size, self.image_size
868 crop_ratio = [1, 1]
869
870 """process the global view"""
871 # image = image.resize((base_size, base_size))
872 global_view = ImageOps.pad(image, (base_size, base_size),
873 color=tuple(int(x * 255) for x in image_transform.mean))
874
875 if base_size == 1024:
876 valid_img_tokens += int(256 * ratio)
877 elif base_size == 1280:
878 valid_img_tokens += int(400 * ratio)
879 # elif base_size == 640:
880 # valid_img_tokens += int(100 * ratio)
881
882
883
884
885
886 images_list.append(image_transform(global_view).to(torch.bfloat16))
887
888 # global_view_tensor = image_transform(global_view).to(torch.bfloat16)
889
890 width_crop_num, height_crop_num = crop_ratio
891
892 images_spatial_crop.append([width_crop_num, height_crop_num])
893
894
895 if width_crop_num > 1 or height_crop_num > 1:
896 """process the local views"""
897
898 for i in range(len(images_crop_raw)):
899 images_crop_list.append(image_transform(images_crop_raw[i]).to(torch.bfloat16))
900
901 if image_size == 640:
902 valid_img_tokens += len(images_crop_list) * 100
903
904 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
905 num_queries_base = math.ceil((base_size // patch_size) / downsample_ratio)
906
907
908
909 """add image tokens"""
910
911
912
913 tokenized_image = ([image_token_id] * num_queries_base + [image_token_id]) * num_queries_base
914 tokenized_image += [image_token_id]
915 if width_crop_num > 1 or height_crop_num > 1:
916 tokenized_image += ([image_token_id] * (num_queries * width_crop_num) + [image_token_id]) * (
917 num_queries * height_crop_num)
918 tokenized_str += tokenized_image
919 images_seq_mask += [True] * len(tokenized_image)
920 # num_image_tokens.append(len(tokenized_image))
921
922 else:
923 # best_width, best_height = self.image_size, self.image_size
924 # print(image.size, (best_width, best_height)) # check the select_best_resolutions func
925
926 """process the global view"""
927 if image_size <= 640:
928 print('directly resize')
929 image = image.resize((image_size, image_size))
930 # else:
931 global_view = ImageOps.pad(image, (image_size, image_size),
932 color=tuple(int(x * 255) for x in image_transform.mean))
933 images_list.append(image_transform(global_view).to(torch.bfloat16))
934
935 if base_size == 1024:
936 valid_img_tokens += int(256 * ratio)
937 elif base_size == 1280:
938 valid_img_tokens += int(400 * ratio)
939 elif base_size == 640:
940 valid_img_tokens += int(100 * 1)
941 elif base_size == 512:
942 valid_img_tokens += int(64 * 1)
943
944 width_crop_num, height_crop_num = 1, 1
945
946 images_spatial_crop.append([width_crop_num, height_crop_num])
947
948
949 """add image tokens"""
950 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
951
952 tokenized_image = ([image_token_id] * num_queries + [image_token_id]) * num_queries
953 tokenized_image += [image_token_id]
954 # tokenized_image += ([self.image_token_id] * (num_queries * width_crop_num) + [self.image_token_id]) * (
955 # num_queries * height_crop_num)
956 tokenized_str += tokenized_image
957 images_seq_mask += [True] * len(tokenized_image)
958 # num_image_tokens.append(len(tokenized_image))
959
960
961 """process the last text split"""
962 tokenized_sep = text_encode(tokenizer, text_splits[-1], bos=False, eos=False)
963 tokenized_str += tokenized_sep
964 images_seq_mask += [False] * len(tokenized_sep)
965
966 """add the bos tokens"""
967 bos_id = 0
968 tokenized_str = [bos_id] + tokenized_str
969 images_seq_mask = [False] + images_seq_mask
970
971
972
973 input_ids = torch.LongTensor(tokenized_str)
974
975
976
977
978 images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)
979
980
981 if len(images_list) == 0:
982 images_ori = torch.zeros((1, 3, image_size, image_size))
983 images_spatial_crop = torch.zeros((1, 2), dtype=torch.long)
984 images_crop = torch.zeros((1, 3, base_size, base_size))
985
986 else:
987 images_ori = torch.stack(images_list, dim=0)
988 images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)
989 if images_crop_list:
990 images_crop = torch.stack(images_crop_list, dim=0)
991 else:
992 images_crop = torch.zeros((1, 3, base_size, base_size))
993
994
995
996 if not eval_mode:
997 streamer = TPSTextStreamer(tokenizer, interval=tps_interval, skip_prompt=True, skip_special_tokens=False)
998 _orig_sw = getattr(self.config, 'sliding_window_size', None) or getattr(self.config, 'sliding_window', None)
999 self.config._ring_window = _orig_sw
1000 self.config.sliding_window = None
1001 # Build logits processors for ngram
1002 gen_kwargs = dict(
1003 input_ids=input_ids.unsqueeze(0).cuda(),
1004 images=[(images_crop.cuda(), images_ori.cuda())],
1005 images_seq_mask=images_seq_mask.unsqueeze(0).cuda(),
1006 images_spatial_crop=images_spatial_crop,
1007 do_sample=temperature > 0,
1008 temperature=temperature if temperature > 0 else None,
1009 eos_token_id=tokenizer.eos_token_id,
1010 streamer=streamer,
1011 max_length=max_length,
1012 use_cache=True
1013 )
1014 if no_repeat_ngram_size > 0 and ngram_window > 0:
1015 gen_kwargs['logits_processor'] = [SlidingWindowNoRepeatNgramProcessor(no_repeat_ngram_size, ngram_window)]
1016 elif no_repeat_ngram_size > 0:
1017 gen_kwargs['no_repeat_ngram_size'] = no_repeat_ngram_size
1018 with torch.autocast("cuda", dtype=torch.bfloat16):
1019 with torch.no_grad():
1020 output_ids = self.generate(**gen_kwargs)
1021 self.config.sliding_window = _orig_sw
1022
1023 else:
1024 _orig_sw = getattr(self.config, 'sliding_window_size', None) or getattr(self.config, 'sliding_window', None)
1025 self.config._ring_window = _orig_sw
1026 self.config.sliding_window = None
1027 gen_kwargs = dict(
1028 input_ids=input_ids.unsqueeze(0).cuda(),
1029 images=[(images_crop.cuda(), images_ori.cuda())],
1030 images_seq_mask=images_seq_mask.unsqueeze(0).cuda(),
1031 images_spatial_crop=images_spatial_crop,
1032 do_sample=temperature > 0,
1033 temperature=temperature if temperature > 0 else None,
1034 eos_token_id=tokenizer.eos_token_id,
1035 max_length=max_length,
1036 use_cache=True
1037 )
1038 if no_repeat_ngram_size > 0 and ngram_window > 0:
1039 gen_kwargs['logits_processor'] = [SlidingWindowNoRepeatNgramProcessor(no_repeat_ngram_size, ngram_window)]
1040 elif no_repeat_ngram_size > 0:
1041 gen_kwargs['no_repeat_ngram_size'] = no_repeat_ngram_size
1042 with torch.autocast("cuda", dtype=torch.bfloat16):
1043 with torch.no_grad():
1044 output_ids = self.generate(**gen_kwargs)
1045 self.config.sliding_window = _orig_sw
1046
1047
1048 if '<image>' in conversation[0]['content'] and eval_mode:
1049 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
1050 stop_str = '<｜end▁of▁sentence｜>'
1051 if outputs.endswith(stop_str):
1052 outputs = outputs[:-len(stop_str)]
1053 # re_match
1054 outputs = outputs.strip()
1055
1056 return outputs
1057
1058 if '<image>' in conversation[0]['content'] and test_compress:
1059 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
1060 pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False))
1061 print('='*50)
1062 print('image size: ', (w, h))
1063 print('valid image tokens: ', int(valid_img_tokens))
1064 print('output texts tokens (valid): ', pure_texts_outputs_token_length)
1065 print('compression ratio: ', round(pure_texts_outputs_token_length/valid_img_tokens, 2))
1066 print('='*50)
1067
1068
1069 if '<image>' in conversation[0]['content'] and save_results:
1070 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
1071 stop_str = '<｜end▁of▁sentence｜>'
1072
1073 print('='*15 + 'save results:' + '='*15)
1074
1075 # # # # conv.messages[-1][-1] = outputs
1076 if outputs.endswith(stop_str):
1077 outputs = outputs[:-len(stop_str)]
1078 outputs = outputs.strip()
1079
1080 matches_ref, matches_images, mathes_other = re_match(outputs)
1081 # print(matches_ref)
1082 result = process_image_with_refs(image_draw, matches_ref, output_path)
1083
1084
1085 for idx, a_match_image in enumerate(tqdm(matches_images, desc="image")):
1086 outputs = outputs.replace(a_match_image, '![](images/' + str(idx) + '.jpg)\n')
1087
1088 for idx, a_match_other in enumerate(tqdm(mathes_other, desc="other")):
1089 outputs = outputs.replace(a_match_other, '').replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')
1090
1091
1092 # if 'structural formula' in conversation[0]['content']:
1093 # outputs = '<smiles>' + outputs + '</smiles>'
1094 with open(f'{output_path}/result.md', 'w', encoding = 'utf-8') as afile:
1095 afile.write(outputs)
1096
1097 if 'line_type' in outputs:
1098 import matplotlib.pyplot as plt
1099 lines = eval(outputs)['Line']['line']
1100
1101 line_type = eval(outputs)['Line']['line_type']
1102 # print(lines)
1103
1104 endpoints = eval(outputs)['Line']['line_endpoint']
1105
1106 fig, ax = plt.subplots(figsize=(3,3), dpi=200)
1107 ax.set_xlim(-15, 15)
1108 ax.set_ylim(-15, 15)
1109
1110 for idx, line in enumerate(lines):
1111 try:
1112 p0 = eval(line.split(' -- ')[0])
1113 p1 = eval(line.split(' -- ')[-1])
1114
1115 if line_type[idx] == '--':
1116 ax.plot([p0[0], p1[0]], [p0[1], p1[1]], linewidth=0.8, color='k')
1117 else:
1118 ax.plot([p0[0], p1[0]], [p0[1], p1[1]], linewidth = 0.8, color = 'k')
1119
1120 ax.scatter(p0[0], p0[1], s=5, color = 'k')
1121 ax.scatter(p1[0], p1[1], s=5, color = 'k')
1122 except:
1123 pass
1124
1125 for endpoint in endpoints:
1126
1127 label = endpoint.split(': ')[0]
1128 (x, y) = eval(endpoint.split(': ')[1])
1129 ax.annotate(label, (x, y), xytext=(1, 1), textcoords='offset points',
1130 fontsize=5, fontweight='light')
1131
1132
1133 plt.savefig(f'{output_path}/geo.jpg')
1134 plt.close()
1135
1136 result.save(f"{output_path}/result_with_boxes.jpg")
1137
1138
1139 def infer_multi(self, tokenizer, prompt='', image_files=None, output_path='', image_size=640, save_results=False, max_length=32768, tps_interval=0, no_repeat_ngram_size=0, ngram_window=0, temperature=0.0):
1140 """
1141 Multi-image inference. Does NOT support crop mode.
1142 Prompt uses a single <image> token (e.g. "<image>Multi page parsing.").
1143 All images' token sequences are concatenated at that single <image> position,
1144 separated by a single image_token_id between each image (same as crop mode separator).
1145
1146 Args:
1147 prompt: text prompt with one <image> token, e.g. "<image>Multi page parsing."
1148 image_files: list of image file paths
1149 image_size: size to resize each image to
1150 save_results: whether to save output to file
1151 """
1152 self.disable_torch_init()
1153
1154 if image_files is None or len(image_files) == 0:
1155 assert False, 'image_files must be a non-empty list for multi-image inference!'
1156
1157 os.makedirs(output_path, exist_ok=True)
1158 os.makedirs(f'{output_path}/images', exist_ok=True)
1159
1160 # Prompt contains a single <image>, all image files go into "images" list
1161 conversation = [
1162 {
1163 "role": "<|User|>",
1164 "content": f'{prompt}',
1165 "images": image_files,
1166 },
1167 {"role": "<|Assistant|>", "content": ""},
1168 ]
1169
1170 formatted_prompt = format_messages(conversations=conversation, sft_format='plain', system_prompt='')
1171
1172 patch_size = 16
1173 downsample_ratio = 4
1174
1175 # Load all images
1176 images = load_pil_images(conversation)
1177
1178 image_transform = BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True)
1179
1180 image_token = '<image>'
1181 image_token_id = 128815
1182
1183 # Split on the single <image> token -> 2 parts: before and after
1184 text_splits = formatted_prompt.split(image_token)
1185
1186 images_list, images_seq_mask = [], []
1187 tokenized_str = []
1188 images_spatial_crop = []
1189
1190 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
1191
1192 # Tokenize text before <image>
1193 tokenized_sep = text_encode(tokenizer, text_splits[0], bos=False, eos=False)
1194 tokenized_str += tokenized_sep
1195 images_seq_mask += [False] * len(tokenized_sep)
1196
1197 # Process all images at the single <image> position
1198 for idx, image in enumerate(images):
1199 # Match single-image logic: if image_size <= 640, resize all images
1200 if image_size <= 640:
1201 image = image.resize((image_size, image_size))
1202 global_view = ImageOps.pad(image, (image_size, image_size),
1203 color=tuple(int(x * 255) for x in image_transform.mean))
1204
1205 images_list.append(image_transform(global_view).to(torch.bfloat16))
1206 images_spatial_crop.append([1, 1])
1207
1208 # Image tokens for this image (same structure as single-image non-crop mode)
1209 tokenized_image = ([image_token_id] * num_queries + [image_token_id]) * num_queries
1210 tokenized_image += [image_token_id] # separator token between images
1211 tokenized_str += tokenized_image
1212 images_seq_mask += [True] * len(tokenized_image)
1213
1214 # Tokenize text after <image>
1215 tokenized_sep = text_encode(tokenizer, text_splits[1], bos=False, eos=False)
1216 tokenized_str += tokenized_sep
1217 images_seq_mask += [False] * len(tokenized_sep)
1218
1219 # Add bos token
1220 bos_id = 0
1221 tokenized_str = [bos_id] + tokenized_str
1222 images_seq_mask = [False] + images_seq_mask
1223
1224 input_ids = torch.LongTensor(tokenized_str)
1225 images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)
1226
1227 # Stack all images as image_ori; dummy_crop is zeros (triggers no-crop branch)
1228 images_ori = torch.stack(images_list, dim=0) # [N, 3, H, W]
1229 images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)
1230 dummy_crop = torch.zeros((1, 3, image_size, image_size))
1231
1232 streamer = TPSTextStreamer(tokenizer, interval=tps_interval, skip_prompt=True, skip_special_tokens=False)
1233 # Disable config.sliding_window to prevent DynamicCache from truncating prefill tokens.
1234 # The ring buffer in SlidingWindowLlamaAttention handles sliding window manually.
1235 _orig_sw = getattr(self.config, 'sliding_window_size', None) or getattr(self.config, 'sliding_window', None)
1236 self.config._ring_window = _orig_sw # Save for ring buffer to read
1237 self.config.sliding_window = None
1238 with torch.autocast("cuda", dtype=torch.bfloat16):
1239 with torch.no_grad():
1240 gen_kwargs = dict(
1241 input_ids=input_ids.unsqueeze(0).cuda(),
1242 images=[(dummy_crop.cuda(), images_ori.cuda())],
1243 images_seq_mask=images_seq_mask.unsqueeze(0).cuda(),
1244 images_spatial_crop=images_spatial_crop,
1245 do_sample=temperature > 0,
1246 temperature=temperature if temperature > 0 else None,
1247 eos_token_id=tokenizer.eos_token_id,
1248 streamer=streamer,
1249 max_length=max_length,
1250 use_cache=True
1251 )
1252 if no_repeat_ngram_size > 0 and ngram_window > 0:
1253 gen_kwargs['logits_processor'] = [SlidingWindowNoRepeatNgramProcessor(no_repeat_ngram_size, ngram_window)]
1254 elif no_repeat_ngram_size > 0:
1255 gen_kwargs['no_repeat_ngram_size'] = no_repeat_ngram_size
1256 output_ids = self.generate(**gen_kwargs)
1257 self.config.sliding_window = _orig_sw # Restore
1258
1259 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
1260 stop_str = '<｜end▁of▁sentence｜>'
1261 if outputs.endswith(stop_str):
1262 outputs = outputs[:-len(stop_str)]
1263 outputs = outputs.strip()
1264
1265 output_tokens = len(text_encode(tokenizer, outputs, bos=False, eos=False))
1266
1267 if save_results:
1268 print('=' * 15 + 'save results:' + '=' * 15)
1269 pages = outputs.split('<PAGE>')[1:]
1270 processed_pages = []
1271 for page_idx, page_output in enumerate(pages):
1272 page_output = page_output.strip()
1273 if page_idx >= len(images):
1274 processed_pages.append(page_output)
1275 continue
1276
1277 matches_ref, matches_images, mathes_other = re_match(page_output)
1278 image_prefix = f'page_{page_idx}_'
1279 result = process_image_with_refs(
1280 images[page_idx].copy(),
1281 matches_ref,
1282 output_path,
1283 image_prefix=image_prefix,
1284 )
1285 result.save(f"{output_path}/result_with_boxes_{page_idx}.jpg")
1286
1287 for idx, a_match_image in enumerate(tqdm(matches_images, desc=f"image_page_{page_idx}")):
1288 page_output = page_output.replace(a_match_image, f'![](images/{image_prefix}{idx}.jpg)\n')
1289
1290 for idx, a_match_other in enumerate(tqdm(mathes_other, desc=f"other_page_{page_idx}")):
1291 page_output = page_output.replace(a_match_other, '').replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')
1292
1293 processed_pages.append(page_output)
1294
1295 outputs = '<PAGE>\n' + '\n<PAGE>\n'.join(processed_pages)
1296 with open(f'{output_path}/result.md', 'w', encoding='utf-8') as afile:
1297 afile.write(outputs)
1298
1299 return outputs, output_tokens