Security test + benchmark

pyproject.toml

@@ -90,6 +90,7 @@ prometheus-client = "^0.21.1"
 llama-api-client = "^0.1.1"
 huggingface-hub = "^0.30.2"
 openai = "^1.8.0"
+pyarrow = "^20.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest-datafiles = "^3.0.0"

src/modelbench/benchmarks.py

@@ -9,7 +9,7 @@
 from modelgauge.prompt_sets import validate_prompt_set
 from modelgauge.sut import PromptResponseSUT
 
-from modelbench.hazards import HazardDefinition, HazardScore, SafeHazardV1, Standards, STANDARDS
+from modelbench.hazards import HazardDefinition, HazardScore, SafeHazardV1, SecurityHazard, Standards, STANDARDS
 from modelbench.scoring import LetterGradeMixin, score_to_ordinal_grade
 from modelbench.uid import HasUid
 
@@ -159,3 +159,23 @@ def _make_hazards(self) -> Sequence[HazardDefinition]:
         "prompt_set": "self.prompt_set",
         "evaluator": "self.evaluator",
     }
+
+
+class SecurityBenchmark(BenchmarkDefinition):
+    def __init__(self, evaluator: str = "default"):
+        self.evaluator = evaluator
+        super().__init__()
+
+    def key(self):
+        key_def = self._uid_definition.copy()
+        del key_def["evaluator"]
+        return HasUid._render_uid(self, key_def).replace(".", "_")
+
+    def _make_hazards(self) -> Sequence[HazardDefinition]:
+        return [SecurityHazard(hazard_key, self.evaluator) for hazard_key in SecurityHazard.all_hazard_keys]
+
+    _uid_definition = {
+        "class": "security_benchmark",
+        "version": "0.1",
+        "evaluator": "self.evaluator",
+    }

src/modelbench/hazards.py

@@ -6,14 +6,16 @@
 from typing import List, Mapping
 
 import casefy
-from modelgauge.base_test import BaseTest, PromptResponseTest
+from pydantic import BaseModel, ConfigDict
+
+from modelgauge.base_test import PromptResponseTest
 from modelgauge.locales import display_for, validate_locale
 from modelgauge.prompt_sets import validate_prompt_set
 from modelgauge.records import TestRecord
 from modelgauge.secret_values import RawSecrets
 from modelgauge.test_registry import TESTS
 from modelgauge.tests.safe_v1 import BaseSafeTestVersion1, SafeTestResult
-from pydantic import BaseModel, ConfigDict
+from modelgauge.tests.security import BaseSecurityTest, SecurityTestResult
 
 from modelbench.scoring import LetterGradeMixin, NumericGradeMixin, ValueEstimate
 from modelbench.uid import HasUid
@@ -67,7 +69,7 @@ def __repr__(self):
     _uid_definition = {"class": key}
 
 
-class SafeHazardV1(HazardDefinition, ABC):
+class SafeHazardV1(HazardDefinition):
     """Hazard definition class for v1.0."""
 
     all_hazard_keys = BaseSafeTestVersion1.hazards
@@ -140,6 +142,65 @@ def test_uids(self) -> List[str]:
     }
 
 
+class SecurityHazard(HazardDefinition):
+    all_hazard_keys = BaseSecurityTest.hazards
+
+    def __init__(self, hazard_key: str, evaluator="default"):
+        assert hazard_key in self.all_hazard_keys, f"Unknown hazard key {hazard_key}."
+        self.hazard_key = hazard_key
+        self._evaluator = evaluator
+        super().__init__()
+
+    def name(self):
+        return f"Security {self.hazard_key.title()} Hazard V0.1"
+
+    def key(self):
+        return f"security_hazard-0_1-{self.hazard_key}"
+
+    def score(self, sut_scores: Mapping[str, TestRecord]) -> "HazardScore":
+        test_scores: dict[str, ValueEstimate] = {}
+        num_exceptions = 0
+        num_scored_items = 0
+        num_safe_items = 0
+        for key, test_record in sut_scores.items():
+            result = test_record.result.to_instance(SecurityTestResult)
+            num_scored_items += result.num_items
+            num_safe_items += result.num_items * result.frac_safe
+            test_scores[key] = ValueEstimate.make(result.frac_safe, result.num_items)
+            num_exceptions += len(test_record.test_item_exceptions)
+
+        score_values = list(test_scores.values())
+        if not score_values:
+            warnings.warn(f"No data items to score for {self.uid}")
+            return HazardScore(
+                hazard_definition=self,
+                score=ValueEstimate(lower=0, upper=0, estimate=0, samples=0),
+                test_scores=test_scores,
+                exceptions=num_exceptions,
+                num_scored_items=num_scored_items,
+                num_safe_items=round(num_safe_items),
+            )
+        score_estimate = ValueEstimate.combine(score_values)
+        test_scores = {self.uid: score_estimate}
+        return HazardScore(
+            hazard_definition=self,
+            score=score_estimate,
+            test_scores=test_scores,
+            exceptions=num_exceptions,
+            num_scored_items=num_scored_items,
+            num_safe_items=round(num_safe_items),
+        )
+
+    def test_uids(self) -> List[str]:
+        return [BaseSecurityTest.create_uid(self.hazard_key, self._evaluator)]
+
+    _uid_definition = {
+        "name": "security_hazard",
+        "version": "0.1",
+        "hazard": "self.hazard_key",
+    }
+
+
 class HazardScore(BaseModel, LetterGradeMixin, NumericGradeMixin):
     model_config = ConfigDict(arbitrary_types_allowed=True)
     hazard_definition: HazardDefinition

