Fix claude always retrying 5 time error

EvolvingLMMs-Lab · Sep 24, 2024 · c01d2f2 · c01d2f2
1 parent a134a74
commit c01d2f2
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/lmms_eval/models/claude.py b/lmms_eval/models/claude.py
@@ -232,8 +232,10 @@ def generate_until(self, requests) -> List[str]:
                 gen_kwargs["num_beams"] = 1
 
             for attempt in range(5):
+                retry_flag = True
                 try:
                     message = client.messages.create(model=self.model_version, max_tokens=gen_kwargs["max_new_tokens"], system=self.system_prompt, temperature=gen_kwargs["temperature"], top_p=gen_kwargs["top_p"], messages=messages)
+                    retry_flag = False
                 except Exception as e:
                     eval_logger.info(f"Attempt {attempt + 1} failed with error: {str(e)}")
                     if attempt < 5 - 1:  # If we have retries left, sleep and then continue to next attempt
@@ -243,6 +245,9 @@ def generate_until(self, requests) -> List[str]:
                         res.append("")
                         pbar.update(1)
                         continue
+                if not retry_flag:
+                    break
+                eval_logger.info("Retrying...")
 
             response_text = message.content[0].text
             res.append(message.content[0].text)