Benchmark small tweaks

GitOrigin-RevId: 80900bceef0034c7ce6bdb5965f48c21df0882f1
gretelai · Sep 9, 2024 · c6986dd · c6986dd
1 parent 01aae3c
commit c6986dd
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/src/gretel_trainer/benchmark/__init__.py b/src/gretel_trainer/benchmark/__init__.py
@@ -15,5 +15,6 @@
     GretelGPTX,
     GretelLSTM,
     GretelModel,
+    GretelNavigatorFT,
 )
 from gretel_trainer.benchmark.log import set_log_level
diff --git a/src/gretel_trainer/benchmark/entrypoints.py b/src/gretel_trainer/benchmark/entrypoints.py
@@ -92,7 +92,7 @@ def _entrypoint(
 def _verify_client_config(session: ClientConfig):
     try:
         current_user_email = session.email
-        logger.info(f"Using gretel client configured with {current_user_email}!r")
+        logger.info(f"Using gretel client configured with {current_user_email!r}")
     except Exception as e:
         raise BenchmarkException(
             "Invalid gretel client configuration, please make sure to configure the "

diff --git a/src/gretel_trainer/benchmark/session.py b/src/gretel_trainer/benchmark/session.py
@@ -30,6 +30,8 @@
     "PCS": "principal_component_stability",
     "FDS": "field_distribution_stability",
     "PPL": "privacy_protection_level",
+    "AIA": "attribute_inference_attack_score",
+    "DPS": "data_privacy_score",
 }
 
 FutureKeyT = Union[str, RunKey]