lint

2024-12-06 14:37:32 -05:00
parent a3a4d22709
commit eab1638686
1 changed files with 5 additions and 0 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -30,6 +30,7 @@ from transformers import (  # noqa: F401
    AddedToken,
    AutoConfig,
    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
    AutoModelForVision2Seq,
    AutoProcessor,
    AutoTokenizer,
@@ -553,6 +554,10 @@ class ModelLoader:
                self.AutoModelLoader = (  # pylint: disable=invalid-name
                    MllamaForConditionalGeneration
                )
+            elif self.model_config.model_type == "qwen2_vl":
+                self.AutoModelLoader = (  # pylint: disable=invalid-name
+                    AutoModelForImageTextToText
+                )
            else:
                self.AutoModelLoader = (
                    AutoModelForVision2Seq  # pylint: disable=invalid-name