From fcb9464e345f48888f2419e0913b543e4d030b3d Mon Sep 17 00:00:00 2001
From: kcz358 <kaichenzhang358@outlook.com>
Date: Tue, 26 Nov 2024 05:28:31 +0000
Subject: [PATCH] Fix wavcaps bugs

---
 lmms_eval/tasks/wavcaps/utils.py | 34 +++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/lmms_eval/tasks/wavcaps/utils.py b/lmms_eval/tasks/wavcaps/utils.py
index 020a8f8b..09a4333f 100644
--- a/lmms_eval/tasks/wavcaps/utils.py
+++ b/lmms_eval/tasks/wavcaps/utils.py
@@ -32,8 +32,8 @@ def wavcaps_doc_to_text(doc, lmms_eval_specific_kwargs):
     config = yaml.safe_load("".join(safe_data))
 
 
-NUM_SECONDS_TO_SLEEP = 2
-GPT_EVAL_MODEL_NAME = config["metadata"]["api_version"]
+NUM_SECONDS_TO_SLEEP = 5
+GPT_EVAL_MODEL_NAME = config["metadata"]["gpt_eval_model_name"]
 API_TYPE = os.getenv("API_TYPE", "azure")
 
 if API_TYPE == "openai":
@@ -80,24 +80,30 @@ def wavcaps_doc_to_text(doc, lmms_eval_specific_kwargs):
 def get_eval(max_tokens: int, content: str):
     global headers
 
-    time.sleep(NUM_SECONDS_TO_SLEEP)
     messages = [
         {"role": "user", "content": content},
     ]
 
     payload = {"model": GPT_EVAL_MODEL_NAME, "messages": messages, "temperature": 0, "max_tokens": max_tokens, "n": 1}
 
-    try:
-        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
-        response.raise_for_status()
-        response_data = response.json()
-        content = response_data["choices"][0]["message"]["content"].strip()
-
-        if content != "":
-            return content, payload["model"]
-    except Exception as e:
-        eval_logger.info(f"Attempt failed with error: {e}")
-        return "", ""
+    for attempt in range(5):
+        try:
+            response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
+            response.raise_for_status()
+            response_data = response.json()
+
+            content = response_data["choices"][0]["message"]["content"].strip()
+            if content != "":
+                return content, response_data["model"]
+            break  # If successful, break out of the loop
+
+        except Exception as e:
+            eval_logger.info(f"Attempt {attempt + 1} failed with error: {e}")
+            if attempt < 5:  # If we have retries left, sleep and then continue to next attempt
+                time.sleep(NUM_SECONDS_TO_SLEEP)
+            else:  # If this was the last attempt, log and return empty
+                eval_logger.error(f"All 5 attempts failed. Last error message: {e}")
+                return "", ""
     return "", ""