Update vision_transformer.py

vision_transformer.py

@@ -1,3 +1,22 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Mostly copy-paste from timm library.
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+
+Modified to work with Hugging Face Transformers
+"""
 import math
 from functools import partial
 import torch
@@ -304,6 +323,24 @@ class VisionTransformer(PreTrainedModel):
             attentions=None,
         )
     
+    def forward_features(self, x):
+        """
+        Feature extraction method - EXACT match to local vision_transformer.py
+        This ensures HuggingFace and local models give identical results
+        """
+        x = self.prepare_tokens(x)  # Tokenize input
+        
+        for blk in self.blocks:
+            x = blk(x)
+        
+        x_norm = self.norm(x)  # Normalize tokens
+        
+        return {
+            "x_norm_clstoken": x_norm[:, 0],  # CLS token
+            "x_norm_patchtokens": x_norm[:, 1:],  # Patch tokens
+            "x_prenorm": x,  # Before norm
+        }
+    
     def get_last_selfattention(self, x):
         """Get attention from last block"""
         x = self.prepare_tokens(x)