use tqdm

google · Jul 6, 2024 · b6fd96c · b6fd96c
1 parent cbf5d5a
commit b6fd96c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/DiarizationLM/postprocess_completions.py b/DiarizationLM/postprocess_completions.py
@@ -9,6 +9,7 @@
 
 from absl import app
 from absl import flags
+import tqdm
 
 from diarizationlm import utils
 
@@ -40,7 +41,7 @@ def main(argv: Sequence[str]) -> None:
   po = utils.PromptOptions(completion_suffix=FLAGS.completion_suffix)
 
   # Process utterances.
-  for utt in data_dict["utterances"]:
+  for utt in tqdm.tqdm(data_dict["utterances"]):
     utils.postprocess_completions_for_utt(
         utt,
         llm_text_field="llm_text",

diff --git a/DiarizationLM/run_finetuned_gpt.py b/DiarizationLM/run_finetuned_gpt.py
@@ -4,6 +4,7 @@
 import json
 from absl import app
 from absl import flags
+import tqdm
 import openai
 from diarizationlm import utils
 
@@ -49,7 +50,7 @@ def main(argv: Sequence[str]) -> None:
       prompt_suffix=FLAGS.prompt_suffix,
   )
 
-  for utt in data_dict["utterances"]:
+  for utt in tqdm.tqdm(ata_dict["utterances"]):
     prompts = utils.generate_prompts(utt, po=po)
 
     utt["completions"] = []