[WIP] Generate base class for better integration of distributed inference

torchchat/distributed/dist_run.py

@@ -282,16 +282,6 @@ def _cleanup():
     dist.destroy_process_group()
 
 
-prompts = [
-    "What is Snow?",
-    # "Can you explain what is the purpose of back propagation in neural networks?",
-    "Who is Santa Claus?",
-    "Where does Santa live?",
-    "Who is Abraham Lincoln?",
-    # "How are models trained?",
-]
-
-
 def main(
     model_name,
     builder_args,

torchchat/distributed/generate.py

@@ -21,6 +21,7 @@
 from torchchat.cli.builder import BuilderArgs, TokenizerArgs
 from torchchat.distributed.dist_run import NAME_TO_DISTRIBUTION_AND_DTYPE
 from torchchat.distributed.logging_utils import SingletonLogger
+from torchchat.utils.generator import Generator, GeneratorArgs
 
 logger = SingletonLogger.get_logger()
 
@@ -194,19 +195,19 @@ def step(self) -> List[Output]:
         return outputs
 
 
-class DistributedGenerator(object):
+class DistributedGenerator(Generator):
     def __init__(
         self,
         # TODO: switch this to torchchat method
         model_name: str,
         builder_args: BuilderArgs,
         tokenizer_args: TokenizerArgs,
-        # TODO: move GeneratorArgs into a different module
-        generator_args,
+        generator_args: GeneratorArgs,
         profile: Optional[Path],
         quantize: bool,
         draft_quantize: bool,
     ):
+        super().__init__(builder_args, tokenizer_args, generator_args)
         self.model_name = model_name
         self.builder_args = builder_args
         self.generate_args = generator_args