Fix AUROC Bug and allow custom Example Selectors

src/autolabel/dataset/dataset.py

@@ -110,10 +110,6 @@ def process_labels(
             x.successfully_labeled for x in llm_labels
         ]
 
-        self.df[self.generate_label_name("annotation")] = [
-            pickle.dumps(x) for x in llm_labels
-        ]
-
         # Add row level LLM metrics to the dataframe
         if metrics is not None:
             for metric in metrics:

src/autolabel/labeler.py

@@ -15,7 +15,7 @@
 from autolabel.dataset import AutolabelDataset
 from autolabel.data_models import AnnotationModel, TaskRunModel
 from autolabel.database import StateManager
-from autolabel.few_shot import ExampleSelectorFactory
+from autolabel.few_shot import ExampleSelectorFactory, BaseExampleSelector
 from autolabel.models import BaseModel, ModelFactory
 from autolabel.metrics import BaseMetric
 from autolabel.transforms import BaseTransform, TransformFactory
@@ -52,6 +52,7 @@ def __init__(
         self,
         config: Union[AutolabelConfig, str, dict],
         cache: Optional[bool] = True,
+        example_selector: Optional[BaseExampleSelector] = None,
     ) -> None:
         self.db = StateManager()
         self.generation_cache = SQLAlchemyGenerationCache() if cache else None
@@ -67,6 +68,7 @@ def __init__(
         self.confidence = ConfidenceCalculator(
             score_type="logprob_average", llm=self.llm
         )
+        self.example_selector = example_selector
 
         if in_notebook():
             import nest_asyncio
@@ -139,12 +141,13 @@ def run(
                 f"Explanation column {self.config.explanation_column()} not found in dataset.\nMake sure that explanations were generated using labeler.generate_explanations(seed_file)."
             )
 
-        self.example_selector = ExampleSelectorFactory.initialize_selector(
-            self.config,
-            seed_examples,
-            dataset.df.keys().tolist(),
-            cache=self.generation_cache is not None,
-        )
+        if self.example_selector is None:
+            self.example_selector = ExampleSelectorFactory.initialize_selector(
+                self.config,
+                seed_examples,
+                dataset.df.keys().tolist(),
+                cache=self.generation_cache is not None,
+            )
 
         num_failures = 0
         current_index = self.task_run.current_index

src/autolabel/metrics/auroc.py

@@ -1,6 +1,7 @@
 from typing import List
 
 from sklearn.metrics import roc_auc_score
+import numpy as np
 
 from autolabel.metrics import BaseMetric
 from autolabel.schema import LLMAnnotation, MetricResult, MetricType
@@ -26,7 +27,11 @@ def compute(
         ]
         confidence = [llm_label.confidence_score for llm_label in filtered_llm_labels]
 
-        auroc = roc_auc_score(match, confidence)
+        if np.unique(match).shape[0] == 1:
+            # all labels are the same
+            auroc = 1 if match[0] == 1 else 0
+        else:
+            auroc = roc_auc_score(match, confidence)
 
         value = [
             MetricResult(