Refactors EnsembleAnnotatorSet initialization. (#1071)

* Refactors `EnsembleAnnotatorSet` initialization.

* Fix typing.

* Renames and relocates safety model response module.

Author: superdosh

src/modelgauge/ensemble_annotator_set.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+from modelgauge.annotator_set import AnnotatorSet
+from modelgauge.safety_model_response import EnsembleSafetyModelResponse, SafetyModelResponse
+from modelgauge.single_turn_prompt_response import SUTResponseAnnotations
+
+
+class EnsembleStrategy(ABC):
+
+    @abstractmethod
+    def compute_response(self, annotations: Dict[str, SafetyModelResponse]) -> EnsembleSafetyModelResponse:
+        pass
+
+
+# Dictionary to hold registered ensemble strategies
+ENSEMBLE_STRATEGIES: Dict[str, EnsembleStrategy] = {}
+
+
+class EnsembleAnnotatorSet(AnnotatorSet):
+    annotators: List[str] = []
+
+    def __init__(self, annotators: List[str], strategy: EnsembleStrategy):
+        self.strategy = strategy
+        self.annotators = annotators
+
+    def evaluate(self, item: SUTResponseAnnotations) -> Dict[str, float]:
+        individual_annotations = {}
+        for annotator_id in self.annotators:
+            if annotator_id in item.annotations:
+                individual_annotations[annotator_id] = item.get_annotation(annotator_id, SafetyModelResponse)
+        annotation = self.strategy.compute_response(individual_annotations)
+        return {"is_safe": float(annotation.is_safe)}

src/modelgauge/private_ensemble_annotator_set.py

@@ -1,8 +1,4 @@
-import os
-from typing import Any, Dict, List
-
-from modelgauge.annotator_set import AnnotatorSet
-from modelgauge.single_turn_prompt_response import SUTResponseAnnotations
+from modelgauge.ensemble_annotator_set import EnsembleAnnotatorSet
 
 try:
     from modelgauge.annotators.llama_3_70b_chat_annotator import (
@@ -23,25 +19,12 @@
     raise NotImplementedError("Private annotators are not available.")
 
 
-class EnsembleAnnotatorSet(AnnotatorSet):
-    annotators: List[str] = []
-
-    def __init__(self):
-        self.strategy = MajorityVoteEnsembleStrategy()
-        self.annotators = [
-            MISTRAL_7B_RUBY_CONFIG.name,
-            MISTRAL_8x22B_CONFIG.name,
-            LLAMA_3_70B_CONFIG.name,
-            LG3_LORA_CONFIG.name,
-        ]
-
-    def evaluate(self, item: SUTResponseAnnotations) -> Dict[str, float]:
-        individual_annotations = {}
-        for annotator_id in self.annotators:
-            if annotator_id in item.annotations:
-                individual_annotations[annotator_id] = item.get_annotation(annotator_id, SafetyModelResponse)
-        annotation = self.strategy.compute_response(individual_annotations)
-        return {"is_safe": float(annotation.is_safe)}
-
-
-PRIVATE_ANNOTATOR_SET = EnsembleAnnotatorSet()
+PRIVATE_ANNOTATOR_SET = EnsembleAnnotatorSet(
+    annotators=[
+        MISTRAL_7B_RUBY_CONFIG.name,
+        MISTRAL_8x22B_CONFIG.name,
+        LLAMA_3_70B_CONFIG.name,
+        LG3_LORA_CONFIG.name,
+    ],
+    strategy=MajorityVoteEnsembleStrategy(),
+)

src/modelgauge/safety_model_response.py

@@ -0,0 +1,27 @@
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+class SafetyModelResponse(BaseModel):
+    """Dataclass that all safety model responses should adhere to"""
+
+    is_safe: bool
+    safety_categories: List[str]
+    is_valid: bool = True
+    raw_input: Optional[str] = None
+    raw_response: Optional[str] = None
+    reasoning: Optional[str] = None
+    total_tokens_used: Optional[float] = None
+    logprobs: Optional[Any] = None
+
+
+class EnsembleSafetyModelResponse(SafetyModelResponse):
+    """Dataclass for safety model responses computed by an ensemble
+
+    Stores additional metadata on the original responses
+
+    **joined_responses**: Store the individual SafetyModelResponse objects
+    """
+
+    joined_responses: Dict[str, SafetyModelResponse]