optimum-rbln 0.9.3__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/transformers/configuration_generic.py

@@ -118,3 +118,30 @@ class RBLNModelForImageClassificationConfig(RBLNImageModelConfig):
 
 class RBLNModelForDepthEstimationConfig(RBLNImageModelConfig):
     pass
+
+
+class RBLNModelForAudioClassificationConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        max_length: Optional[int] = None,
+        num_mel_bins: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            max_length (Optional[int]): Maximum length of the audio input in time dimension.
+            num_mel_bins (Optional[int]): Number of Mel frequency bins for audio processing.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.max_length = max_length
+        self.num_mel_bins = num_mel_bins

optimum/rbln/transformers/modeling_attention_utils.py

@@ -1,18 +1,18 @@
 import math
-from collections import Counter, defaultdict
-from typing import TYPE_CHECKING, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
-import rebel
+from optimum.rbln.transformers.models.decoderonly.configuration_decoderonly import (
+    RBLNDecoderOnlyModelForCausalLMConfig,
+)
 
 from ..utils.logging import get_logger
-from ..utils.runtime_utils import get_available_dram
-from .models.decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
 
 
 logger = get_logger()
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedModel
+    from rebel import RBLNCompiledModel
+    from transformers import PretrainedConfig
 
 
 DEFAULT_FLASH_ATTN_PARTITION_LENGTH = 16_384
