From 878c62639ecfc30daad2fcff53f6b0a0c0b7c8ed Mon Sep 17 00:00:00 2001
From: Barbara Korycki <bkorycki@gmail.com>
Date: Mon, 13 Jan 2025 15:02:19 -0500
Subject: [PATCH] Private ensemble doesn't worry about secrets. (#786)

* LG3 Lora and mistral 7b ruby don't need manual secret injection

* prompt engineered annotators do their own secret injection
---
 src/modelgauge/auth/vllm_keys.py              | 41 -------------------
 .../private_ensemble_annotator_set.py         | 35 ++--------------
 2 files changed, 4 insertions(+), 72 deletions(-)
 delete mode 100644 src/modelgauge/auth/vllm_keys.py

diff --git a/src/modelgauge/auth/vllm_keys.py b/src/modelgauge/auth/vllm_keys.py
deleted file mode 100644
index a772ba0f..00000000
--- a/src/modelgauge/auth/vllm_keys.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from modelgauge.secret_values import RequiredSecret, SecretDescription
-
-
-class Lg3LoraVllmApiKey(RequiredSecret):
-    @classmethod
-    def description(cls) -> SecretDescription:
-        return SecretDescription(
-            scope="lg3_lora_vllm",
-            key="api_key",
-            instructions="Contact MLCommons admin for access.",
-        )
-
-
-class Lg3LoraVllmEndpointUrl(RequiredSecret):
-    @classmethod
-    def description(cls) -> SecretDescription:
-        return SecretDescription(
-            scope="lg3_lora_vllm",
-            key="endpoint_url",
-            instructions="Contact MLCommons admin for access.",
-        )
-
-
-class Mistral7bVllmApiKey(RequiredSecret):
-    @classmethod
-    def description(cls) -> SecretDescription:
-        return SecretDescription(
-            scope="mistral_7b_vllm",
-            key="api_key",
-            instructions="Contact MLCommons admin for access.",
-        )
-
-
-class Mistral7bVllmEndpointUrl(RequiredSecret):
-    @classmethod
-    def description(cls) -> SecretDescription:
-        return SecretDescription(
-            scope="mistral_7b_vllm",
-            key="endpoint_url",
-            instructions="Contact MLCommons admin for access.",
-        )
diff --git a/src/modelgauge/private_ensemble_annotator_set.py b/src/modelgauge/private_ensemble_annotator_set.py
index 5783335f..0a249c70 100644
--- a/src/modelgauge/private_ensemble_annotator_set.py
+++ b/src/modelgauge/private_ensemble_annotator_set.py
@@ -3,12 +3,6 @@
 
 from modelgauge.annotator_set import AnnotatorSet
 from modelgauge.auth.together_key import TogetherApiKey
-from modelgauge.auth.vllm_keys import (
-    Lg3LoraVllmApiKey,
-    Lg3LoraVllmEndpointUrl,
-    Mistral7bVllmApiKey,
-    Mistral7bVllmEndpointUrl,
-)
 from modelgauge.dependency_injection import _replace_with_injected
 from modelgauge.secret_values import InjectSecret
 from modelgauge.single_turn_prompt_response import TestItemAnnotations
@@ -17,21 +11,19 @@
     from modelgauge.annotators.llama_3_70b_chat_annotator import (
         LLAMA_3_70B_PE_TAMALPAIS_2024_09_09_CONFIG as LLAMA_3_70B_CONFIG,
     )  # type: ignore
-    from modelgauge.annotators.llama_guard_3_lora_annotator import (
-        LlamaGuard3LoRAAnnotator,
-    )  # type: ignore
     from modelgauge.annotators.template_lg3_lora_annotator import (
         config as LG3_LORA_CONFIG,
+        Lg3LoraVllmApiKey,
+        Lg3LoraVllmEndpointUrl,
     )  # type: ignore
     from modelgauge.annotators.mistral_7b_ruby_annotator import (
         config as MISTRAL_7B_RUBY_CONFIG,
+        Mistral7bVllmApiKey,
+        Mistral7bVllmEndpointUrl,
     )  # type: ignore
     from modelgauge.annotators.mistral_8x22b_instruct_annotator import (
         MISTRAL_8x22B_PE_TAMALPAIS_2024_09_09_CONFIG as MISTRAL_8x22B_CONFIG,
     )  # type: ignore
-    from modelgauge.annotators.prompt_engineered_annotator import (
-        PromptEngineeredAnnotator,
-    )  # type: ignore
     from modelgauge.ensemble.majority_vote_ensemble_strategy import MajorityVoteEnsembleStrategy  # type: ignore
     from modelgauge.safety_model_response import SafetyModelResponse  # type: ignore
 except:
@@ -53,12 +45,6 @@ class EnsembleAnnotatorSet(AnnotatorSet):
     def __init__(self, secrets):
         self.secrets = secrets
         self.strategy = MajorityVoteEnsembleStrategy()
-        self.__configure_vllm_annotators()
-        self.__configure_huggingface_annotators()
-        self.__configure_together_annotators()
-        self._register_annotators()
-
-    def _register_annotators(self):
         self.annotators = [
             MISTRAL_7B_RUBY_CONFIG.name,
             MISTRAL_8x22B_CONFIG.name,
@@ -66,19 +52,6 @@ def _register_annotators(self):
             LG3_LORA_CONFIG.name,
         ]
 
-    def __configure_vllm_annotators(self):
-        LG3_LORA_CONFIG.api_key = self.secrets["lg3_vllm_api_key"].value
-        LG3_LORA_CONFIG.base_url = self.secrets["lg3_vllm_endpoint_url"].value
-        MISTRAL_7B_RUBY_CONFIG.api_key = self.secrets["mistral_7b_vllm_api_key"].value
-        MISTRAL_7B_RUBY_CONFIG.base_url = self.secrets["mistral_7b_vllm_endpoint_url"].value
-
-    def __configure_huggingface_annotators(self):
-        return
-
-    def __configure_together_annotators(self):
-        MISTRAL_8x22B_CONFIG.llm_config.api_key = self.secrets["together_api_key"]
-        LLAMA_3_70B_CONFIG.llm_config.api_key = self.secrets["together_api_key"]
-
     def evaluate(self, item: TestItemAnnotations) -> Dict[str, float]:
         annotated_completion = item.interactions[0].response.completions[0]
         individual_annotations = {}