diff --git a/surya/foundation/__init__.py b/surya/foundation/__init__.py index a04ec2b..ada34b6 100644 --- a/surya/foundation/__init__.py +++ b/surya/foundation/__init__.py @@ -723,6 +723,7 @@ class FoundationPredictor(BasePredictor): drop_repeated_tokens: bool = True, max_lookahead_tokens: Optional[int] = None, top_k: int = 0, + tqdm_desc: str = "Recognizing Text" ) -> tuple: allowed_tasks = self.tasks.keys() assert all([task_name in allowed_tasks for task_name in task_names]), ( @@ -765,7 +766,7 @@ class FoundationPredictor(BasePredictor): pbar = tqdm( total=len(self.prompt_queue), - desc="Recognizing Text", + desc=tqdm_desc, disable=self.disable_tqdm, ) diff --git a/surya/layout/__init__.py b/surya/layout/__init__.py index b8d6005..bd710a6 100644 --- a/surya/layout/__init__.py +++ b/surya/layout/__init__.py @@ -53,7 +53,8 @@ class LayoutPredictor(BasePredictor): max_lookahead_tokens=0, # Do not do MTP for layout top_k=5, max_sliding_window=576, - max_tokens=500 + max_tokens=500, + tqdm_desc="Recognizing Layout" ) ) diff --git a/surya/recognition/__init__.py b/surya/recognition/__init__.py index a0b2947..f812c08 100644 --- a/surya/recognition/__init__.py +++ b/surya/recognition/__init__.py @@ -438,6 +438,7 @@ class RecognitionPredictor(BasePredictor): max_lookahead_tokens=self.foundation_predictor.model.config.multi_output_distance, max_sliding_window=max_sliding_window, max_tokens=max_tokens, + tqdm_desc="Recognizing Text" ) # Get text and bboxes in structured form