Addressing comments

Signed-off-by: Asmita Goswami <[email protected]>

Author: asmigosw

QEfficient/base/common.py

@@ -13,20 +13,19 @@
 """
 
 import importlib
-from collections import OrderedDict
 from typing import Any
 
 import transformers.models.auto.modeling_auto as mapping
 from transformers import AutoConfig
 
 from QEfficient.base.modeling_qeff import QEFFBaseModel
 
-MODEL_CLASS_MAPPING = OrderedDict(
-    [
-        (tuple(mapping.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values()), "QEFFAutoModelForCausalLM"),
-        (tuple(mapping.MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()), "QEFFAutoModelForImageTextToText"),
-    ]
-)
+MODEL_CLASS_MAPPING = {}
+for architecture in mapping.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
+    MODEL_CLASS_MAPPING[architecture] = "QEFFAutoModelForCausalLM"
+
+for architecture in mapping.MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values():
+    MODEL_CLASS_MAPPING[architecture] = "QEFFAutoModelForImageTextToText"
 
 
 class QEFFCommonLoader:
@@ -50,13 +49,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) ->
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
         architecture = config.architectures[0] if config.architectures else None
 
-        model_class = None
-        for key_tuple, class_name in MODEL_CLASS_MAPPING.items():
-            if architecture in key_tuple:
-                module = importlib.import_module("QEfficient.transformers.models.modeling_auto")
-                model_class = getattr(module, class_name)
-                break
-        if model_class is None:
+        class_name = MODEL_CLASS_MAPPING.get(architecture)
+        if class_name:
+            module = importlib.import_module("QEfficient.transformers.models.modeling_auto")
+            model_class = getattr(module, class_name)
+        else:
             raise NotImplementedError(
                 f"Unknown architecture={architecture}, either use specific auto model class for loading the model or raise an issue for support!"
             )

QEfficient/cloud/infer.py

@@ -12,7 +12,7 @@
 
 import requests
 from PIL import Image
-from transformers import AutoConfig, AutoProcessor, TextStreamer
+from transformers import AutoProcessor, TextStreamer
 from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
 
 from QEfficient.base.common import QEFFCommonLoader
@@ -121,16 +121,10 @@ def main(
         **kwargs,
     )
 
-    tokenizer = load_hf_tokenizer(
-        pretrained_model_name_or_path=(local_model_dir if local_model_dir else model_name),
-        cache_dir=cache_dir,
-        hf_token=hf_token,
-    )
-
     #########
     # Execute
     #########
-    config = AutoConfig.from_pretrained(model_name)
+    config = qeff_model.model.config
     architecture = config.architectures[0] if config.architectures else None
 
     if architecture in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values():
@@ -166,13 +160,19 @@ def main(
             add_special_tokens=False,
         )
         streamer = TextStreamer(processor.tokenizer)
-        _ = qeff_model.generate(
+        output = qeff_model.generate(
             inputs=split_inputs,
             streamer=streamer,
             device_ids=device_group,
             generation_len=generation_len,
         )
+        print(output)
     else:
+        tokenizer = load_hf_tokenizer(
+            pretrained_model_name_or_path=(local_model_dir if local_model_dir else model_name),
+            cache_dir=cache_dir,
+            hf_token=hf_token,
+        )
         _ = qeff_model.generate(
             tokenizer,
             prompts=prompt,

QEfficient/transformers/models/modeling_auto.py

@@ -835,8 +835,6 @@ def kv_offload_generate(
                 prefill_time=prefill_time, decode_perf=decode_perf, total_perf=total_perf, total_time=total_time
             ),
         )
-
-        print(exec_info)
         return exec_info
 
 
@@ -1128,8 +1126,6 @@ def cloud_ai_100_generate(
                 prefill_time=prefill_time, decode_perf=decode_perf, total_perf=total_perf, total_time=total_time
             ),
         )
-
-        print(exec_info)
         return exec_info
 
     @property
@@ -1572,7 +1568,7 @@ def compile(
             specializations.append(decode_specialization)
 
         if compiler_options.pop("img_size", None):
-            logger.warning("img_size is not a valid argument for Text-to-Text Model.")
+            logger.warning(f"Skipping img_size as it is not a valid argument for {self.model.config.architectures[0]}.")
 
         if enable_qnn:
             if compiler_options: