✨ add ability to force CPU

Signed-off-by: peter szemraj <[email protected]>

aggregate.py

@@ -54,15 +54,22 @@ class BatchAggregator:
     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
+        self,
+        model_name: str = "pszemraj/bart-large-mnli-dolly_hhrlhf-v1",
+        force_cpu: bool = False,
+        **kwargs,
     ):
         """
         __init__ initializes the BatchAggregator class.
 
         :param str model_name: model name to use, default: "pszemraj/bart-large-mnli-dolly_hhrlhf-v1"
+        :param bool force_cpu: force the model to run on CPU, default: False
         """
         self.device = None
         self.is_compiled = False
+        self.model_name = None
+        self.aggregator = None
+        self.force_cpu = force_cpu
         self.logger = logging.getLogger(__name__)
         self.init_model(model_name)
 
@@ -105,7 +112,7 @@ class BatchAggregator:
 
         :raises Exception: if the pipeline cannot be created
         """
-        self.device = 0 if torch.cuda.is_available() else -1
+        self.device = 0 if torch.cuda.is_available() and not self.force_cpu else -1
         try:
             self.logger.info(
                 f"Creating pipeline with model {model_name} on device {self.device}"