Simplify

Author: Jack-Khuu

torchchat/generate.py

@@ -797,6 +797,10 @@ def _gen_model_input(
             max_new_tokens is not None
         ), "max_new_tokens must be specified for Flamingo models"
 
+        # Wrap string prompts into a list
+        if isinstance(prompt, str):
+            prompt = [{"role": "user", "content": prompt}]
+
         image_found = False
         messages = []
         for message in prompt:
@@ -960,15 +964,8 @@ def chat(
             text_transformer_args.max_seq_length if text_transformer_args else 2048
         )
 
-        # For Llama 3.2 11B (Flamingo), format the input as a dialog
-        # Else for other models, format the input as a single string
-        text_prompt = (
-            generator_args.prompt
-            if self.model.config.model_type != ModelType.Flamingo
-            else [{"role": "user", "content": generator_args.prompt}]
-        )
         encoded, batch = self._gen_model_input(
-            text_prompt,
+            generator_args.prompt,
             generator_args.image_prompts,
             generator_args.max_new_tokens,
             max_seq_length,