Update modeling_ovis.py

modeling_ovis.py

@@ -1,10 +1,12 @@
 import logging
 import os
+from packaging import version
 from importlib import import_module
 from typing import List, Callable, Union, Optional, Dict
 
 import PIL.Image
 import torch
+import transformers
 from torch import Tensor
 from torch.nn import init
 from torch.nn.functional import softmax, gumbel_softmax, pad
@@ -556,25 +558,42 @@ class Ovis(OvisPreTrainedModel):
         cache_cls = HybridCache
         llm = self.get_llm()
 
-        need_new_cache = (
-            not hasattr(llm, "_cache")
-            or (not isinstance(llm._cache, cache_cls))
-            or llm._cache.batch_size != batch_size
-            or llm._cache.max_cache_len < max_cache_len
-        )
+        if version.parse(transformers.__version__) >= version.parse("4.46.0"):
+            need_new_cache = (
+                not hasattr(llm, "_cache")
+                or (not isinstance(llm._cache, cache_cls))
+                or llm._cache.batch_size != batch_size
+                or llm._cache.max_cache_len < max_cache_len
+            )
+        else:
+            need_new_cache = (
+                not hasattr(llm, "_cache")
+                or (not isinstance(llm._cache, cache_cls))
+                or llm._cache.max_batch_size != batch_size
+                or llm._cache.max_cache_len < max_cache_len
+            )
 
         if need_new_cache:
             if hasattr(llm.config, "_pre_quantization_dtype"):
                 cache_dtype = llm.config._pre_quantization_dtype
             else:
                 cache_dtype = llm.dtype
-            llm._cache = cache_cls(
-                config=llm.config,
-                batch_size=batch_size,
-                max_cache_len=max_cache_len,
-                device=llm.device,
-                dtype=cache_dtype,
-            )
+            if version.parse(transformers.__version__) >= version.parse("4.46.0"):
+                llm._cache = cache_cls(
+                    config=llm.config,
+                    batch_size=batch_size,
+                    max_cache_len=max_cache_len,
+                    device=llm.device,
+                    dtype=cache_dtype,
+                )
+            else:
+                llm._cache = cache_cls(
+                    config=llm.config,
+                    max_batch_size=batch_size,
+                    max_cache_len=max_cache_len,
+                    device=llm.device,
+                    dtype=cache_dtype,
+                )
         else:
             llm._cache.reset()
         return llm._cache