add cli arg to 🅱️oost 🅱️eams

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

aggregate.py

@@ -179,6 +179,15 @@ class BatchAggregator:
 
         self.aggregator.model.generation_config.update(**kwargs)
 
+    def get_generation_config(self) -> dict:
+        """
+        Get the current generation configuration.
+
+        Returns:
+            dict: The current generation configuration.
+        """
+        return self.aggregator.model.generation_config.to_dict()
+
     def update_loglevel(self, level: str = "INFO"):
         """
         Update the log level.

app.py

@@ -427,6 +427,14 @@ def parse_args():
         default=None,
         help=f"Add a token batch size to the demo UI options, default: {pp.pformat(TOKEN_BATCH_OPTIONS, compact=True)}",
     )
+    parser.add_argument(
+        "-max_agg",
+        "-2x",
+        "--aggregator_beam_boost",
+        dest="aggregator_beam_boost",
+        action="store_true",
+        help="Double the number of beams for the aggregator during beam search",
+    )
     parser.add_argument(
         "-level",
         "--log_level",
@@ -460,6 +468,12 @@ if __name__ == "__main__":
         logger.info(f"Adding token batch option {args.token_batch_option} to the list")
         TOKEN_BATCH_OPTIONS.append(args.token_batch_option)
 
+    if args.aggregator_beam_boost:
+        logger.info("Doubling aggregator num_beams")
+        _agg_cfg = aggregator.get_generation_config()
+        _agg_cfg["num_beams"] = _agg_cfg["num_beams"] * 2
+        aggregator.update_generation_config(**_agg_cfg)
+
     logger.info("Loading OCR model")
     with contextlib.redirect_stdout(None):
         ocr_model = ocr_predictor(