From 243070d901f43c3b26a8b6bcccb0313d2d77811e Mon Sep 17 00:00:00 2001
From: nikhil <nikhil@DNa1d6487.SUNet>
Date: Mon, 4 Nov 2024 14:35:58 -0800
Subject: [PATCH] Modified cache

---
 lotus/models/lm.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/lotus/models/lm.py b/lotus/models/lm.py
index b86ed8ce..aa7799c0 100644
--- a/lotus/models/lm.py
+++ b/lotus/models/lm.py
@@ -1,20 +1,19 @@
-import logging
-from functools import lru_cache
+import numpy as np
 from typing import Any
-
 import litellm
-import numpy as np
 from litellm import batch_completion, completion_cost
-from litellm.caching import Cache
+from litellm.caching import Cache, LiteLLMCacheType
 from litellm.types.utils import ChatCompletionTokenLogprob, Choices, ModelResponse
 from litellm.utils import token_counter
 from openai import OpenAIError
 from tokenizers import Tokenizer
 
 import lotus
+import logging
+from functools import lru_cache
 from lotus.types import LMOutput, LMStats, LogprobsForCascade, LogprobsForFilterCascade
 
-litellm.cache = Cache(disk_cache_dir=".lotus_cache", type="disk")
+litellm.cache = Cache(disk_cache_dir=".lotus_cache", type=LiteLLMCacheType.DISK)
 
 class LM:
     def __init__(