Fix weight init, deprecation warnings, safetensors (#38)

- Update model and modeling file (f18f8283c0ae5ec93982c132fe94787848ea1429)
- Update modeling file (f285b2c5edf2e5ac4f5eb5db70873d2e850754dc)
- Tie weights correctly (b5c7055a2e5f2db1d629e32c0f64e1ae0e57a7cb)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b4e610c952eef90a836c56cda0f398a672a3a6ca7b4d96b0e09a86dee42e2c3
+size 1540980506

modeling_florence2.py

@@ -26,9 +26,10 @@ import torch.utils.checkpoint as checkpoint
 from torch.nn import CrossEntropyLoss 
 from collections import OrderedDict
 from einops import rearrange
-from timm.models.layers import DropPath, trunc_normal_
+from timm.layers import DropPath, trunc_normal_
 
 from transformers.modeling_utils import PreTrainedModel
+from transformers.generation.utils import GenerationMixin
 from transformers.utils import (
     ModelOutput,
     add_start_docstrings,
@@ -609,29 +610,10 @@ class DaViT(nn.Module):
         self.avgpool = nn.AdaptiveAvgPool1d(1)
         self.head = nn.Linear(self.embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()
 
-        self.apply(self._init_weights)
-
     @property
     def dim_out(self):
         return self.embed_dims[-1]
 
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.Conv2d):
-            nn.init.normal_(m.weight, std=0.02)
-            for name, _ in m.named_parameters():
-                if name in ['bias']:
-                    nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm):
-            nn.init.constant_(m.weight, 1.0)
-            nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.BatchNorm2d):
-            nn.init.constant_(m.weight, 1.0)
-            nn.init.constant_(m.bias, 0)
-
     def forward_features_unpool(self, x):
         """
         forward until avg pooling 
@@ -1450,6 +1432,17 @@ class Florence2LanguagePreTrainedModel(PreTrainedModel):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.Conv2d):
+            nn.init.normal_(module.weight, std=0.02)
+            for name, _ in module.named_parameters():
+                if name == "bias":
+                    nn.init.constant_(module.bias, 0)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.constant_(module.weight, 1.0)
+            nn.init.constant_(module.bias, 0)
+        elif isinstance(module, nn.BatchNorm2d):
+            nn.init.constant_(module.weight, 1.0)
+            nn.init.constant_(module.bias, 0)
 
     @property
     def dummy_inputs(self):
@@ -2059,7 +2052,7 @@ class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
         )
 
 
-class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel):
+class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel, GenerationMixin):
     base_model_prefix = "model"
     _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"]
     _keys_to_ignore_on_load_missing = ["final_logits_bias"]
@@ -2073,6 +2066,12 @@ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel
         # Initialize weights and apply final processing
         self.post_init()
 
+    def _tie_weights(self):
+        if self.config.tie_word_embeddings:
+            self._tie_or_clone_weights(self.model.encoder.embed_tokens, self.model.shared)
+            self._tie_or_clone_weights(self.model.decoder.embed_tokens, self.model.shared)
+            self._tie_or_clone_weights(self.lm_head, self.model.shared)
+
     def get_encoder(self):
         return self.model.get_encoder()
 
@@ -2530,6 +2529,8 @@ class Florence2VisionModelWithProjection(Florence2PreTrainedModel):
     FLORENCE2_START_DOCSTRING,
 )
 class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
+    _tied_weights_keys = ["language_model.encoder.embed_tokens.weight", "language_model.decoder.embed_tokens.weight", "language_model.lm_head.weight"]
+
     def __init__(self, config: Florence2Config):
         super().__init__(config)
         assert config.vision_config.model_type == 'davit', 'only DaViT is supported for now'
@@ -2544,8 +2545,6 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
 
         language_model = Florence2LanguageForConditionalGeneration(config=config.text_config)
 
-        if language_model._tied_weights_keys is not None:
-            self._tied_weights_keys = [f"language_model.{k}" for k in language_model._tied_weights_keys]
         self.language_model = language_model
 
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1