add model.ocr

modelling_h2ovl_chat.py

@@ -13,6 +13,7 @@ import transformers
 from .conversation import get_conv_template
 from .configuration_h2ovl_chat import H2OVLChatConfig
 from .image_process import load_single_image, load_multi_images
+import re
 
 logger = logging.get_logger(__name__)
 
@@ -338,3 +339,136 @@ class H2OVLChatModel(PreTrainedModel):
         )
 
         return outputs
+
+    def ocr(self, tokenizer, image_files, question, generation_config , max_tiles=6, history=None, return_history=False,
+             num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
+             verbose=False):
+
+        from transformers import LogitsProcessor
+        class SuppressConsecutiveSpacesLogitsProcessor(LogitsProcessor):
+            def __init__(self, tokenizer):
+                self.tokenizer = tokenizer
+            def __call__(self, input_ids, scores):
+                logits = scores[-1].squeeze()
+                _, topk_indices = torch.topk(logits, 30)
+                if input_ids.shape[1] > 1:
+                    if len(self.tokenizer.decode(input_ids[0, -1]).strip()) == 0 and topk_indices[0] == input_ids[0, -1]:
+                        for i in range(len(topk_indices)):
+                            if len(self.tokenizer.decode(topk_indices[i]).strip()) == 0:
+                                scores[0, topk_indices[i]] = -99999999.
+                            else:
+                                break
+                return scores
+
+        if image_files:
+            if isinstance(image_files, list):
+                pixel_values, num_patches_list = load_multi_images(image_files, max_num=max_tiles)  # Load multiple images
+            else:
+                pixel_values, num_patches_list = load_single_image(image_files, max_num=max_tiles, msac=self.use_msac)  # Load single image
+        else:
+            pixel_values = None
+            num_patches_list = []
+
+
+        if history is None and pixel_values is not None and '<image>' not in question:
+            question = '<image>\n' + question
+
+        if num_patches_list is None:
+            num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []
+
+        assert pixel_values is None or len(pixel_values) == sum(num_patches_list)
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep)
+
+        space_suppressor = SuppressConsecutiveSpacesLogitsProcessor(tokenizer)
+
+        history = [] if history is None else history
+        for (old_question, old_answer) in history:
+            template.append_message(template.roles[0], old_question)
+            template.append_message(template.roles[1], old_answer)
+        template.append_message(template.roles[0], question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+
+        if verbose and pixel_values is not None:
+            image_bs = pixel_values.shape[0]
+            print(f'dynamic ViT batch size: {image_bs}')
+
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+
+        model_inputs = tokenizer(query, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].cuda()
+        attention_mask = model_inputs['attention_mask'].cuda()
+        generation_config['eos_token_id'] = eos_token_id
+        generation_output = self.generate_ocr(
+            space_suppressor=space_suppressor,
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **generation_config
+        )
+        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
+        response = response.split(template.sep)[0].strip()
+        response = re.sub(' +', ' ', response)
+        history.append((question, response))
+        if return_history:
+            return response, history
+        else:
+            query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
+            query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
+            if verbose:
+                print(query_to_print, response)
+            return response
+
+    @torch.no_grad()
+    def generate_ocr(
+            self,
+            space_suppressor,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            input_ids: Optional[torch.FloatTensor] = None,
+            attention_mask: Optional[torch.LongTensor] = None,
+            visual_features: Optional[torch.FloatTensor] = None,
+            generation_config: Optional[GenerationConfig] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            **generate_kwargs,
+    ) -> torch.LongTensor:
+
+        assert self.img_context_token_id is not None
+        if pixel_values is not None:
+            if visual_features is not None:
+                vit_embeds = visual_features
+            else:
+                vit_embeds = self.extract_feature(pixel_values)
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+            B, N, C = input_embeds.shape
+            input_embeds = input_embeds.reshape(B * N, C)
+
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.img_context_token_id)
+            assert selected.sum() != 0
+            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
+
+            input_embeds = input_embeds.reshape(B, N, C)
+        else:
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        outputs = self.language_model.generate(
+            logits_processor=[space_suppressor],
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            use_cache=True,
+            **generate_kwargs,
+        )
+
+        return outputs