@@ -115,261 +115,128 @@ def validate_sliding_window(rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         raise ValueError("`use_attention_mask` must be set to False when `cache_impl` is set to 'sliding_window'.")
 
 
-def align(x: int, nbytes: int) -> int:
-    return int(math.ceil(x / nbytes) * nbytes)
-
-
-def align_2MB(x: int) -> int:
-    return align(x, 2**21)
-
-
-def get_alloc_memory_by_key(compiled_models: Dict[str, "rebel.RBLNCompiledModel"]) -> Dict[str, int]:
-    alloc_memory_by_key = defaultdict(int)
-    # Get the actual memory allocation of each node by key
-    for compiled_model in compiled_models.values():
-        alloc_per_node_by_key = compiled_model.get_alloc_per_node_by_key()
-        for key, memory_per_node in alloc_per_node_by_key.items():
-            alloc_memory_by_key[key] += sum(memory_per_node)
-
-    return alloc_memory_by_key
-
-
-def format_byte_size(nbytes: int) -> str:
-    if nbytes < 1024:
-        return f"{nbytes} B"
-    elif nbytes < 1024**2:
-        return f"{nbytes / 1024:.2f} KB"
-    elif nbytes < 1024**3:
-        return f"{nbytes / 1024**2:.2f} MB"
-    else:
-        return f"{nbytes / 1024**3:.2f} GB"
-
-
 class RBLNDecoderOnlyFlashAttentionMixin:
     @classmethod
-    def get_maximum_num_blocks_by_model(
+    def get_maximum_num_blocks(
         cls,
-        model: "PreTrainedModel",
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
+        config: "PretrainedConfig",
+        tensor_parallel_size: int,
+        kvcache_block_size: int,
+        nbits_per_param: Optional[int] = None,
+        n_model_params: Optional[int] = None,
+        kernel_size: Optional[int] = None,
+        buffer: Optional[int] = None,
+        num_runtimes: int = 2,
     ) -> int:
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        available_dram = get_available_dram(rbln_config.npu) * tensor_parallel_size
-
-        kernel_memory = cls._get_kernel_memory(model, model_config=model_config, rbln_config=rbln_config)
-        buffer = cls._get_buffer(rbln_config)
-
-        remaining_dram = available_dram - kernel_memory - buffer
-        if remaining_dram <= 0:
-            raise ValueError(
-                "Insufficient available DRAM after accounting for kernel memory and buffer. "
-                "Cannot allocate any KV cache blocks."
-                f" (Available DRAM: {format_byte_size(available_dram)}, "
-                f"Kernel Memory: {format_byte_size(kernel_memory)}, "
-                f"Buffer: {format_byte_size(buffer)})"
+        # We are finding max_n_blocks(x) that satisfies the following equation:
+
+        # available_dram - kernel_size - buffer
+        #     - num_layers * 2 * tensor_parallel_size
+        #     * align_2MB(
+        #         x
+        #         * block_size
+        #         * align_64(head_dim)
+        #         * math.ceil(num_key_value_heads / tensor_parallel_size)
+        #         * 2
+        #     ) > 0
+
+        # This inequality can be rewritten as follows:
+
+        # a - c * align_2MB(b * x) > 0
+        # where
+        #    a = available_dram - kernel_size - buffer
+        #    b = block_size * align_64(head_dim) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+        #    c = num_layers * 2 * tensor_parallel_size
+
+        # We can rewrite the inequality as follows:
+        # k > align_2MB(b*x)
+        # where
+        #    k = a / c
+
+        # After that, we can derive the following equation:
+        # x = floor(2**21 / b * floor((k - 1) / 2**21))
+
+        def align(x: int, nbytes: int) -> int:
+            return int(math.ceil(x / nbytes) * nbytes)
+
+        def align_2MB(x: int) -> int:
+            return align(x, 2**21)
+
+        num_attention_heads = getattr(config, "n_head", None) or getattr(config, "num_attention_heads")
+        num_layers = getattr(config, "n_layer", None) or getattr(config, "num_hidden_layers")
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // num_attention_heads
+        vocab_size = config.vocab_size
+        hidden_size = getattr(config, "n_embd", None) or getattr(config, "hidden_size")
+        num_key_value_heads = getattr(config, "num_key_value_heads", None) or num_attention_heads
+
+        # TODO(jongho): Update if target npu is REBEL.
+        ATOM_DRAM_NBYTES = 16 * 2**30
+        ATOM_SYS_DRAM_NBYTES = 288 * 2**20
+        available_dram = tensor_parallel_size * (ATOM_DRAM_NBYTES - ATOM_SYS_DRAM_NBYTES)
+
+        if kernel_size is None:
+            if n_model_params is None:
+                raise ValueError("`n_model_params` should be specified to estimate the kernel memory.")
+            # Get estimated kernel size (approximated)
+            lm_heads_params = align(vocab_size, 64) * hidden_size
+            lm_heads_nbytes = (
+                align_2MB(lm_heads_params * nbits_per_param // 8 / tensor_parallel_size) * tensor_parallel_size
             )
-        estimated_num_blocks = cls._estimate_num_blocks(
-            remaining_dram, model_config=model_config, rbln_config=rbln_config
-        )
-
-        return estimated_num_blocks
-
-    @classmethod
-    def _get_kernel_memory(
-        cls,
-        model: "PreTrainedModel",
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> int:
-        if model.get_output_embeddings() is None:
-            lm_head_nbytes = 0
-        else:
-            lm_head_nbytes = cls._get_lm_head_memory(model_config, rbln_config)
-
-        layer_nbytes = cls._get_layer_memory(model, model_config, rbln_config)
-        return lm_head_nbytes + layer_nbytes
-
-    @classmethod
-    def _get_lm_head_memory(
-        cls, model_config: "PretrainedConfig", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ) -> int:
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        vocab_size = model_config.vocab_size
-        hidden_size = getattr(model_config, "n_embd", None) or getattr(model_config, "hidden_size")
-        lm_head_params = align(vocab_size, 64) * hidden_size
-
-        nbytes_per_param = 2  # Assuming lm_head is always not quantized
-        lm_head_memory_in_bytes = (
-            align_2MB(lm_head_params * nbytes_per_param / tensor_parallel_size) * tensor_parallel_size
-        )
-
-        return lm_head_memory_in_bytes
-
-    @classmethod
-    def _get_layer_memory(
-        cls,
-        model: "PreTrainedModel",
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> int:
-        # This is an *APPROXIMATE* calculation based on the number of parameters
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        num_hidden_layers = getattr(model_config, "n_layer", None) or getattr(model_config, "num_hidden_layers")
-
-        n_model_params = sum(p.numel() for p in model.parameters())
-        embed_token_params = sum(p.numel() for p in model.get_input_embeddings().parameters())
-
-        # Check : `embed_token` is same as `lm_head`
-        if model.get_output_embeddings() is not None:
-            params = n_model_params - 2 * embed_token_params
-        else:
-            params = n_model_params - embed_token_params
-
-        # Assuming all layers have the same number of parameters
-        # and all linear layers are quantized if quantization is enabled (This is not always true)
-        # TODO(jongho): More accurate calculation
-        nbits_per_param = rbln_config.nbits_per_param
-        layer_nbytes = (
-            (align_2MB(params // num_hidden_layers * nbits_per_param // 8 / tensor_parallel_size))
-            * num_hidden_layers
-            * tensor_parallel_size
-        )
-
-        return layer_nbytes
-
-    @classmethod
-    def _get_buffer(cls, rbln_config) -> int:
-        # TODO(jongho): Accurate buffer estimation
-        buffer_per_runtime_per_core = 2**28  # 256MB per runtime
-        num_runtimes = 1 if not rbln_config.can_generate else 1 + len(rbln_config.decoder_batch_sizes)
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-
-        buffer_per_core = buffer_per_runtime_per_core * num_runtimes
-        buffer = buffer_per_core * tensor_parallel_size
-        return buffer
-
-    @classmethod
-    def get_maximum_num_blocks_by_compiled_model(
-        cls,
-        compiled_models: Dict[str, "rebel.RBLNCompiledModel"],
-        model_config: "PretrainedConfig",
-        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
-    ) -> int:
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-        available_dram = get_available_dram(rbln_config.npu) * tensor_parallel_size
-
-        alloc_memory_by_key = get_alloc_memory_by_key(compiled_models)
-        alloc_memory_by_key.pop("PortRecur", None)  # Old compiler's kv-cache Key
-        alloc_memory_by_key.pop("DramTensor", None)  # kv-cache
-        used_memory = sum(alloc_memory_by_key.values())
-
-        remaining_dram = available_dram - used_memory
-
-        if remaining_dram <= 0:
-            logger.warning(
-                "Insufficient available DRAM after accounting for kernel memory and buffer. "
-                "Model cannot allocate any KV cache blocks."
-            )
-
-        estimated_num_blocks = cls._estimate_num_blocks(
-            remaining_dram, model_config=model_config, rbln_config=rbln_config
-        )
-
-        return estimated_num_blocks
-
-    @classmethod
-    def _estimate_num_blocks(
-        cls, available_dram: int, model_config: "PretrainedConfig", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
-    ) -> int:
-        """
-        Estimate the maximum number of KV cache blocks that can be allocated.
-
-        if all of the layers are full attention, the dram_per_block can be calculated simply as follows:
-            num_blocks = available_dram // dram_per_block
-
-        However, if the model contains a mix of full attention and sliding window attention layers,
-        we need to consider the memory occupied by the sliding window attention layers first,
-        since their memory usage is constant regardless of the number of blocks.
-            num_blocks = (available_dram - swa_kv_nbytes) // dram_per_block
-
-        """
-
-        def get_dram_per_block(seq_len: int, num_key_value_heads: int, tensor_parallel_size: int) -> int:
-            nbytes_per_param = 2  # Assuming kv-cache is always not quantized
-            dram_per_block = (
-                seq_len
-                * align(head_dim, 64)
-                * math.ceil(num_key_value_heads / tensor_parallel_size)
-                * nbytes_per_param
+            params = n_model_params - lm_heads_params
+            layer_nbytes = (
+                align_2MB(params * nbits_per_param // 8 / num_layers / tensor_parallel_size)
+                * num_layers
                 * tensor_parallel_size
-                * 2
-            )  # *2 for key and value
-
-            return dram_per_block
-
-        num_attention_heads = getattr(model_config, "n_head", None) or getattr(model_config, "num_attention_heads")
-        head_dim = getattr(model_config, "head_dim", None) or model_config.hidden_size // num_attention_heads
-        num_hidden_layers = getattr(model_config, "n_layer", None) or getattr(model_config, "num_hidden_layers")
-        num_key_value_heads = getattr(model_config, "num_key_value_heads", None) or num_attention_heads
-        tensor_parallel_size = rbln_config.tensor_parallel_size or 1
-
-        # Consider layer types if available
-        # If layer types are not found, assume all layers are full attention
-        layer_types = getattr(model_config, "layer_types", None)
-        if layer_types:
-            layer_types_dict = Counter(layer_types)
-            num_full_attention = layer_types_dict.pop("full_attention", 0)
-            num_sliding_window_attention = layer_types_dict.pop("sliding_attention", 0)
-            if len(layer_types_dict) > 0:
-                raise ValueError(f"Unknown layer types found in the config: {layer_types_dict.keys()}")
-
-        else:
-            num_full_attention = num_hidden_layers
-            num_sliding_window_attention = 0
-
-        # Reduce available DRAM by sliding window attention kv-cache
-        # Since memory occupation of swa layer is constant regardless of num_blocks
-        swa_kv_nbytes = 0
-        if num_sliding_window_attention > 0:
-            sliding_window = getattr(model_config, "sliding_window", None)
-            if sliding_window is None:
-                logger.warning(
-                    "`sliding_window` is not found in the config while `sliding_attention` layers are present. "
-                    "Assuming maximum sliding window size for estimation."
-                )
-                sliding_window = rbln_config.kvcache_block_size
-
-            swa_kv_nbytes = num_sliding_window_attention * get_dram_per_block(
-                seq_len=sliding_window,
-                num_key_value_heads=num_key_value_heads,
-                tensor_parallel_size=tensor_parallel_size,
             )
+            kernel_size = layer_nbytes + lm_heads_nbytes
+        elif n_model_params is not None:
+            raise ValueError("Both `n_model_params` and `kernel_size` cannot be specified.")
 
-            available_dram -= swa_kv_nbytes
+        available_dram -= kernel_size
 
-        dram_per_block = num_full_attention * get_dram_per_block(
-            seq_len=rbln_config.kvcache_block_size,
-            num_key_value_heads=num_key_value_heads,
-            tensor_parallel_size=tensor_parallel_size,
-        )
+        if buffer is None:
+            # TODO: Accurate buffer estimation
+            buffer_per_runtime_per_core = 2**28  # 256MB per runtime
+            buffer_per_core = buffer_per_runtime_per_core * num_runtimes  # 1 for prefill, 1 for decoder
+            buffer = buffer_per_core * tensor_parallel_size
+        available_dram -= buffer
 
-        if dram_per_block == 0:
-            raise ValueError("DRAM per block is calculated as zero, cannot estimate maximum number of blocks.")
+        b = kvcache_block_size * align(head_dim, 64) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+        c = num_layers * 2 * tensor_parallel_size
+        k = available_dram / c
+        max_n_blocks = math.floor(2**21 / b * math.floor((k - 1) / 2**21))
 
-        max_n_blocks = available_dram // dram_per_block
         return max_n_blocks
 
     @classmethod
     def maybe_suggest_kvcache_num_blocks(
         cls,
-        compiled_models: Dict[str, "rebel.RBLNCompiledModel"],
+        compiled_models: Dict[str, "RBLNCompiledModel"],
         model_config: "PretrainedConfig",
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
     ) -> None:
-        max_num_blocks = cls.get_maximum_num_blocks_by_compiled_model(
-            compiled_models=compiled_models,
-            model_config=model_config,
-            rbln_config=rbln_config,
+        # Get the actual memory allocation of each node by key
+        alloc_memory_per_node_by_key: Dict[str, List[int]] = compiled_models["prefill"].get_alloc_per_node_by_key()
+        alloc_memory_by_key: Dict[str, int] = {
+            key: sum(memory_per_node) for key, memory_per_node in alloc_memory_per_node_by_key.items()
+        }
+        for batch_size in rbln_config.decoder_batch_sizes:
+            for key, memory_per_node in (
+                compiled_models[f"decoder_batch_{batch_size}"].get_alloc_per_node_by_key().items()
+            ):
+                alloc_memory_by_key[key] += sum(memory_per_node)
+        alloc_memory_by_key.pop("PortRecur", None)  # Old compiler's kv-cache Key
+        alloc_memory_by_key.pop("DramTensor", None)  # kv-cache
+        kernel_size = alloc_memory_by_key.pop("Kernel")  # model weight
+
+        # Get the maximum number of blocks that can be allocated
+        buffer = sum(alloc_memory_by_key.values())
+        max_num_blocks = cls.get_maximum_num_blocks(
+            config=model_config,
+            tensor_parallel_size=rbln_config.tensor_parallel_size,
+            kvcache_block_size=rbln_config.kvcache_block_size,
+            kernel_size=kernel_size,
+            buffer=buffer,
         )
 
         # Since our estimation logic is not always accurate,

optimum/rbln/transformers/modeling_generic.py

@@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Optional, Union
 from torch import nn
 from transformers import (
     AutoModel,
+    AutoModelForAudioClassification,
     AutoModelForDepthEstimation,
     AutoModelForImageClassification,
     AutoModelForMaskedLM,
@@ -41,6 +42,7 @@ from ..modeling import RBLNModel
 from ..utils.logging import get_logger
 from .configuration_generic import (
     RBLNImageModelConfig,
+    RBLNModelForAudioClassificationConfig,
     RBLNTransformerEncoderConfig,
 )
 
@@ -57,7 +59,7 @@ class RBLNTransformerEncoder(RBLNModel):
     rbln_dtype = "int64"
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNTransformerEncoderConfig) -> nn.Module:
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNTransformerEncoderConfig) -> nn.Module:
         class TransformerEncoderWrapper(nn.Module):
             # Parameters to disable for RBLN compilation
             DISABLED_PARAMS = {"return_dict", "use_cache"}
@@ -266,7 +268,7 @@ class RBLNModelForDepthEstimation(RBLNImageModel):
     auto_model_class = AutoModelForDepthEstimation
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
         class ImageModelWrapper(nn.Module):
             def __init__(self, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
                 super().__init__()
@@ -278,3 +280,60 @@ class RBLNModelForDepthEstimation(RBLNImageModel):
                 return output.predicted_depth
 
         return ImageModelWrapper(model, rbln_config).eval()
+
+
+class RBLNModelForAudioClassification(RBLNModel):
+    """
+    This is a generic model class that will be instantiated as one of the model classes of the library (with a audio classification head) when created with the from_pretrained() class method
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based AudioClassification models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers AudioClassification model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    Currently, this model class only supports the 'AST' model from the transformers library. Future updates may include support for additional model types.
+    """
+
+    auto_model_class = AutoModelForAudioClassification
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: "AutoFeatureExtractor" = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "PretrainedConfig" = None,
+        rbln_config: Optional[RBLNModelForAudioClassificationConfig] = None,
+    ) -> RBLNModelForAudioClassificationConfig:
+        if rbln_config.num_mel_bins is None:
+            rbln_config.num_mel_bins = getattr(model_config, "num_mel_bins", None)
+            if rbln_config.num_mel_bins is None:
+                for feature_extractor in preprocessors:
+                    if hasattr(feature_extractor, "num_mel_bins"):
+                        rbln_config.num_mel_bins = feature_extractor.num_mel_bins
+                        break
+
+        if rbln_config.num_mel_bins is None:
+            raise ValueError("`num_mel_bins` should be specified!")
+
+        if rbln_config.max_length is None:
+            rbln_config.max_length = getattr(model_config, "max_length", None)
+            for feature_extractor in preprocessors:
+                if hasattr(feature_extractor, "max_length"):
+                    rbln_config.max_length = feature_extractor.max_length
+                    break
+
+        if rbln_config.max_length is None:
+            raise ValueError("`max_length` should be specified!")
+
+        input_info = [
+            (
+                "input_values",
+                [rbln_config.batch_size, rbln_config.max_length, rbln_config.num_mel_bins],
+                "float32",
+            ),
+        ]
+
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
+        return rbln_config

optimum/rbln/transformers/models/audio_spectrogram_transformer/configuration_audio_spectrogram_transformer.py

@@ -12,36 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional
+from ...configuration_generic import RBLNModelForAudioClassificationConfig
 
-from ....configuration_utils import RBLNModelConfig
-from ....utils.deprecation import deprecate_kwarg
 
-
-class RBLNASTForAudioClassificationConfig(RBLNModelConfig):
+class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
     """
     Configuration class for RBLNASTForAudioClassification.
     """
-
-    @deprecate_kwarg(old_name="num_mel_bins", version="0.10.0")
-    def __init__(
-        self,
-        batch_size: Optional[int] = None,
-        max_length: Optional[int] = None,
-        **kwargs: Any,
-    ):
-        """
-        Args:
-            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            max_length (Optional[int]): Maximum length of the audio input in time dimension.
-            kwargs: Additional arguments passed to the parent RBLNModelConfig.
-
-        Raises:
-            ValueError: If batch_size is not a positive integer.
-        """
-        super().__init__(**kwargs)
-        self.batch_size = batch_size or 1
-        if not isinstance(self.batch_size, int) or self.batch_size < 0:
-            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
-
-        self.max_length = max_length

optimum/rbln/transformers/models/audio_spectrogram_transformer/modeling_audio_spectrogram_transformer.py

@@ -12,80 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional
+from ...modeling_generic import RBLNModelForAudioClassification
 
-import torch
-from transformers import AutoModelForAudioClassification
-from transformers.modeling_outputs import SequenceClassifierOutput
 
-from ....configuration_utils import RBLNCompileConfig
-from ....modeling import RBLNModel
-from .configuration_audio_spectrogram_transformer import RBLNASTForAudioClassificationConfig
-
-
-if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, PretrainedConfig, PreTrainedModel
-
-
-class RBLNASTForAudioClassification(RBLNModel):
+class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
     """
     Audio Spectrogram Transformer model with an audio classification head on top (a linear layer on top of the pooled output) e.g. for datasets like AudioSet, Speech Commands v2.
-    This model inherits from [RBLNModelForAudioClassification]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNModelForAudioClassification`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
-    A class to convert and run pre-trained transformer-based ASTForAudioClassification models on RBLN devices.
-    It implements the methods to convert a pre-trained transformers ASTForAudioClassification model into a RBLN transformer model by:
+    A class to convert and run pre-trained transformer-based `ASTForAudioClassification` models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `ASTForAudioClassification` model into a RBLN transformer model by:
 
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN Compiler.
     """
-
-    auto_model_class = AutoModelForAudioClassification
-
-    @classmethod
-    def _update_rbln_config(
-        cls,
-        preprocessors: "AutoFeatureExtractor" = None,
-        model: Optional["PreTrainedModel"] = None,
-        model_config: "PretrainedConfig" = None,
-        rbln_config: Optional[RBLNASTForAudioClassificationConfig] = None,
-    ) -> RBLNASTForAudioClassificationConfig:
-        num_mel_bins = getattr(model_config, "num_mel_bins", None)
-
-        if rbln_config.max_length is None:
-            rbln_config.max_length = getattr(model_config, "max_length", None)
-            for feature_extractor in preprocessors:
-                if hasattr(feature_extractor, "max_length"):
-                    rbln_config.max_length = feature_extractor.max_length
-                    break
-
-        if rbln_config.max_length is None:
-            raise ValueError("max_length should be specified!")
-
-        input_info = [
-            (
-                "input_values",
-                [rbln_config.batch_size, rbln_config.max_length, num_mel_bins],
-                "float32",
-            ),
-        ]
-
-        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
-        return rbln_config
-
-    def forward(self, input_values: torch.Tensor, **kwargs) -> SequenceClassifierOutput:
-        """
-        Forward pass for the RBLN-optimized Audio Spectrogram Transformer model for audio classification.
-
-        Args:
-            input_values (torch.FloatTensor of shape (batch_size, max_length, num_mel_bins)):
-                Float values mel features extracted from the raw audio waveform. Raw audio waveform can be obtained by
-                loading a .flac or .wav audio file into an array of type list[float], a numpy.ndarray or a torch.Tensor, *e.g.* via
-                the torchcodec library (pip install torchcodec) or the soundfile library (pip install soundfile).
-                To prepare the array into input_features, the [AutoFeatureExtractor] should be used for extracting the
-                mel features, padding and conversion into a tensor of type torch.FloatTensor.
-
-        Returns:
-            Returns a SequenceClassifierOutput object.
-        """
-
-        return super().forward(input_values, **kwargs)