optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a6__py3-none-any.whl

optimum/rbln/transformers/modeling_generic.py

@@ -21,155 +21,242 @@ different model architectures.
 """
 
 import inspect
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from transformers import (
+    AutoModel,
     AutoModelForAudioClassification,
+    AutoModelForDepthEstimation,
     AutoModelForImageClassification,
     AutoModelForMaskedLM,
     AutoModelForQuestionAnswering,
     AutoModelForSequenceClassification,
+    AutoModelForTextEncoding,
     PretrainedConfig,
 )
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    DepthEstimatorOutput,
+    ImageClassifierOutput,
+    MaskedLMOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+)
 
+from ..configuration_utils import RBLNCompileConfig
 from ..modeling import RBLNModel
-from ..modeling_config import RBLNCompileConfig, RBLNConfig
 from ..utils.logging import get_logger
+from .configuration_generic import (
+    RBLNModelForAudioClassificationConfig,
+    _RBLNImageModelConfig,
+    _RBLNTransformerEncoderConfig,
+)
 
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
 logger = get_logger()
 
 
-class RBLNModelForQuestionAnswering(RBLNModel):
-    auto_model_class = AutoModelForQuestionAnswering
+class _RBLNTransformerEncoder(RBLNModel):
+    auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
+    rbln_dtype = "int64"
+    output_class = BaseModelOutput
+    output_key = "last_hidden_state"
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+        rbln_config: Optional[_RBLNTransformerEncoderConfig] = None,
+    ) -> _RBLNTransformerEncoderConfig:
+        return cls.update_rbln_config_for_transformers_encoder(
+            preprocessors=preprocessors,
+            model=model,
+            model_config=model_config,
+            rbln_config=rbln_config,
+        )
 
-        if rbln_max_seq_len is None:
-            for tokenizer in preprocessors:
-                if hasattr(tokenizer, "model_max_length"):
-                    rbln_max_seq_len = tokenizer.model_max_length
-                    break
-            if rbln_max_seq_len is None:
-                raise ValueError("`rbln_max_seq_len` should be specified!")
+    @classmethod
+    def update_rbln_config_for_transformers_encoder(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[_RBLNTransformerEncoderConfig] = None,
+    ) -> _RBLNTransformerEncoderConfig:
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        if rbln_config.max_seq_len is None:
+            rbln_config.max_seq_len = max_position_embeddings
+            if rbln_config.max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_config.max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_config.max_seq_len is None:
+                    raise ValueError("`max_seq_len` should be specified!")
 
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
+        if max_position_embeddings is not None and rbln_config.max_seq_len > max_position_embeddings:
+            raise ValueError("`max_seq_len` should be less or equal than max_position_embeddings!")
 
-        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
+        signature_params = inspect.signature(model.forward).parameters.keys()
 
-        if rbln_model_input_names is None:
+        if rbln_config.model_input_names is None:
             for tokenizer in preprocessors:
                 if hasattr(tokenizer, "model_input_names"):
-                    rbln_model_input_names = [name for name in signature_params if name in tokenizer.model_input_names]
+                    rbln_config.model_input_names = [
+                        name for name in signature_params if name in tokenizer.model_input_names
+                    ]
 
-                    invalid_params = set(rbln_model_input_names) - set(signature_params)
+                    invalid_params = set(rbln_config.model_input_names) - set(signature_params)
                     if invalid_params:
                         raise ValueError(f"Invalid model input names: {invalid_params}")
                     break
-            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
-                rbln_model_input_names = cls.rbln_model_input_names
-            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                raise ValueError(
-                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
-                    f"and be sure to make the order of the inputs same as QuestionAnswering forward() arguments like ({list(signature_params)})"
-                )
+            if rbln_config.model_input_names is None and cls.rbln_model_input_names is not None:
+                rbln_config.model_input_names = cls.rbln_model_input_names
+
         else:
-            invalid_params = set(rbln_model_input_names) - set(signature_params)
+            invalid_params = set(rbln_config.model_input_names) - set(signature_params)
             if invalid_params:
                 raise ValueError(f"Invalid model input names: {invalid_params}")
-            rbln_model_input_names = [name for name in signature_params if name in rbln_model_input_names]
+            rbln_config.model_input_names = [
+                name for name in signature_params if name in rbln_config.model_input_names
+            ]
+
+        if rbln_config.model_input_names is None or len(rbln_config.model_input_names) == 0:
+            raise ValueError(
+                "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`. "
+                "This is an internal error. Please report it to the developers."
+            )
 
         input_info = [
-            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in rbln_model_input_names
+            (model_input_name, [rbln_config.batch_size, rbln_config.max_seq_len], cls.rbln_dtype)
+            for model_input_name in rbln_config.model_input_names
         ]
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config
 
 
-class RBLNModelForImageClassification(RBLNModel):
-    """
-    This is a generic model class that will be instantiated as one of the model classes of the library (with a image classification head) when created with the from_pretrained() class method
-    """
-
-    auto_model_class = AutoModelForImageClassification
+class _RBLNImageModel(RBLNModel):
+    auto_model_class = AutoModel
+    main_input_name = "pixel_values"
+    output_class = BaseModelOutput
+    output_key = "last_hidden_state"
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_image_size = rbln_kwargs.get("image_size", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_config: Optional[_RBLNImageModelConfig] = None,
+    ) -> _RBLNImageModelConfig:
+        return cls.update_rbln_config_for_image_model(
+            preprocessors=preprocessors,
+            model=model,
+            model_config=model_config,
+            rbln_config=rbln_config,
+        )
 
-        if rbln_image_size is None:
+    @classmethod
+    def update_rbln_config_for_image_model(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[_RBLNImageModelConfig] = None,
+    ) -> _RBLNImageModelConfig:
+        if rbln_config.image_size is None:
             for processor in preprocessors:
                 if hasattr(processor, "size"):
                     if all(required_key in processor.size.keys() for required_key in ["height", "width"]):
-                        rbln_image_size = (processor.size["height"], processor.size["width"])
+                        rbln_config.image_size = (processor.size["height"], processor.size["width"])
                     elif "shortest_edge" in processor.size.keys():
-                        rbln_image_size = (processor.size["shortest_edge"], processor.size["shortest_edge"])
+                        rbln_config.image_size = (processor.size["shortest_edge"], processor.size["shortest_edge"])
                     elif "longest_edge" in processor.size.keys():
-                        rbln_image_size = (processor.size["longest_edge"], processor.size["longest_edge"])
+                        rbln_config.image_size = (processor.size["longest_edge"], processor.size["longest_edge"])
                     break
 
-            if rbln_image_size is None:
-                rbln_image_size = model_config.image_size
-
-            if rbln_image_size is None:
-                raise ValueError("`rbln_image_size` should be specified!")
+            if rbln_config.image_size is None:
+                rbln_config.image_size = model_config.image_size
 
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        if isinstance(rbln_image_size, int):
-            rbln_image_height, rbln_image_width = rbln_image_size, rbln_image_size
-        elif isinstance(rbln_image_size, (list, tuple)):
-            rbln_image_height, rbln_image_width = rbln_image_size[0], rbln_image_size[1]
-        elif isinstance(rbln_image_size, dict):
-            rbln_image_height, rbln_image_width = rbln_image_size["height"], rbln_image_size["width"]
-        else:
-            raise ValueError(
-                "`rbln_image_size` should be `int` (ex. 224), `tuple` (ex. 224, 224), `dict` (ex. {'height': 224, 'width': 224}) format"
-            )
+            if rbln_config.image_size is None:
+                raise ValueError("`image_size` should be specified!")
 
         input_info = [
             (
-                "pixel_values",
-                [rbln_batch_size, 3, rbln_image_height, rbln_image_width],
+                cls.main_input_name,
+                [rbln_config.batch_size, 3, rbln_config.image_height, rbln_config.image_width],
                 "float32",
             )
         ]
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        return RBLNConfig(rbln_cls=cls.__name__, compile_cfgs=[rbln_compile_config], rbln_kwargs=rbln_kwargs)
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
+        return rbln_config
+
+
+class RBLNModelForQuestionAnswering(_RBLNTransformerEncoder):
+    auto_model_class = AutoModelForQuestionAnswering
+    rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
+    output_class = QuestionAnsweringModelOutput
+
+    def _prepare_output(self, output, return_dict):
+        """
+        Prepare QuestionAnswering specific output format.
+        """
+        start_logits, end_logits = output
+
+        if not return_dict:
+            return (start_logits, end_logits)
+        else:
+            return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits)
+
+
+class RBLNModelForSequenceClassification(_RBLNTransformerEncoder):
+    auto_model_class = AutoModelForSequenceClassification
+    rbln_model_input_names = ["input_ids", "attention_mask"]
+    output_class = SequenceClassifierOutput
+    output_key = "logits"
+
+
+class RBLNModelForMaskedLM(_RBLNTransformerEncoder):
+    auto_model_class = AutoModelForMaskedLM
+    rbln_model_input_names = ["input_ids", "attention_mask"]
+    output_class = MaskedLMOutput
+    output_key = "logits"
+
+
+class RBLNModelForTextEncoding(_RBLNTransformerEncoder):
+    auto_model_class = AutoModelForTextEncoding
+    rbln_model_input_names = ["input_ids", "attention_mask"]
+
+
+class RBLNTransformerEncoderForFeatureExtraction(_RBLNTransformerEncoder):
+    # TODO: RBLNModel is also for feature extraction.
+    auto_model_class = AutoModel
+    rbln_model_input_names = ["input_ids", "attention_mask"]
+    output_class = BaseModelOutput
+    output_key = "last_hidden_state"
+
+
+class RBLNModelForImageClassification(_RBLNImageModel):
+    auto_model_class = AutoModelForImageClassification
+    output_class = ImageClassifierOutput
+    output_key = "logits"
+
+
+class RBLNModelForDepthEstimation(_RBLNImageModel):
+    auto_model_class = AutoModelForDepthEstimation
+    output_class = DepthEstimatorOutput
+    output_key = "predicted_depth"
 
 
 class RBLNModelForAudioClassification(RBLNModel):
@@ -186,219 +273,45 @@ class RBLNModelForAudioClassification(RBLNModel):
     """
 
     auto_model_class = AutoModelForAudioClassification