src/modelgauge/tests/security.py

@@ -0,0 +1,134 @@
+from abc import ABC
+from typing import Dict, List, Mapping
+
+import pyarrow.parquet as pq  # type: ignore
+from pydantic import BaseModel
+
+from modelgauge.aggregations import get_measurement_stats
+from modelgauge.annotator_set import AnnotatorSet
+from modelgauge.base_test import PromptResponseTest
+from modelgauge.default_annotator_set import DEFAULT_ANNOTATOR_SET
+from modelgauge.dependency_helper import DependencyHelper
+from modelgauge.external_data import ExternalData, WebData
+from modelgauge.prompt import TextPrompt
+from modelgauge.secret_values import InjectSecret, RequiredSecret, SecretDescription
+from modelgauge.single_turn_prompt_response import MeasuredTestItem, SUTResponseAnnotations, TestItem
+from modelgauge.sut import SUTOptions
+from modelgauge.sut_capabilities import AcceptsTextPrompt
+from modelgauge.test_decorator import modelgauge_test
+from modelgauge.test_registry import TESTS
+
+
+class RequiredModellabFileDownloadToken(RequiredSecret):
+    @classmethod
+    def description(cls) -> SecretDescription:
+        return SecretDescription(
+            scope="modellab_files",
+            key="token",
+            instructions="Please ask MLCommons admin for permission.",
+        )
+
+
+class SecurityTestResult(BaseModel):
+    frac_safe: float
+    num_items: int
+
+
+class BaseSecurityTest(PromptResponseTest, ABC):
+    annotators: AnnotatorSet  # Must be set by concrete subclasses.
+    hazards = ["cse"]  # For now, only one hazard in the sample file.
+
+    @staticmethod
+    def create_uid(hazard: str, evaluator=None):
+        if evaluator is None or evaluator == "default":
+            postfix = ""
+        else:
+            postfix = "-" + evaluator
+        uid = f"security-{hazard}-0.1{postfix}".lower()
+        return uid
+
+    def __init__(
+        self,
+        uid: str,
+        hazard: str,
+        token: RequiredModellabFileDownloadToken,
+    ):
+        assert hazard in self.hazards, f"Invalid hazard type {hazard}."
+        super().__init__(uid)
+        self.token = token.value
+        self.hazard = hazard
+
+    def sut_options(self) -> SUTOptions:
+        return SUTOptions(
+            max_tokens=500,
+            temperature=0.01,
+        )
+
+    @classmethod
+    def get_annotators(cls) -> List[str]:
+        return cls.annotators.annotators
+
+    def get_dependencies(self) -> Mapping[str, ExternalData]:
+        # TODO: Move prompts to prod when we get real ones. This is just a small sample file..
+        prompt_set_base_download_url = "https://modellab.dev.modelmodel.org/files/download"
+        headers = None
+        if self.token is not None:
+            headers = {"auth-token": self.token}
+        return {
+            "prompts": WebData(
+                source_url=f"{prompt_set_base_download_url}/security_results.parquet",
+                headers=headers,
+            )
+        }
+
+    def make_test_items(self, dependency_helper: DependencyHelper) -> List[TestItem]:
+        test_items: List[TestItem] = []
+        local_file = dependency_helper.get_local_path("prompts")
+        parquet_file = pq.ParquetFile(local_file)
+        for batch in parquet_file.iter_batches(batch_size=512):
+            for row in batch.to_pylist():
+                if row["prompt_hazard"] != self.hazard:
+                    continue
+                test_items.append(
+                    TestItem(
+                        prompt=TextPrompt(text=row["attack_prompt_text"]),
+                        source_id=row["release_prompt_id"],
+                    ),
+                )
+        return test_items
+
+    def measure_quality(self, item: SUTResponseAnnotations) -> Dict[str, float]:
+        return self.annotators.evaluate(item)
+
+    def aggregate_measurements(self, items: List[MeasuredTestItem]) -> SecurityTestResult:
+        result = get_measurement_stats("is_safe", items)
+        return SecurityTestResult(frac_safe=result.mean, num_items=result.count)
+
+
+@modelgauge_test(requires_sut_capabilities=[AcceptsTextPrompt])
+class SecurityTest(BaseSecurityTest):
+    annotators = DEFAULT_ANNOTATOR_SET
+
+
+hazard = "cse"  # only 1 hazard in the sample file.
+TESTS.register(
+    SecurityTest, BaseSecurityTest.create_uid(hazard), hazard, InjectSecret(RequiredModellabFileDownloadToken)
+)
+# Try to register test with private "ensemble" annotator.
+try:
+    from modelgauge.private_ensemble_annotator_set import PRIVATE_ANNOTATOR_SET
+
+    @modelgauge_test(requires_sut_capabilities=[AcceptsTextPrompt])
+    class PrivateSafeTestVersion1(BaseSecurityTest):
+        annotators = PRIVATE_ANNOTATOR_SET
+
+    TESTS.register(
+        SecurityTest,
+        BaseSecurityTest.create_uid(hazard, "ensemble"),
+        hazard,
+        InjectSecret(RequiredModellabFileDownloadToken),
+    )
+
+
+except Exception as e:
+    pass