tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

lint

Browse Source
This commit is contained in:
bursteratom
2024-12-06 14:37:32 -05:00
parent a3a4d22709
commit eab1638686
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/axolotl/utils/models.py
View File

@@ -30,6 +30,7 @@ from transformers import ( # noqa: F401
AddedToken, AddedToken,
AutoConfig, AutoConfig,
AutoModelForCausalLM, AutoModelForCausalLM,
AutoModelForImageTextToText,
AutoModelForVision2Seq, AutoModelForVision2Seq,
AutoProcessor, AutoProcessor,
AutoTokenizer, AutoTokenizer,
@@ -553,6 +554,10 @@ class ModelLoader:
self.AutoModelLoader = ( # pylint: disable=invalid-name self.AutoModelLoader = ( # pylint: disable=invalid-name
MllamaForConditionalGeneration MllamaForConditionalGeneration
) )
elif self.model_config.model_type == "qwen2_vl":
self.AutoModelLoader = ( # pylint: disable=invalid-name
AutoModelForImageTextToText
)
else: else:
self.AutoModelLoader = ( self.AutoModelLoader = (
AutoModelForVision2Seq # pylint: disable=invalid-name AutoModelForVision2Seq # pylint: disable=invalid-name