Added Model mapping for swiftKV model to __init__.py (#413)

`__init__.py` was calling modelling_utils leading to creation of
`MODEL_CLASS_MAPPING` before init. As a result, registered swiftKV
model's config class and QEff Class was not present in the
`MODEL_CLASS_MAPPING` and hence giving error of unknown architecture.

Updated this `__init__` to be initialized first and then all other
classes by removing the modelling_utils call and adding
model_class_mapping in __init__.

---------

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

Author: asmigosw

QEfficient/__init__.py

@@ -12,19 +12,10 @@
 # hf_transfer is imported (will happen on line 15 via leading imports)
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
-from transformers import AutoConfig
-
-from QEfficient.transformers.modeling_utils import MODEL_TYPE_TO_CONFIG_CLS_AND_ARCH_CLS
+# Placeholder for all non-transformer models registered in QEfficient
+import QEfficient.utils.model_registery  # noqa: F401
 from QEfficient.utils.logging_utils import logger
 
-# loop over all the model types which are not present in transformers and register them
-for model_type, model_cls in MODEL_TYPE_TO_CONFIG_CLS_AND_ARCH_CLS.items():
-    # Register the model config class based on the model type. This will be first element in the tuple
-    AutoConfig.register(model_type, model_cls[0])
-
-    # Register the non transformer library Class and config class using AutoModelClass
-    model_cls[2].register(model_cls[0], model_cls[1])
-
 
 def check_qaic_sdk():
     """Check if QAIC SDK is installed"""

QEfficient/base/common.py

@@ -41,15 +41,14 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) ->
         Downloads HuggingFace model if already doesn't exist locally, returns QEFFAutoModel object based on type of model.
         """
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
-        architecture = config.architectures[0] if config.architectures else None
 
-        class_name = MODEL_CLASS_MAPPING.get(architecture)
+        class_name = MODEL_CLASS_MAPPING.get(config.__class__.__name__, None)
         if class_name:
             module = __import__("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!"
+                f"Unknown architecture={config.__class__.__name__}, either use specific auto model class for loading the model or raise an issue for support!"
             )
 
         local_model_dir = kwargs.pop("local_model_dir", None)

QEfficient/transformers/modeling_utils.py

@@ -11,7 +11,6 @@
 import torch
 import torch.nn as nn
 import transformers.models.auto.modeling_auto as mapping
-from transformers import AutoModelForCausalLM
 from transformers.models.codegen.modeling_codegen import (
     CodeGenAttention,
     CodeGenBlock,
@@ -91,11 +90,6 @@
 from QEfficient.customop import CustomRMSNormAIC
 
 # Placeholder for all non-transformer models
-from QEfficient.transformers.models.llama_swiftkv.modeling_llama_swiftkv import (
-    QEffLlamaSwiftKVConfig,
-    QEffLlamaSwiftKVForCausalLM,
-)
-
 from .models.codegen.modeling_codegen import (
     QEffCodeGenAttention,
     QeffCodeGenBlock,
@@ -279,18 +273,19 @@
     WhisperForConditionalGeneration: QEffWhisperForConditionalGeneration,
 }
 
-# Map of model type to config class, Modelling class and transformer model architecture class
-MODEL_TYPE_TO_CONFIG_CLS_AND_ARCH_CLS = {
-    "llama_swiftkv": [QEffLlamaSwiftKVConfig, QEffLlamaSwiftKVForCausalLM, AutoModelForCausalLM],
-}
+
+def build_model_class_mapping(auto_model_class, qeff_class_name):
+    """
+    Build a mapping of model config class names to QEfficient model class names.
+    """
+    return {
+        config_class.__name__: qeff_class_name for config_class, model_class in auto_model_class._model_mapping.items()
+    }
 
 
 MODEL_CLASS_MAPPING = {
-    **{architecture: "QEFFAutoModelForCausalLM" for architecture in mapping.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values()},
-    **{
-        architecture: "QEFFAutoModelForImageTextToText"
-        for architecture in mapping.MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()
-    },
+    **build_model_class_mapping(mapping.AutoModelForCausalLM, "QEFFAutoModelForCausalLM"),
+    **build_model_class_mapping(mapping.AutoModelForImageTextToText, "QEFFAutoModelForImageTextToText"),
 }
 
 

QEfficient/utils/model_registery.py

@@ -0,0 +1,28 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# -----------------------------------------------------------------------------
+
+
+from transformers import AutoConfig, AutoModelForCausalLM
+
+# Placeholder for all non-transformer models
+from QEfficient.transformers.models.llama_swiftkv.modeling_llama_swiftkv import (
+    QEffLlamaSwiftKVConfig,
+    QEffLlamaSwiftKVForCausalLM,
+)
+
+# Map of model type to config class, Modelling class and transformer model architecture class
+MODEL_TYPE_TO_CONFIG_CLS_AND_ARCH_CLS = {
+    "llama_swiftkv": [QEffLlamaSwiftKVConfig, QEffLlamaSwiftKVForCausalLM, AutoModelForCausalLM],
+}
+
+# loop over all the model types which are not present in transformers and register them
+for model_type, model_cls in MODEL_TYPE_TO_CONFIG_CLS_AND_ARCH_CLS.items():
+    # Register the model config class based on the model type. This will be first element in the tuple
+    AutoConfig.register(model_type, model_cls[0])
+
+    # Register the non transformer library Class and config class using AutoModelClass
+    model_cls[2].register(model_cls[0], model_cls[1])