From 7789de4278453a369d06d4dbb6d4c8026534eef1 Mon Sep 17 00:00:00 2001
From: Jimil Patel <jimilp7@gmail.com>
Date: Fri, 29 Sep 2023 16:44:07 -0400
Subject: [PATCH] Issue#17 - Add Progress Bar to show progress on perturbations
 and scores

---
 auditor/evaluation/discriminative.py    |  5 +++++
 auditor/evaluation/expected_behavior.py |  7 +++++++
 auditor/evaluation/generative.py        |  6 ++++++
 auditor/utils/progress_logger.py        | 15 +++++++++++++++
 4 files changed, 33 insertions(+)
 create mode 100644 auditor/utils/progress_logger.py

diff --git a/auditor/evaluation/discriminative.py b/auditor/evaluation/discriminative.py
index fcd2cd7..e1a00a5 100644
--- a/auditor/evaluation/discriminative.py
+++ b/auditor/evaluation/discriminative.py
@@ -15,6 +15,7 @@
 )
 from auditor.reporting import generate_robustness_report
 from auditor.utils.logging import get_logger
+from auditor.utils.progress_logger import ProgressLogger
 
 LOG = get_logger(__name__)
 
@@ -90,6 +91,8 @@ def evaluate(
             f'Started model evaluation with perturbation type '
             f'{self.perturbed_dataset.perturbation_type}'
         )
+        progress_bar = ProgressLogger(total_steps=min(len(self.perturbed_dataset.data), len(self.perturbed_dataset.metadata)), description="Starting Model Evaluation")
+
         for perturbed_samples, metadata_samples in zip(
             self.perturbed_dataset.data, self.perturbed_dataset.metadata
         ):
@@ -126,12 +129,14 @@ def evaluate(
                         metadata=mdata,
                     )
                 )
+            progress_bar.update()
         robust_accuracy = self.compute_accuracy(test_results)
         LOG.info(f'Robust Accuracy: {robust_accuracy*100.}')
         LOG.info(
             'Completed model evaluation with perturbation type '
             f'{self.perturbed_dataset.perturbation_type}'
         )
+        progress_bar.close()
         self.test_results = TestSummary(
             results=test_results,
             robust_accuracy=robust_accuracy,
diff --git a/auditor/evaluation/expected_behavior.py b/auditor/evaluation/expected_behavior.py
index 7f692d1..95bace8 100644
--- a/auditor/evaluation/expected_behavior.py
+++ b/auditor/evaluation/expected_behavior.py
@@ -4,6 +4,7 @@
 import numpy as np
 from sentence_transformers.SentenceTransformer import SentenceTransformer
 
+from auditor.utils.progress_logger import ProgressLogger
 from auditor.utils.similarity import compute_similarity
 from auditor.utils.logging import get_logger
 
@@ -163,6 +164,8 @@ def check(
         reference_generation: str,
     ) -> List[Tuple[bool, Dict[str, float]]]:
         test_results = []
+        progress_bar = ProgressLogger(total_steps=len(perturbed_generations), description="Fetching Scores")
+
         for peturbed_gen in perturbed_generations:
             try:
                 score = compute_similarity(
@@ -178,9 +181,13 @@ def check(
                     self.similarity_metric_key: round(score, ndigits=2)
                 }
                 test_results.append((test_status, score_dict))
+                progress_bar.update()
             except Exception as e:
                 LOG.error('Unable to complete semanatic similarity checks')
+                progress_bar.close()
                 raise e
+
+        progress_bar.close()
         return test_results
 
     def behavior_description(self):
diff --git a/auditor/evaluation/generative.py b/auditor/evaluation/generative.py
index 42b1971..ad681f9 100644
--- a/auditor/evaluation/generative.py
+++ b/auditor/evaluation/generative.py
@@ -12,6 +12,7 @@
 from auditor.utils.logging import get_logger
 from auditor.perturbations import Paraphrase
 from auditor.perturbations import TransformBase
+from auditor.utils.progress_logger import ProgressLogger
 
 LOG = get_logger(__name__)
 
@@ -98,6 +99,8 @@ def _evaluate_generations(
         else:
             evaluate_prompts = prompt_perturbations
 
+        progress_bar = ProgressLogger(total_steps=len(evaluate_prompts), description="Applying Perturbations")
+
         # generations for each of the perturbed prompts
         alternative_generations = []
         for alt_prompt in evaluate_prompts:
@@ -107,6 +110,9 @@ def _evaluate_generations(
                 post_context,
             )
             alternative_generations.append(resp)
+            progress_bar.update()
+
+        progress_bar.close()
 
         # create test result
         metric = self.expected_behavior.check(
diff --git a/auditor/utils/progress_logger.py b/auditor/utils/progress_logger.py
new file mode 100644
index 0000000..3699632
--- /dev/null
+++ b/auditor/utils/progress_logger.py
@@ -0,0 +1,15 @@
+import tqdm
+
+class ProgressLogger:
+    """class to show progress bar"""
+    def __init__(self, total_steps, description="Logging..."):
+        self.total_steps = total_steps
+        self.description = description
+
+        self.pbar = tqdm.tqdm(total=total_steps, desc=description)
+
+    def update(self, incremental=1):
+        self.pbar.update(incremental)
+
+    def close(self):
+        self.pbar.close()
\ No newline at end of file