From 0ee78d455c7b9131ef468848e31cab489f729306 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Sep 2024 13:28:05 -0700
Subject: [PATCH] Update _utils.py

---
 unsloth/models/_utils.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 09b448a2f6..6144efe485 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -141,24 +141,24 @@ pass
 
 # =============================================
 # Fix KeyError: 'Cache only has 0 layers, attempted to access layer with index 0'
-# import transformers.cache_utils
-# if hasattr(transformers.cache_utils, "DynamicCache") and \
-#     transformers.cache_utils.DynamicCache.__getitem__.__name__ != "__cache_utils_getitem__":
+import transformers.cache_utils
+if hasattr(transformers.cache_utils, "DynamicCache") and \
+    transformers.cache_utils.DynamicCache.__getitem__.__name__ != "__cache_utils_getitem__":
 
-#     source = inspect.getsource(transformers.cache_utils.DynamicCache.__getitem__)
-#     start = source.find("def")
-#     spaces = start*" "
-#     source = source.split("\n")
-#     source = "\n".join(x[start:] for x in source)
-#     where = source.find("raise KeyError")
-#     source = source[:where] + \
-#         f"if len(self) == 0:\n{spaces}{spaces}"\
-#         "    raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \
-#         f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:]
-#     source = source.replace("__getitem__", "__cache_utils_getitem__", 1)
-#     exec(source)
-#     transformers.cache_utils.DynamicCache.__getitem__ = __cache_utils_getitem__
-# pass
+    source = inspect.getsource(transformers.cache_utils.DynamicCache.__getitem__)
+    start = source.find("def")
+    spaces = start*" "
+    source = source.split("\n")
+    source = "\n".join(x[start:] for x in source)
+    where = source.find("raise KeyError")
+    # source = source[:where] + \
+    #     f"if len(self) == 0:\n{spaces}{spaces}"\
+    #     "    raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \
+    #     f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:]
+    source = source.replace("__getitem__", "__cache_utils_getitem__", 1)
+    exec(source)
+    transformers.cache_utils.DynamicCache.__getitem__ = __cache_utils_getitem__
+pass
 # =============================================
 
 # =============================================