optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.4a2__py3-none-any.whl

optimum/rbln/modeling_base.py

@@ -15,7 +15,6 @@
 import importlib
 import os
 import shutil
-from abc import ABC
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
@@ -39,7 +38,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
-class PreTrainedModel(ABC):  # noqa: F811
+class PreTrainedModel:  # noqa: F811
     pass
 
 
@@ -63,7 +62,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         rbln_compiled_models: Optional[rebel.RBLNCompiledModel] = None,
-        rbln_submodules: List["RBLNBaseModel"] = [],
+        rbln_submodules: Optional[List["RBLNBaseModel"]] = None,
         **kwargs,
     ):
         self.model = models
@@ -71,7 +70,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         self.rbln_config = rbln_config
         if not rbln_config.is_frozen():
             raise RuntimeError("`rbln_config` must be frozen. Please call `rbln_config.freeze()` first.")
-
         self.compiled_models = rbln_compiled_models
 
         # Registers the RBLN classes into the transformers AutoModel classes to avoid warnings when creating
@@ -107,6 +105,8 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             self.model_save_dir = model_save_dir
         self.subfolder = subfolder
 
+        if rbln_submodules is None:
+            rbln_submodules = []
         self.rbln_submodules = rbln_submodules
         self.__post_init__(**kwargs)
 
@@ -182,7 +182,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         # passed from compile function
         rbln_config: Optional[RBLNModelConfig] = None,
         rbln_compiled_models: Optional[Dict[str, rebel.RBLNCompiledModel]] = None,
-        rbln_submodules: List["RBLNBaseModel"] = [],
+        rbln_submodules: Optional[List["RBLNBaseModel"]] = None,
         **kwargs,
     ) -> "RBLNBaseModel":
         if rbln_compiled_models is None:
@@ -218,8 +218,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 )
 
             if len(cls._rbln_submodules) > 0:
-                rbln_submodules = cls._load_submodules(model_save_dir=model_id, rbln_config=rbln_config, **kwargs)
-            else:
+                if rbln_submodules is None:
+                    rbln_submodules = cls._load_submodules(model_save_dir=model_id, rbln_config=rbln_config, **kwargs)
+            elif rbln_submodules is None:
                 rbln_submodules = []
 
             rbln_config.freeze()
@@ -280,9 +281,12 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         config: "PretrainedConfig",
         model_save_dir: Union[Path, str],
         subfolder: Union[Path, str],
-        rbln_submodules: List["RBLNBaseModel"] = [],
+        rbln_submodules: Optional[List["RBLNBaseModel"]] = None,
         **kwargs,
     ):
+        if rbln_submodules is None:
+            rbln_submodules = []
+
         if isinstance(model_save_dir, str):
             model_save_dir = Path(model_save_dir)
 

optimum/rbln/transformers/configuration_generic.py

@@ -118,30 +118,3 @@ 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 typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+from collections import Counter, defaultdict
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
-from optimum.rbln.transformers.models.decoderonly.configuration_decoderonly import (
-    RBLNDecoderOnlyModelForCausalLMConfig,
-)
+import rebel
 
 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 rebel import RBLNCompiledModel
-    from transformers import PretrainedConfig
+    from transformers import PretrainedConfig, PreTrainedModel
 
 
 DEFAULT_FLASH_ATTN_PARTITION_LENGTH = 16_384
@@ -115,128 +115,261 @@ 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(
+    def get_maximum_num_blocks_by_model(
         cls,
-        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,
+        model: "PreTrainedModel",
+        model_config: "PretrainedConfig",
+        rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
     ) -> int:
-        # 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
+        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)})"
             )
-            params = n_model_params - lm_heads_params
-            layer_nbytes = (
-                align_2MB(params * nbits_per_param // 8 / num_layers / tensor_parallel_size)
-                * num_layers
+        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 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 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
                 * tensor_parallel_size
+                * 2
+            )  # *2 for key and value
+
+            return dram_per_block
+
+        num_attention_heads = getattr(model_config, "n_head", None) or 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 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 -= kernel_size
+            available_dram -= swa_kv_nbytes
 
-        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
+        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,
+        )
 
-        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))
+        if dram_per_block == 0:
+            raise ValueError("DRAM per block is calculated as zero, cannot estimate maximum number of blocks.")
 
+        max_n_blocks = available_dram // dram_per_block
         return max_n_blocks
 
     @classmethod
     def maybe_suggest_kvcache_num_blocks(
         cls,
-        compiled_models: Dict[str, "RBLNCompiledModel"],
+        compiled_models: Dict[str, "rebel.RBLNCompiledModel"],
         model_config: "PretrainedConfig",
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
     ) -> None:
-        # 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,
+        max_num_blocks = cls.get_maximum_num_blocks_by_compiled_model(
+            compiled_models=compiled_models,
+            model_config=model_config,
+            rbln_config=rbln_config,
         )
 
         # Since our estimation logic is not always accurate,

optimum/rbln/transformers/modeling_generic.py

@@ -26,7 +26,6 @@ from typing import TYPE_CHECKING, Optional, Union
 from torch import nn
 from transformers import (
     AutoModel,
-    AutoModelForAudioClassification,
     AutoModelForDepthEstimation,
     AutoModelForImageClassification,
     AutoModelForMaskedLM,
@@ -42,7 +41,6 @@ from ..modeling import RBLNModel
 from ..utils.logging import get_logger
 from .configuration_generic import (
     RBLNImageModelConfig,
-    RBLNModelForAudioClassificationConfig,
     RBLNTransformerEncoderConfig,
 )
 
@@ -59,7 +57,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"}
@@ -268,7 +266,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__()
@@ -280,60 +278,3 @@ 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/modeling_outputs.py

@@ -24,6 +24,7 @@ class RBLNDecoderOnlyOutput(ModelOutput):
     logits: torch.FloatTensor = None
     generate_idx: torch.Tensor = None
     padded_cache_lengths: int = None
+    hidden_states: Tuple[torch.FloatTensor] = None
 
 
 @dataclass

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

@@ -12,10 +12,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...configuration_generic import RBLNModelForAudioClassificationConfig
+from typing import Any, Optional
 
+from ....configuration_utils import RBLNModelConfig
+from ....utils.deprecation import deprecate_kwarg
 
-class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
+
+class RBLNASTForAudioClassificationConfig(RBLNModelConfig):
     """
     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