🎨 📝

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

aggregate.py

@@ -1,3 +1,12 @@
+"""
+aggregate.py is a module for aggregating text from multiple sources, or multiple parts of a single source.
+    Primary usage is through the BatchAggregator class.
+
+How it works:
+1. We tell the language model to do it.
+2. The language model does it.
+3. Yaay!
+"""
 import pprint as pp
 import logging
 import time
@@ -14,10 +23,15 @@ logging.basicConfig(
 
 
 class BatchAggregator:
-    CONFIGURED_MODELS = [
-        "pszemraj/bart-large-mnli-dolly_hhrlhf-v1"
-    ]  # TODO: Add models here
-    DEFAULT_INSTRUCTION = "Write a comprehensive yet concise summary that pulls together the main points of the following text:"
+    """
+    BatchAggregator is a class for aggregating text from multiple sources.
+
+    Usage:
+    >>> from aggregate import BatchAggregator
+    >>> aggregator = BatchAggregator()
+    >>> aggregator.aggregate(["This is a test", "This is another test"])
+    """
+
     GENERIC_CONFIG = GenerationConfig(
         num_beams=8,
         early_stopping=True,
@@ -29,10 +43,23 @@ class BatchAggregator:
         no_repeat_ngram_size=4,
         encoder_no_repeat_ngram_size=5,
     )
+    CONFIGURED_MODELS = [
+        "pszemraj/bart-large-mnli-dolly_hhrlhf-v1",
+        "pszemraj/bart-base-instruct-dolly_hhrlhf",
+        "pszemraj/flan-t5-large-instruct-dolly_hhrlhf",
+        "pszemraj/flan-t5-base-instruct-dolly_hhrlhf",
+    ]  # these have generation configs defined for this task in their model repos
+
+    DEFAULT_INSTRUCTION = "Write a comprehensive yet concise summary that pulls together the main points of the following text:"
 
     def __init__(
         self, model_name: str = "pszemraj/bart-large-mnli-dolly_hhrlhf-v1", **kwargs
     ):
+        """
+        __init__ initializes the BatchAggregator class.
+
+        :param str model_name: model name to use, default: "pszemraj/bart-large-mnli-dolly_hhrlhf-v1"
+        """
         self.device = None
         self.is_compiled = False
         self.logger = logging.getLogger(__name__)
@@ -125,20 +152,6 @@ class BatchAggregator:
         """
         self.aggregator.model.generation_config = self.GENERIC_CONFIG
 
-        if "bart" in self.model_name.lower():
-            self.logger.info("Using BART model, updating generation config")
-            upd = {
-                "num_beams": 8,
-                "repetition_penalty": 1.3,
-                "length_penalty": 1.0,
-                "_from_model_config": False,
-                "max_new_tokens": 256,
-                "min_new_tokens": 32,
-                "no_repeat_ngram_size": 3,
-                "encoder_no_repeat_ngram_size": 6,
-            }  # TODO: clean up
-            self.aggregator.model.generation_config.update(**upd)
-
         if (
             "large"
             or "xl" in self.model_name.lower()