+    output_class = SequenceClassifierOutput
+    output_key = "logits"
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
-        preprocessors: "AutoFeatureExtractor",
-        model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_max_length = rbln_kwargs.get("max_length", None)
-        rbln_num_mel_bins = rbln_kwargs.get("num_mel_bins", None)
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        if rbln_num_mel_bins is None:
-            rbln_num_mel_bins = getattr(model_config, "num_mel_bins", None)
-            if rbln_num_mel_bins is None:
+        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_num_mel_bins = feature_extractor.num_mel_bins
+                        rbln_config.num_mel_bins = feature_extractor.num_mel_bins
                         break
 
-        if rbln_num_mel_bins is None:
-            raise ValueError("`rbln_num_mel_bins` should be specified!")
+        if rbln_config.num_mel_bins is None:
+            raise ValueError("`num_mel_bins` should be specified!")
 
-        if rbln_max_length is None:
-            rbln_max_length = getattr(model_config, "max_length", 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_max_length = feature_extractor.max_length
+                    rbln_config.max_length = feature_extractor.max_length
                     break
 
-        if rbln_max_length is None:
-            raise ValueError("`rbln_max_length` should be specified!")
+        if rbln_config.max_length is None:
+            raise ValueError("`max_length` should be specified!")
 
         input_info = [
             (
                 "input_values",
-                [rbln_batch_size, rbln_max_length, rbln_num_mel_bins],
+                [rbln_config.batch_size, rbln_config.max_length, rbln_config.num_mel_bins],
                 "float32",
             ),
         ]
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-        rbln_config.model_cfg.update(
-            {
-                "batch_size": rbln_batch_size,
-                "max_length": rbln_max_length,
-                "num_mel_bins": rbln_num_mel_bins,
-            }
-        )
-        return rbln_config
-
-
-class RBLNModelForSequenceClassification(RBLNModel):
-    """
-    This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence 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 SequenceClassification models on RBLN devices.
-    It implements the methods to convert a pre-trained transformers SequenceClassification 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 supports the 'XLMRoberta' and 'Roberta' model from the transformers library. Future updates may include support for additional model types.
-    """
-
-    auto_model_class = AutoModelForSequenceClassification
-
-    @classmethod
-    def _get_rbln_config(
-        cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-
-        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
-            model_config, "max_position_embeddings", None
-        )
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = max_position_embeddings
-            if rbln_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_max_seq_len is None:
-                    raise ValueError("`rbln_max_seq_len` should be specified!")
-
-        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
-
-        if rbln_model_input_names is None:
-            for tokenizer in preprocessors:
-                if hasattr(tokenizer, "model_input_names"):
-                    rbln_model_input_names = [name for name in signature_params if name in tokenizer.model_input_names]
-
-                    invalid_params = set(rbln_model_input_names) - set(signature_params)
-                    if invalid_params:
-                        raise ValueError(f"Invalid model input names: {invalid_params}")
-                    break
-            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
-                rbln_model_input_names = cls.rbln_model_input_names
-            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                raise ValueError(
-                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
-                    f"and be sure to make the order of the inputs same as SequenceClassification forward() arguments like ({list(signature_params)})"
-                )
-        else:
-            invalid_params = set(rbln_model_input_names) - set(signature_params)
-            if invalid_params:
-                raise ValueError(f"Invalid model input names: {invalid_params}")
-            rbln_model_input_names = [name for name in signature_params if name in rbln_model_input_names]
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        input_info = [
-            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in rbln_model_input_names
-        ]
-
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
-        return rbln_config
-
-
-class RBLNModelForMaskedLM(RBLNModel):
-    auto_model_class = AutoModelForMaskedLM
-
-    @classmethod
-    def _get_rbln_config(
-        cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-
-        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
-            model_config, "max_position_embeddings", None
-        )
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = max_position_embeddings
-            if rbln_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_max_seq_len is None:
-                    raise ValueError("`rbln_max_seq_len` should be specified!")
-
-        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        signature_params = inspect.signature(cls.get_hf_class().forward).parameters.keys()
-
-        if rbln_model_input_names is None:
-            for tokenizer in preprocessors:
-                if hasattr(tokenizer, "model_input_names"):
-                    rbln_model_input_names = [name for name in signature_params if name in tokenizer.model_input_names]
-
-                    invalid_params = set(rbln_model_input_names) - set(signature_params)
-                    if invalid_params:
-                        raise ValueError(f"Invalid model input names: {invalid_params}")
-                    break
-            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
-                rbln_model_input_names = cls.rbln_model_input_names
-            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                raise ValueError(
-                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
-                    f"and be sure to make the order of the inputs same as MaskedLM forward() arguments like ({list(signature_params)})"
-                )
-        else:
-            invalid_params = set(rbln_model_input_names) - set(signature_params)
-            if invalid_params:
-                raise ValueError(f"Invalid model input names: {invalid_params}")
-            rbln_model_input_names = [name for name in signature_params if name in rbln_model_input_names]
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        input_info = [
-            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in rbln_model_input_names
-        ]
-
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config