added individual retrying for failed batches (#443)

* added individual retrying for failed batches * consolidated _label_individually to base llm
refuel-ai · Jul 11, 2023 · d8f1b15 · d8f1b15
1 parent 82ebd5d
commit d8f1b15
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 21 deletions.
diff --git a/src/autolabel/models/anthropic.py b/src/autolabel/models/anthropic.py
@@ -5,7 +5,7 @@
 from autolabel.models import BaseModel
 from autolabel.cache import BaseCache
 from langchain.chat_models import ChatAnthropic
-from langchain.schema import Generation, LLMResult, HumanMessage
+from langchain.schema import LLMResult, HumanMessage
 
 
 class AnthropicLLM(BaseModel):
@@ -40,12 +40,10 @@ def __init__(self, config: AutolabelConfig, cache: BaseCache = None) -> None:
     def _label(self, prompts: List[str]) -> LLMResult:
         prompts = [[HumanMessage(content=prompt)] for prompt in prompts]
         try:
-            response = self.llm.generate(prompts)
-            return response
+            return self.llm.generate(prompts)
         except Exception as e:
-            print(f"Error generating from LLM: {e}, returning empty result")
-            generations = [[Generation(text="")] for _ in prompts]
-            return LLMResult(generations=generations)
+            print(f"Error generating from LLM: {e}, retrying each prompt individually")
+            return self._label_individually(prompts)
 
     def get_cost(self, prompt: str, label: Optional[str] = "") -> float:
         num_prompt_toks = tokenizer.count_tokens(prompt)

diff --git a/src/autolabel/models/base.py b/src/autolabel/models/base.py
@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from typing import List, Optional, Dict, Tuple
 
-from langchain.schema import LLMResult
+from langchain.schema import LLMResult, Generation
 
 from autolabel.configs import AutolabelConfig
 from autolabel.schema import CacheEntry
@@ -52,6 +52,26 @@ def label(self, prompts: List[str]) -> Tuple[LLMResult, float]:
         generations = [existing_prompts[i] for i in range(len(prompts))]
         return LLMResult(generations=generations, llm_output=llm_output), cost
 
+    def _label_individually(self, prompts: List[str]) -> LLMResult:
+        """Label each prompt individually. Should be used only after trying as a batch first.
+
+        Args:
+            prompts (List[str]): List of prompts to label
+
+        Returns:
+            LLMResult: LLMResult object with generations
+        """
+        generations = []
+        for prompt in prompts:
+            try:
+                response = self.llm.generate([prompt])
+                generations.append(response.generations[0])
+            except Exception as e:
+                print(f"Error generating from LLM: {e}, returning empty generation")
+                generations.append([Generation(text="")])
+
+        return LLMResult(generations=generations)
+
     @abstractmethod
     def _label(self, prompts: List[str]) -> LLMResult:
         # TODO: change return type to do parsing in the Model class

diff --git a/src/autolabel/models/cohere.py b/src/autolabel/models/cohere.py
@@ -41,9 +41,8 @@ def _label(self, prompts: List[str]) -> LLMResult:
         try:
             return self.llm.generate(prompts)
         except Exception as e:
-            print(f"Error generating from LLM: {e}, returning empty result")
-            generations = [[Generation(text="")] for _ in prompts]
-            return LLMResult(generations=generations)
+            print(f"Error generating from LLM: {e}, retrying each prompt individually")
+            return self._label_individually(prompts)
 
     def get_cost(self, prompt: str, label: Optional[str] = "") -> float:
         num_prompt_toks = len(self.co.tokenize(prompt).tokens)

diff --git a/src/autolabel/models/hf_pipeline.py b/src/autolabel/models/hf_pipeline.py
@@ -1,6 +1,6 @@
 from typing import List, Optional
 from langchain.llms import HuggingFacePipeline
-from langchain.schema import LLMResult, Generation
+from langchain.schema import LLMResult
 
 from autolabel.models import BaseModel
 from autolabel.configs import AutolabelConfig
@@ -69,9 +69,8 @@ def _label(self, prompts: List[str]) -> LLMResult:
         try:
             return self.llm.generate(prompts)
         except Exception as e:
-            print(f"Error generating from LLM: {e}, returning empty result")
-            generations = [[Generation(text="")] for _ in prompts]
-            return LLMResult(generations=generations)
+            print(f"Error generating from LLM: {e}, retrying each prompt individually")
+            return self._label_individually(prompts)
 
     def get_cost(self, prompt: str, label: Optional[str] = "") -> float:
         # Model inference for this model is being run locally

diff --git a/src/autolabel/models/openai.py b/src/autolabel/models/openai.py
@@ -4,7 +4,7 @@
 
 from langchain.chat_models import ChatOpenAI
 from langchain.llms import OpenAI
-from langchain.schema import LLMResult, HumanMessage, Generation
+from langchain.schema import LLMResult, HumanMessage
 import tiktoken
 
 from autolabel.models import BaseModel
@@ -146,9 +146,8 @@ def _label(self, prompts: List[str]) -> LLMResult:
         try:
             return self.llm.generate(prompts)
         except Exception as e:
-            print(f"Error generating from LLM: {e}, returning empty result")
-            generations = [[Generation(text="")] for _ in prompts]
-            return LLMResult(generations=generations)
+            print(f"Error generating from LLM: {e}, retrying each prompt individually")
+            return self._label_individually(prompts)
 
     def get_cost(self, prompt: str, label: Optional[str] = "") -> float:
         encoding = tiktoken.encoding_for_model(self.model_name)

diff --git a/src/autolabel/models/palm.py b/src/autolabel/models/palm.py
@@ -66,6 +66,30 @@ def __init__(self, config: AutolabelConfig, cache: BaseCache = None) -> None:
     def _label_with_retry(self, prompts: List[str]) -> LLMResult:
         return self.llm.generate(prompts)
 
+    def _label_individually(self, prompts: List[str]) -> LLMResult:
+        """Label each prompt individually. Should be used only after trying as a batch first.
+
+        Args:
+            prompts (List[str]): List of prompts to label
+
+        Returns:
+            LLMResult: LLMResult object with generations
+        """
+        generations = []
+        for i, prompt in enumerate(prompts):
+            try:
+                response = self._label_with_retry([prompt])
+                for generation in response.generations[0]:
+                    generation.text = generation.text.replace(
+                        self.SEP_REPLACEMENT_TOKEN, "\n"
+                    )
+                generations.append(response.generations[0])
+            except Exception as e:
+                print(f"Error generating from LLM: {e}, returning empty generation")
+                generations.append([Generation(text="")])
+
+        return LLMResult(generations=generations)
+
     def _label(self, prompts: List[str]) -> LLMResult:
         for prompt in prompts:
             if self.SEP_REPLACEMENT_TOKEN in prompt:
@@ -94,9 +118,8 @@ def _label(self, prompts: List[str]) -> LLMResult:
                     )
             return result
         except Exception as e:
-            logger.error(f"Error generating from LLM: {e}.")
-        generations = [[Generation(text="")] for _ in prompts]
-        return LLMResult(generations=generations)
+            print(f"Error generating from LLM: {e}, retrying each prompt individually")
+            self._label_individually(prompts)
 
     def get_cost(self, prompt: str, label: Optional[str] = "") -> float:
         if self.model_name is None: