support_llava_hf_vlm_example (#381)

* support_llava_hf_vlm_example Signed-off-by: Zhang, Weiwei1 <[email protected]> * skip import check Signed-off-by: Zhang, Weiwei1 <[email protected]> --------- Signed-off-by: Zhang, Weiwei1 <[email protected]>
intel · Dec 11, 2024 · e6f89d1 · e6f89d1
1 parent 2990616
commit e6f89d1
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/auto_round/script/mllm.py b/auto_round/script/mllm.py
@@ -282,18 +282,21 @@ def tune(args):
 
     # load_model
     processor, image_processor = None, None
-    if "llava" in model_name:
-        from llava.model.builder import load_pretrained_model  # pylint: disable=E0401
+    config = AutoConfig.from_pretrained(model_name, trust_remote_code=not args.disable_trust_remote_code)
+    if "llava" in model_name and config.architectures[0] != "LlavaForConditionalGeneration":
+        from llava.model.builder import load_pretrained_model   # pylint: disable=E0401
         tokenizer, model, image_processor, _ = load_pretrained_model(
             model_name, model_base=None, model_name=model_name,
             torch_dtype=torch_dtype)
         model_type = "llava"
     else:
-        config = AutoConfig.from_pretrained(model_name, trust_remote_code=not args.disable_trust_remote_code)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=not args.disable_trust_remote_code)
         processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=not args.disable_trust_remote_code)
         model_type = config.model_type
-        if "qwen2_vl" in model_type:
+        if "llava" in model_type:
+            from transformers import LlavaForConditionalGeneration
+            cls = LlavaForConditionalGeneration
+        elif "qwen2_vl" in model_type:
             from transformers import Qwen2VLForConditionalGeneration
             cls = Qwen2VLForConditionalGeneration
         elif "mllama" in model_type:
@@ -511,3 +514,4 @@ def lmms_eval(args):
         apply_chat_template=False,
     )
     return results
+