From 02ed7bfd7fb4688aaba9968327dfee0a9f79ce03 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 21 Nov 2024 03:58:18 -0800
Subject: [PATCH] Update vision.py

---
 unsloth/models/vision.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 8d0f63c44d..237673f10c 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -36,6 +36,32 @@ __all__ = [
     "FastBaseVisionModel",
 ]
 
+def _wrap_fast_inference(generate, device_type, dtype, model):
+    # Wraps inference with bfloat16 / float16
+    @torch.inference_mode
+    def _fast_generate(*args, **kwargs):
+        # For num_logits_to_keep
+        kwargs["num_logits_to_keep"] = 1
+
+        # Remove token_type_ids
+        kwargs.pop("token_type_ids", None)
+
+        # Check pad_token
+        model_eos_token_id = getattr(model.config, "eos_token_id", None)
+        if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"):
+            model_eos_token_id = model_eos_token_id[0]
+
+        kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id)
+
+        # Autocasted
+        with torch.autocast(device_type = device_type, dtype = dtype):
+            output = generate(*args, **kwargs)
+        pass
+        return output
+    pass
+    return _fast_generate
+pass
+
 
 class FastBaseVisionModel: