Upload model

config.json

@@ -1,14 +1,12 @@
 {
   "architectures": [
-    "VisFocusModel",
-    "VisFocusForLocalizedMaskedLanguageModeling",
-    "VisFocusForImageTextToText"
+    "VisFocusModelForImageTextToText"
   ],
   "auto_map": {
     "AutoConfig": "configuration_visfocus.VisFocusConfig",
     "AutoModel": "configuration_visfocus.VisFocusPreTrainedModel",
     "AutoModelForConditionalGeneration": "configuration_visfocus.VisFocusModelForImageTextToText",
-    "AutoModelForImageTextToText": "configuration_visfocus.VisFocusModelForImageTextToText"
+    "AutoModelForImageTextToText": "modeling_visfocus.VisFocusModelForImageTextToText"
   },
   "cache_dir": null,
   "do_lower_case": true,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f644d82b1150eba66c88fcb62fed2cdd1a871f0ea44bfb136ea0bc182b8c9fae
+oid sha256:e631e85b53c7ccd3df8c70a10c528c8582394914427f9cb0ba185b81d9b8ed22
 size 1047109288

modeling_visfocus.py

@@ -4,7 +4,7 @@ from torch import nn
 from torch.nn import LayerNorm, CrossEntropyLoss, L1Loss
 from torch.nn import functional as F
 
-from transformers import PreTrainedModel, T5Tokenizer, T5Model, logging
+from transformers import PreTrainedModel, AutoTokenizer, GenerationMixin, logging
 from transformers.models.t5.modeling_t5 import T5Stack
 from transformers.modeling_outputs import Seq2SeqLMOutput, BaseModelOutput
 from transformers.file_utils import ModelOutput
@@ -17,8 +17,10 @@ import yaml
 import copy
 from easydict import EasyDict
 
-from .configuration_visfocus import VisFocusConfig
-from .modeling_vilmaswin import VilmaSwinTransformerV2
+from configuration_visfocus import VisFocusConfig
+from modeling_vilmaswin import VilmaSwinTransformerV2
+from image_processing_visfocus import VisFocusImageProcessor
+from processing_visfocus import VisFocusProcessor
 
 logger = logging.get_logger(__name__)
 
@@ -148,6 +150,7 @@ def load_vision_pretrained(configs, model):
 
 class T5_Encoder(nn.Module):
     def __init__(self, t5_variant='base', freeze=True):
+        from transformers import T5Tokenizer, T5Model
         super().__init__()
         self.tokenizer = T5Tokenizer.from_pretrained(f'{t5_variant}')
         model = T5Model.from_pretrained(f'{t5_variant}')
@@ -255,7 +258,7 @@ class MLP(nn.Module):
         return x
 
 
-class VisFocusModel(PreTrainedModel):
+class VisFocusPreTrainedModel(PreTrainedModel, GenerationMixin):
     config_class = VisFocusConfig
 
     def __init__(self, config):
@@ -413,7 +416,7 @@ class VisFocusModel(PreTrainedModel):
             
             if self.config.vl_l1_loss:
                 labels_ = labels.clone()
-                labels_[labels_ == -100] = self.input_tokenizer.pad_token_id # -> replace the ignore_index with the pad_token id to calculate the text target for the vl loss
+                labels_[labels_ == -100] = 0 # -> replace the ignore_index with the pad_token id to calculate the text target for the vl loss
                 with torch.no_grad():
                     target = self.encoder(input_ids=labels_).last_hidden_state
                 if target.shape[1] != hidden_states.shape[1]:
@@ -567,15 +570,6 @@ class VisFocusModel(PreTrainedModel):
             inputs_embeds=inputs_tensor, **encoder_kwargs)
 
         return model_kwargs
-
-    def add_task_tokens(self):
-        self.input_tokenizer.add_tokens('<OCR>', special_tokens=True)
-        self.task_token_ids = torch.nn.ParameterDict([['ocr', self.register_token('<OCR>')]])
-
-    def register_token(self, token: str):
-        self.input_tokenizer.add_tokens(token, special_tokens=True)
-        token_ids = self.input_tokenizer.encode(token)
-        return torch.nn.Parameter(torch.tensor(token_ids), requires_grad=False)
     
     def set_task_name(self, task_name):
         if task_name:
@@ -585,7 +579,7 @@ class VisFocusModel(PreTrainedModel):
         return torch.ones((inp.shape[:2]), dtype=torch.int32).to(self.device)
 
 
-class VisFocusModelForLocalizedMaskedLanguageModeling(VisFocusModel):
+class VisFocusModelForLocalizedMaskedLanguageModeling(VisFocusPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.set_task_name('mpm')
@@ -604,6 +598,7 @@ class VisFocusModelForLocalizedMaskedLanguageModeling(VisFocusModel):
                 **kwargs):        
         if not kwargs.get('encoder_outputs'):
             if self.task_name == 'ocr':
+                # NOTE: not supported yet
                 input_ids = None
                 if not hasattr(self, 'prompt_embeds'):
                     prompt = 'what is written in this document?'
@@ -681,11 +676,6 @@ class VisFocusModelForLocalizedMaskedLanguageModeling(VisFocusModel):
         inputs = self._maybe_initialize_input_ids_for_generation(inputs, bos_token_id, model_kwargs)
         return inputs, input_name, model_kwargs
 
-    def add_task_tokens(self):
-        super().add_task_tokens()
-        self.input_tokenizer.add_tokens('<MPM>', special_tokens=True)
-        self.task_token_ids.update({'mpm': self.register_token('<MPM>')})
-
 
 class VisFocusModelForImageTextToText(VisFocusModelForLocalizedMaskedLanguageModeling):
     def __init__(self, config):
@@ -759,11 +749,6 @@ class VisFocusModelForImageTextToText(VisFocusModelForLocalizedMaskedLanguageMod
         text_embeds = self.shared(input_ids) # for concat, use direct the T5 nn.embeddings
         return text_embeds, vision_embeds, attention_mask
 
-    def add_task_tokens(self):
-        super().add_task_tokens()
-        self.input_tokenizer.add_tokens('<LMPM_VQA_CONCAT>', special_tokens=True)
-        self.task_token_ids.update({'pm_vqa_concat': self.register_token('<LMPM_VQA_CONCAT>')})
-
 
 def _to_cuda(sample, device=torch.device('cuda')):
     if isinstance(sample, torch.Tensor):
@@ -806,5 +791,13 @@ if __name__ == '__main__':
     cfg = VisFocusConfig.from_pretrained('configs/config.json')
     cfg.push_to_hub('ofirab/visfocus-base-docvqa')
     model = VisFocusModelForImageTextToText(cfg)
+
+    VisFocusConfig.register_for_auto_class()
+    VisFocusPreTrainedModel.register_for_auto_class("AutoModel")
+    VisFocusModelForImageTextToText.register_for_auto_class("AutoModelForImageTextToText")
+
     model.push_to_hub('ofirab/visfocus-base-docvqa')
+    pr = VisFocusImageProcessor(is_train=False)
+    tokenizer = AutoTokenizer.from_pretrained('ofirab/visfocus-base-docvqa')
+    prr = VisFocusProcessor(pr, tokenizer)
     model.to(DEVICE)