rebased

Signed-off-by: Onkar Chougule <[email protected]>

Author: ochougul

QEfficient/transformers/modeling_utils.py

@@ -378,6 +378,5 @@ def _create_causal_mask(
 # While onboarding new models make sure to add the new SwiftKV model card names to this dictionary.
 SwiftKVModelTypeToConfigClassAndModelArchClassDict = {
     # LlamaSwiftKV Model
-    "llama_swiftkv" : [LlamaSwiftKVConfig, LlamaSwiftKVForCausalLM]
+    "llama_swiftkv": [LlamaSwiftKVConfig, LlamaSwiftKVForCausalLM]
 }
-

QEfficient/transformers/models/llama_swiftkv/config_llama_swiftkv.py

@@ -9,8 +9,6 @@
 
 """Inference-only LLaMA model compatible with HuggingFace weights."""
 
-
-
 from typing import Optional
 from transformers import LlamaConfig
 
@@ -40,6 +38,4 @@ def __init__(
         self.swiftkv = swiftkv
         self.num_key_value_layers = num_key_value_layers or self.num_hidden_layers
         self.key_value_group_size = key_value_group_size or 1
-        assert (
-            self.num_hidden_layers - self.num_key_value_layers
-        ) % self.key_value_group_size == 0
+        assert (self.num_hidden_layers - self.num_key_value_layers) % self.key_value_group_size == 0

QEfficient/transformers/models/llama_swiftkv/modeling_llama_swiftkv.py

@@ -29,6 +29,7 @@
 )
 from QEfficient.transformers.models.llama_swiftkv.config_llama_swiftkv import LlamaSwiftKVConfig
 
+
 class LlamaSwiftKVAttention(nn.Module):
     def __init__(self, config: LlamaSwiftKVConfig, layer_idx) -> None:
         super().__init__()

QEfficient/transformers/models/modeling_auto.py

@@ -80,7 +80,6 @@ def __repr__(self) -> str:
     @classmethod
     @with_replaced_quantizers
     def from_pretrained(cls, pretrained_model_name_or_path: str, is_tlm: bool = False, *args, **kwargs):
-
         # Load the SwiftKV model if supported
         QEFFLoadSwiftKVModels(pretrained_model_name_or_path)
 

QEfficient/utils/_utils.py

@@ -8,18 +8,30 @@
 import json
 import os
 import subprocess
+import sys
+import warnings
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import requests
 import torch
 from huggingface_hub import login, snapshot_download
 from requests.exceptions import HTTPError
-from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
-
+from transformers import (
+    AutoConfig,
+    AutoProcessor,
+    AutoTokenizer,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+
+from QEfficient.transformers.modeling_utils import (
+    SwiftKVModelCardNameToSwiftKVModelTypeDict,
+    SwiftKVModelTypeToConfigClassAndModelArchClassDict,
+)
 from QEfficient.utils.constants import QEFF_MODELS_DIR, Constants
 from QEfficient.utils.logging_utils import logger
-from QEfficient.transformers.modeling_utils import SwiftKVModelCardNameToSwiftKVModelTypeDict, SwiftKVModelTypeToConfigClassAndModelArchClassDict
+
 
 class DownloadRetryLimitExceeded(Exception):
     """
@@ -442,3 +454,63 @@ class IOInfo:
 
     def __repr__(self):
         return f"input_name:{self.name}\tdatatype:{self.datatype}\tshape:{self.shape}"
+
+
+def convert_str_to_class(className):
+    """
+    Convert the string to class name
+    ---------
+    :className: `str`- Class name string.
+    Return:
+        Class Name
+    """
+    return getattr(sys.modules[__name__], className)
+
+
+def register_swiftKV_model(model_type, SwiftkvConfigCls, SwiftKVModelCls):
+    """
+    Register the SwiftKV Models
+    ---------------------------------------
+    : model_type: str: name of the swiftKVModel for example llama_swiftkv
+    : SwiftkVConfigCls: SwiftKV Config class for example LlamaSwiftKVConfig
+    : SwiftKVModelCls: SwiftKV model class name for example LlamaSwiftKVForCausalLM
+    """
+
+    # Register the SwiftKV Config class using AutoConfig
+    AutoConfig.register(model_type, SwiftkvConfigCls)
+
+    # Construct the AutoModel class name using SwiftKVModel Class name, this code is written to make things generic
+    swiftKvModelName = SwiftKVModelCls.__name__
+    start_index = swiftKvModelName.find("SwiftKVFor")
+
+    # Calculate the index after "SwiftKVFor"
+    substring_start = start_index + len("SwiftKVFor")
+
+    # Get the substring after "SwiftKVFor"
+    swiftKVModel = swiftKvModelName[substring_start:]
+
+    AutoModelName = "AutoModelFor" + swiftKVModel
+
+    # Convert the string to class name
+    AutoModelClassName = convert_str_to_class(AutoModelName)
+
+    # Register the SwiftKVModel Class and config class using AutoModelClass
+    AutoModelClassName.register(SwiftkvConfigCls, SwiftKVModelCls)
+
+
+def QEFFLoadSwiftKVModels(pretrained_model_name_or_path):
+    """
+    Load the SwiftKV Models
+    ---------------------------------------
+    : pretrained_model_name_or_path: str: name of the swiftKVModel for example Snowflake/Llama-3.1-SwiftKV-8B-Instruct
+    """
+    try:
+        modelType = SwiftKVModelCardNameToSwiftKVModelTypeDict[pretrained_model_name_or_path]
+
+        SwiftKVConfigCls = SwiftKVModelTypeToConfigClassAndModelArchClassDict[modelType][0]
+        SwiftKVModelArchCls = SwiftKVModelTypeToConfigClassAndModelArchClassDict[modelType][1]
+
+        register_swiftKV_model(modelType, SwiftKVConfigCls, SwiftKVModelArchCls)
+
+    except KeyError:
+        warnings.warn("Requested SwiftKVModel is currently not supported... stay tuned for future releases", Warning)