Adding TF 5.5 support

QEfficient/transformers/cache_utils.py

@@ -10,7 +10,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch
-from transformers.cache_utils import Cache, EncoderDecoderCache
+from transformers.cache_utils import Cache, CacheLayerMixin, EncoderDecoderCache
 
 try:
     # transformers<5.3 had these hybrid cache classes
@@ -66,11 +66,13 @@ def _get_invalid_idx_value(cls):
         else:
             return 0
 
+
 def _match_invalid_mask(invalid_mask: torch.Tensor, target_len: int) -> torch.Tensor:
     if invalid_mask.shape[-1] == target_len:
         return invalid_mask
     return invalid_mask[..., :target_len]
 
+
 class QEffDynamicLayer(CacheLayerMixin):
     is_compileable = False
 
@@ -447,21 +449,6 @@ def append_new_layers(self, layer_idx: int) -> None:
         while len(self.layers) <= layer_idx:
             self.layers.append(QEffDynamicLayer())
 
-    @classmethod
-    def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "QEffDynamicCache":
-        cache = cls()
-        if past_key_values is not None:
-            for layer_idx in range(len(past_key_values)):
-                key_states, value_states = past_key_values[layer_idx]
-                cache.update(key_states, value_states, layer_idx)
-        return cache
-
-    def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
-        legacy_cache = ()
-        for layer in self.layers:
-            legacy_cache += ((layer.keys, layer.values),)
-        return legacy_cache
-
     def get_seq_length(self, layer_idx: Optional[int] = 0, cache_position: Optional[torch.LongTensor] = None) -> int:
         """
         Keep backward-compatible call shape while deferring to upstream implementation.
@@ -491,17 +478,6 @@ def __iter__(self):
         for idx in range(len(self.layers)):
             yield self[idx]
 
-    def get_seq_length(self, layer_idx: Optional[int] = 0, *args, **kwargs) -> int:
-        if layer_idx is None:
-            layer_idx = 0
-        is_empty_layer = (
-            len(self.layers) == 0
-            or len(self.layers) <= layer_idx
-            or getattr(self.layers[layer_idx], "keys", None) is None
-            or len(self.layers[layer_idx].keys) == 0
-        )
-        return self.layers[layer_idx].keys.shape[-2] if not is_empty_layer else 0
-
     def read_only_blockedKV(self, start_index, end_index, layer_idx, cache_kwargs):
         """
         Reads the `key_states` and `value_states` for the layer `layer_idx`.

QEfficient/transformers/models/gemma3/modeling_gemma3.py

@@ -982,7 +982,9 @@ def get_dummy_inputs(
         fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS
 
         # Add data for KV
-        pkv_dtype = next(self.language_model.parameters()).dtype if hasattr(self, "language_model") else self.config.torch_dtype
+        pkv_dtype = (
+            next(self.language_model.parameters()).dtype if hasattr(self, "language_model") else self.config.torch_dtype
+        )
         lang_inputs["past_key_values"] = self.get_dummy_pkv_cache(
             config=self.language_model.config,
             batch_size=fbs if continuous_batching else bs,

QEfficient/transformers/models/gpt2/modeling_gpt2.py

@@ -31,7 +31,7 @@ def eager_attention_forward(module, query, key, value, attention_mask, head_mask
     if attention_mask is not None:
         if attention_mask.dtype == torch.bool:
             attn_weights = torch.where(
-                attention_mask, torch.tensor(MIN_MASKED_ATTENTION_VALUE,  dtype=module.config.torch_dtype), attn_weights
+                attention_mask, torch.tensor(MIN_MASKED_ATTENTION_VALUE, dtype=module.config.torch_dtype), attn_weights
             )
         else:
             attn_weights = attn_weights + attention_mask