sem_join + sem_filter pbars

lotus/models/lm.py

@@ -42,7 +42,10 @@ def __init__(
         self.cache = Cache(max_cache_size)
 
     def __call__(
-        self, messages: list[list[dict[str, str]]], safe_mode: bool = False, **kwargs: dict[str, Any]
+        self,
+        messages: list[list[dict[str, str]]],
+        show_pbar: bool = True,
+        **kwargs: dict[str, Any],
     ) -> LMOutput:
         all_kwargs = {**self.kwargs, **kwargs}
 
@@ -59,7 +62,7 @@ def __call__(
         self.stats.total_usage.cache_hits += len(messages) - len(uncached_data)
 
         # Process uncached messages in batches
-        uncached_responses = self._process_uncached_messages(uncached_data, all_kwargs)
+        uncached_responses = self._process_uncached_messages(uncached_data, all_kwargs, show_pbar)
 
         # Add new responses to cache
         for resp, (_, hash) in zip(uncached_responses, uncached_data):
@@ -74,12 +77,24 @@ def __call__(
 
         return LMOutput(outputs=outputs, logprobs=logprobs)
 
-    def _process_uncached_messages(self, uncached_data, all_kwargs):
+    def _process_uncached_messages(self, uncached_data, all_kwargs, show_pbar):
         """Processes uncached messages in batches and returns responses."""
         uncached_responses = []
-        for i in tqdm(range(0, len(uncached_data), self.max_batch_size), desc="Processing uncached messages"):
+        total_calls = len(uncached_data)
+
+        pbar = tqdm(
+            total=total_calls,
+            desc="Processing uncached messages",
+            disable=not show_pbar,
+            bar_format="{l_bar}{bar} {n}/{total} LM calls [{elapsed}<{remaining}, {rate_fmt}{postfix}]",
+        )
+        for i in range(0, total_calls, self.max_batch_size):
             batch = [msg for msg, _ in uncached_data[i : i + self.max_batch_size]]
             uncached_responses.extend(batch_completion(self.model, batch, drop_params=True, **all_kwargs))
+
+            pbar.update(len(batch))
+        pbar.close()
+
         return uncached_responses
 
     def _cache_response(self, response, hash):

lotus/sem_ops/sem_filter.py

@@ -24,6 +24,7 @@ def sem_filter(
     strategy: str | None = None,
     logprobs: bool = False,
     safe_mode: bool = False,
+    show_pbar: bool = True,
 ) -> SemanticFilterOutput:
     """
     Filters a list of documents based on a given user instruction using a language model.
@@ -55,7 +56,7 @@ def sem_filter(
         estimated_total_cost = sum(model.count_tokens(input) for input in inputs)
         show_safe_mode(estimated_total_cost, estimated_total_calls)
 
-    lm_output: LMOutput = model(inputs, **kwargs)
+    lm_output: LMOutput = model(inputs, show_pbar=show_pbar, **kwargs)
 
     postprocess_output = filter_postprocess(
         lm_output.outputs, default=default, cot_reasoning=strategy in ["cot", "zs-cot"]
@@ -235,6 +236,7 @@ def __call__(
                 logprobs=True,
                 strategy=helper_strategy,
                 safe_mode=safe_mode,
+                show_pbar=True,
             )
             helper_outputs, helper_logprobs = helper_output.outputs, helper_output.logprobs
             formatted_helper_logprobs: LogprobsForFilterCascade = (
@@ -338,6 +340,7 @@ def __call__(
                 cot_reasoning=cot_reasoning,
                 strategy=strategy,
                 safe_mode=safe_mode,
+                show_pbar=True,
             )
             outputs = output.outputs
             raw_outputs = output.raw_outputs

lotus/sem_ops/sem_join.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 import pandas as pd
+from tqdm import tqdm
 
 import lotus
 from lotus.templates import task_instructions
@@ -68,6 +69,11 @@ def sem_join(
         print("Sem_Join:")
         show_safe_mode(estimated_total_cost, estimated_total_calls)
 
+    pbar = tqdm(
+        total=len(l1),
+        desc="Processing uncached messages",
+        bar_format="{l_bar}{bar} {n}/{total} LM Calls [{elapsed}<{remaining}, {rate_fmt}{postfix}]",
+    )
     # for i1 in enumerate(l1):
     for id1, i1 in zip(ids1, left_multimodal_data):
         # perform llm filter
@@ -81,6 +87,7 @@ def sem_join(
             cot_reasoning=cot_reasoning,
             default=default,
             strategy=strategy,
+            show_pbar=False,
         )
         outputs = output.outputs
         raw_outputs = output.raw_outputs
@@ -98,6 +105,9 @@ def sem_join(
             ]
         )
 
+        pbar.update(1)
+    pbar.close()
+
     lotus.logger.debug(f"outputs: {filter_outputs}")
     lotus.logger.debug(f"explanations: {all_explanations}")