VWA hide huggingface progress bar (#241)

ServiceNow · Nov 5, 2024 · 66cf402 · 66cf402
1 parent 4a76a46
commit 66cf402
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/browsergym/visualwebarena/src/browsergym/visualwebarena/task.py b/browsergym/visualwebarena/src/browsergym/visualwebarena/task.py
@@ -176,6 +176,15 @@ def setup(self, page: playwright.sync_api.Page) -> tuple[str, dict]:
             self.config_file = f.name
 
         # build the evaluator
+        from transformers.utils.logging import (
+            disable_progress_bar,
+            enable_progress_bar,
+            is_progress_bar_enabled,
+        )
+
+        hide_progress_bar = is_progress_bar_enabled()
+        if hide_progress_bar:
+            disable_progress_bar()
         captioning_fn = get_captioning_fn(
             device=self.eval_captioning_model_device,
             dtype=(
@@ -185,6 +194,8 @@ def setup(self, page: playwright.sync_api.Page) -> tuple[str, dict]:
             ),
             model_name="Salesforce/blip2-flan-t5-xl",
         )
+        if hide_progress_bar:
+            enable_progress_bar()
         self.evaluator = evaluator_router(self.config_file, captioning_fn=captioning_fn)
 
         # reset instance if needed (classifieds domain only)