Add diskcache for openai requests. Increase speed of requests to OpenAI API by 10 times. (#236)

* Add diskcache for openai requests. Increase speed by 10 times.

Author: IINemo

requirements.txt

@@ -1,7 +1,7 @@
 datasets>=2.14.2
 rouge-score>=0.0.4
 nlpaug>=1.1.10
-scikit-learn
+scikit-learn>=1.5.1
 tqdm>=4.64.1
 matplotlib>=3.6
 pandas>=1.3.5
@@ -11,7 +11,7 @@ transformers>=4.40
 nltk>=3.6.5
 sacrebleu>=1.5.0
 sentencepiece>=0.1.97
-hf-lfs==0.0.3
+hf-lfs>=0.0.3
 pytest>=4.4.1
 pytreebank>=0.2.7
 setuptools>=60.2.0
@@ -23,15 +23,15 @@ protobuf>=4.23
 fschat>=0.2.3
 hydra-core>=1.3.2
 einops
-accelerate
+accelerate>=0.32.1
 bitsandbytes
-openai==0.28.0
-filelock==3.12.2
+openai>=0.28.0
 wget
 sentence-transformers
-bert-score
+bert-score>=0.3.13
 unbabel-comet==2.2.1
 nltk>=3.7,<4
-evaluate
+evaluate>=0.4.2
 spacy>=3.4.0,<4
 fastchat
+diskcache>=5.6.3

src/lm_polygraph/utils/openai_chat.py

@@ -1,10 +1,8 @@
 import openai
-import json
 import os
 import time
 import logging
-
-from filelock import FileLock
+import diskcache as dc
 
 
 log = logging.getLogger()
@@ -17,7 +15,7 @@ class OpenAIChat:
 
     def __init__(
         self,
-        openai_model: str = "gpt-4",
+        openai_model: str = "gpt-4o",
         cache_path: str = os.path.expanduser("~") + "/.cache",
     ):
         """
@@ -31,22 +29,20 @@ def __init__(
             openai.api_key = api_key
         self.openai_model = openai_model
 
-        self.cache_path = os.path.join(cache_path, "openai_chat_cache.json")
-        self.cache_lock = FileLock(self.cache_path + ".lock")
-        with self.cache_lock:
-            if not os.path.exists(self.cache_path):
-                if not os.path.exists(cache_path):
-                    os.makedirs(cache_path)
-                with open(self.cache_path, "w") as f:
-                    json.dump({}, f)
+        self.cache_path = os.path.join(cache_path, "openai_chat_cache.diskcache")
+        if not os.path.exists(cache_path):
+            os.makedirs(cache_path)
 
     def ask(self, message: str) -> str:
-        # check if the message is cached
-        with open(self.cache_path, "r") as f:
-            openai_responses = json.load(f)
+        cache_settings = dc.DEFAULT_SETTINGS.copy()
+        cache_settings["eviction_policy"] = "none"
+        cache_settings["size_limit"] = int(1e12)
+        cache_settings["cull_limit"] = 0
+        openai_responses = dc.Cache(self.cache_path, **cache_settings)
+
+        if (self.openai_model, message) in openai_responses:
+            reply = openai_responses[(self.openai_model, message)]
 
-        if message in openai_responses.get(self.openai_model, {}).keys():
-            reply = openai_responses[self.openai_model][message]
         else:
             # Ask openai
             if openai.api_key is None:
@@ -59,18 +55,10 @@ def ask(self, message: str) -> str:
                 {"role": "user", "content": message},
             ]
             chat = self._send_request(messages)
-
             reply = chat.choices[0].message.content
 
-            # add reply to cache
-            with self.cache_lock:
-                with open(self.cache_path, "r") as f:
-                    openai_responses = json.load(f)
-                if self.openai_model not in openai_responses.keys():
-                    openai_responses[self.openai_model] = {}
-                openai_responses[self.openai_model][message] = reply
-                with open(self.cache_path, "w") as f:
-                    json.dump(openai_responses, f)
+            openai_responses[(self.openai_model, message)] = reply
+            openai_responses.close()
 
         if "please provide" in reply.lower():
             return ""