optimum-rbln 0.1.9__py3-none-any.whl → 0.1.12__py3-none-any.whl

optimum/rbln/modeling_config.py

@@ -23,24 +23,38 @@
 
 import copy
 import json
-from collections import UserDict
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
+import rebel
 import torch
 
+from .__version__ import __version__
+from .utils.runtime_utils import ContextRblnConfig
+
 
 DEFAULT_COMPILED_MODEL_NAME = "compiled_model"
 DEFAULT_MOD_NAME = "default"
 
 
 @dataclass
-class RBLNRuntimeConfig:
+class RBLNCompileConfig:
+    """
+    Configuration for RBLN compilation.
+
+    Attributes:
+        compiled_model_name (str): Name of the compiled model.
+        mod_name (str): Name of the RBLN module.
+        input_info (List[Tuple[str, Tuple[int], Optional[str]]]): Information about input tensors.
+        fusion (Optional[bool]): Whether to use fusion optimization.
+        npu (Optional[str]): NPU configuration.
+        tensor_parallel_size (Optional[int]): Size for tensor parallelism.
+    """
+
     compiled_model_name: str = DEFAULT_COMPILED_MODEL_NAME
-    rbln_mod_name: str = DEFAULT_MOD_NAME
+    mod_name: str = DEFAULT_MOD_NAME
     input_info: List[Tuple[str, Tuple[int], Optional[str]]] = None
-    batch_size: Optional[int] = None
     fusion: Optional[bool] = None
     npu: Optional[str] = None
     tensor_parallel_size: Optional[int] = None
@@ -48,8 +62,14 @@ class RBLNRuntimeConfig:
     @staticmethod
     def normalize_dtype(dtype):
         """
-        framework's dtype to string.
+        Convert framework-specific dtype to string representation.
         i.e. torch.float32 -> "float32"
+
+        Args:
+            dtype: The input dtype (can be string, torch dtype, or numpy dtype).
+
+        Returns:
+            str: The normalized string representation of the dtype.
         """
         if isinstance(dtype, str):
             return dtype
@@ -60,13 +80,12 @@ class RBLNRuntimeConfig:
             return dtype
 
     def __post_init__(self):
-        self.input_info = [(i[0], i[1], RBLNRuntimeConfig.normalize_dtype(i[2]) or "float32") for i in self.input_info]
+        self.input_info = [(i[0], i[1], RBLNCompileConfig.normalize_dtype(i[2]) or "float32") for i in self.input_info]
 
-    def update(self, **kwargs):
+    def update(self, kwargs: Dict[str, Any]):
         self.compiled_model_name = kwargs.get("compiled_model_name", self.compiled_model_name)
-        self.rbln_mod_name = kwargs.get("rbln_mod_name", self.rbln_mod_name)
+        self.mod_name = kwargs.get("mod_name", self.mod_name)
         self.input_info = kwargs.get("input_info", self.input_info)
-        self.batch_size = kwargs.get("batch_size", self.batch_size)
         self.fusion = kwargs.get("fusion", self.fusion)
         self.npu = kwargs.get("npu", self.npu)
         self.tensor_parallel_size = kwargs.get("tensor_parallel_size", self.tensor_parallel_size)
@@ -86,84 +105,191 @@ class RBLNRuntimeConfig:
         return asdict(self)
 
 
-class RBLNConfig(UserDict):
-    def __init__(self, runtime_cfgs: Dict[str, List[RBLNRuntimeConfig]], _rbln_meta: Dict[str, Any] = None):
-        """Configurations for RBLN model compilation and inference.
+RUNTIME_KEYWORDS = ["create_runtimes", "optimize_host_memory", "device", "device_map"]
+COMPILE_KEYWORDS = ["compiled_model_name", "mod_name", "input_info", "fusion", "npu", "tensor_parallel_size"]
 
-        Args:
-            _rbln_meta (Dict[str, Any], optional):
-                     Any rbln-specific configurations.
-                     (i.e. max_seq_len for language models, image_size for image models).
-                     Defaults to None.
-        """
-        super().__init__(runtime_cfgs)
-        if _rbln_meta:
-            self.meta = _rbln_meta
-        else:
-            self.meta: Dict[str, Any] = {}
 
-    @staticmethod
-    def from_rbln_configs(rbln_configs: List["RBLNConfig"], names: Optional[List[str]] = None) -> "RBLNConfig":
-        # assume each rbln_config has exact one rbln_runtime_config
-        names = [None] * len(rbln_configs) if names is None else names
-        runtime_cfgs = []
-        for name, cfg in zip(names, rbln_configs):
-            if len(cfg) > 1:
-                msg = (
-                    "`from_rbln_configs` requires exact one `RBLNRuntimeConfig` for each `RBLNConfig`."
-                    f"But got {len(cfg)} `RBLNRuntimeConfig`."
-                )
-                raise RuntimeError(msg)
+class RBLNConfig:
+    """
+    Configuration for single RBLN OptimizedModel, representing multiple compiled models.
 
-            runtime_cfg = cfg[list(cfg.keys())[0]][0]
-            runtime_cfg = copy.deepcopy(runtime_cfg)
-            if name is not None:
-                runtime_cfg.compiled_model_name = name
-            runtime_cfgs.append(runtime_cfg)
+    Attributes:
+        compile_cfgs (List[RBLNCompileConfig]): Compilation configurations.
+        meta (dict): Metadata including version and class information.
+        runtime_cfg (dict): Runtime-specific configuration.
+    """
 
-        metas = [cfg.meta for cfg in rbln_configs]
-        merged_meta = {k: v for meta in metas for k, v in meta.items()}
+    # It represents multiple compiled model, one of each can have multiple runtimes.
+    def __init__(
+        self,
+        rbln_cls,
+        compile_cfgs: List[RBLNCompileConfig],
+        rbln_kwargs=None,
+        meta=None,
+    ) -> None:
+        if rbln_kwargs is None:
+            rbln_kwargs = {}
+        else:
+            rbln_kwargs = copy.deepcopy(rbln_kwargs)
 
-        return RBLNConfig.from_rbln_runtime_configs(runtime_cfgs, _rbln_meta=merged_meta)
+        # meta : class, version and other informations.
+        if meta is None:
+            self.meta = {"version": __version__, "cls": rbln_cls}
+        else:
+            self.meta = meta
 
-    @staticmethod
-    def from_rbln_runtime_configs(
-        rbln_runtime_configs: List[RBLNRuntimeConfig],
-        _rbln_meta: Dict[str, Any] = None,
-    ) -> "RBLNConfig":
-        cfgs: Dict[str, List[RBLNRuntimeConfig]] = {}
-        for rbln_runtime_config in rbln_runtime_configs:
-            if rbln_runtime_config.compiled_model_name in cfgs:
-                cfgs[rbln_runtime_config.compiled_model_name].append(rbln_runtime_config)
-            else:
-                cfgs[rbln_runtime_config.compiled_model_name] = [rbln_runtime_config]
-        return RBLNConfig(cfgs, _rbln_meta=_rbln_meta)
+        # compile_cfgs : compile args for each runtime
+        self.compile_cfgs = compile_cfgs
+        for compile_cfg in self.compile_cfgs:
+            compile_cfg.update(rbln_kwargs)
+        for K in COMPILE_KEYWORDS:
+            rbln_kwargs.pop(K, None)
+
+        # runtime_cfg : Values that don't be saved / loaded.
+        self.runtime_cfg = {}
+        for runtime_key in RUNTIME_KEYWORDS:
+            if runtime_key in rbln_kwargs:
+                self.runtime_cfg[runtime_key] = rbln_kwargs.pop(runtime_key)
+
+        # model_cfg : All user-provided values such as "max_seq_len".
+        self.model_cfg: Dict[str, Any] = rbln_kwargs
 
     def save(self, dir_path: str):
         dir_path = Path(dir_path)
-        data = self.asdict()
-        data.update({"rbln_config_meta": self.meta})
+
+        s_json = {}
+        compile_cfgs = [asdict(cfg) for cfg in self.compile_cfgs]
+        s_json["_compile_cfgs"] = compile_cfgs
+        s_json["_meta"] = self.meta
+        s_json.update(self.model_cfg)
+
         with open(dir_path / "rbln_config.json", "w") as jsonf:
-            json.dump(data, jsonf, indent=2)
+            json.dump(s_json, jsonf, indent=2)
 
-    @staticmethod
-    def load(dir_path: str) -> "RBLNConfig":
+    @classmethod
+    def load(cls, dir_path: str) -> "RBLNConfig":
         dir_path = Path(dir_path)
         with open(dir_path / "rbln_config.json", "r") as jsonf:
             config_file = json.load(jsonf)
-        return RBLNConfig.fromdict(config_file)
 
-    def asdict(self):
-        dic = {k: [asdict(cfg) for cfg in cfgs] for k, cfgs in self.data.items()}
-        return dic
+        return cls.fromdict(config_file)
+
+    @classmethod
+    def fromdict(cls, dic: dict):
+        compile_cfgs = dic.pop("_compile_cfgs")
+        compile_cfgs = [RBLNCompileConfig(**cfg) for cfg in compile_cfgs]
+
+        meta = dic.pop("_meta")
+        rbln_cls = meta["cls"]
+
+        rbln_kwargs = dic
+        return cls(rbln_cls=rbln_cls, compile_cfgs=compile_cfgs, rbln_kwargs=rbln_kwargs, meta=meta)
+
+    def update_runtime_cfg(self, rbln_kwargs: Dict[str, Any]):
+        keys = list(rbln_kwargs.keys())
+        for key in keys:
+            if key in RUNTIME_KEYWORDS:
+                self.runtime_cfg[key] = rbln_kwargs[key]
+
+    def __repr__(self):
+        compile_cfgs_repr = [f"\n    {cfg!r}" for cfg in self.compile_cfgs]
+        return (
+            f"RBLNConfig(\n"
+            f"  rbln_cls={self.meta['cls']},\n"
+            f"  version='{self.meta['version']}',\n"
+            f"  compile_cfgs=[{''.join(compile_cfgs_repr)}\n  ],\n"
+            f"  model_cfg={self.model_cfg},\n"
+            f"  runtime_cfg={self.runtime_cfg}\n"
+            f")"
+        )
+
+    @property
+    def create_runtimes(self):
+        context = ContextRblnConfig.get_current_context()["create_runtimes"]
+        if context is not None:
+            return context
+        elif self.runtime_cfg.get("create_runtimes", None) is None:
+            return rebel.npu_is_available()
+        return self.runtime_cfg["create_runtimes"]
+
+    @property
+    def optimize_host_memory(self):
+        context = ContextRblnConfig.get_current_context()["optimize_host_memory"]
+        if context is not None:
+            return context
+        elif self.runtime_cfg.get("optimize_host_memory", None) is None:
+            return True
+        return self.runtime_cfg["optimize_host_memory"]
+
+    @property
+    def device(self):
+        context = ContextRblnConfig.get_current_context()["device"]
+        if context:
+            return context
+        elif self.runtime_cfg.get("device", None) is None:
+            return 0
+        return self.runtime_cfg["device"]
+
+    @property
+    def device_map(self):
+        context = ContextRblnConfig.get_current_context()["device_map"]
+        if context:
+            return context
+        elif self.runtime_cfg.get("device_map", None) is None:
+            rbln_device_map = {}
+            device_val = self.device
+            for cfg in self.compile_cfgs:
+                rbln_device_map[cfg.compiled_model_name] = device_val
+            return rbln_device_map
+        return self.runtime_cfg["device_map"]
+
+
+def use_rbln_config(fn):
+    """
+    If the function uses rbln_config and kwargs,
+        then extract `rbln_` prefix from kwargs.
+
+    If rbln_config is already an instance of RBLNConfig, then pass.
+    """
+
+    def merged_rbln_config_fn(*args, **kwargs):
+        rbln_kwargs = kwargs.pop("rbln_kwargs", None)
+        if rbln_kwargs is not None:
+            raise KeyError("`rbln_kwargs` cannot be specified when using `rbln_config`!")
+
+        rbln_config = kwargs.pop("rbln_config", None)
+
+        keys = list(kwargs.keys())
+        rbln_kwargs = {key[5:]: kwargs.pop(key) for key in keys if key.startswith("rbln_")}
+
+        if isinstance(rbln_config, RBLNConfig):
+            # merge runtime kwargs if exists.
+            runtime_rbln_kwargs = {k: rbln_kwargs.pop(k) for k in RUNTIME_KEYWORDS if k in rbln_kwargs}
+
+            # ignore internal keys and recover "rbln_" prefix
+            RBLN_INTERNAL_KEYS = {"compiled_models", "submodules"}
+            internal_kwargs = {"rbln_" + k: rbln_kwargs.pop(k) for k in RBLN_INTERNAL_KEYS if k in rbln_kwargs}
+
+            if len(rbln_kwargs) > 0:
+                raise KeyError(
+                    f"Failed to merging function argument : {rbln_kwargs.keys()}. "
+                    "If you passed `rbln_config` an instance of `RBLNConfig`, "
+                    "then none `rbln_` prefixes are allowed to be passed."
+                )
+            rbln_config.update_runtime_cfg(runtime_rbln_kwargs)
+            return fn(*args, **kwargs, **internal_kwargs, rbln_config=rbln_config)
+
+        elif rbln_config is None:
+            rbln_config_dict = {}
 
-    @staticmethod
-    def fromdict(dic: dict):
-        runtime_cfgs = {
-            k: [RBLNRuntimeConfig(**cfg) for cfg in cfgs] for k, cfgs in dic.items() if k != "rbln_config_meta"
-        }
-        if "rbln_config_meta" in dic:
-            meta = dic["rbln_config_meta"]
         else:
-            meta = None
-        return RBLNConfig(runtime_cfgs, _rbln_meta=meta)
+            rbln_config_dict = rbln_config
+
+        for key in rbln_config_dict:
+            if key in rbln_kwargs:
+                raise KeyError(f"Duplicated key in both `rbln_config` and rbln_{key}.")
+
+        rbln_kwargs.update(rbln_config_dict)
+        return fn(*args, **kwargs, rbln_config=rbln_kwargs)
+
+    return merged_rbln_config_fn

optimum/rbln/transformers/__init__.py

@@ -30,17 +30,38 @@ _import_structure = {
     "cache_utils": ["RebelDynamicCache"],
     "generation": ["BatchTextIteratorStreamer"],
     "models": [
+        "RBLNAutoModel",
+        "RBLNAutoModelForAudioClassification",
+        "RBLNAutoModelForCausalLM",
+        "RBLNAutoModelForCTC",
+        "RBLNAutoModelForDepthEstimation",
+        "RBLNAutoModelForImageClassification",
+        "RBLNAutoModelForMaskedLM",
+        "RBLNAutoModelForQuestionAnswering",
+        "RBLNAutoModelForSeq2SeqLM",
+        "RBLNAutoModelForSequenceClassification",
+        "RBLNAutoModelForSpeechSeq2Seq",
+        "RBLNAutoModelForVision2Seq",
+        "RBLNBartForConditionalGeneration",
+        "RBLNBartModel",
+        "RBLNBertModel",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
+        "RBLNCLIPVisionModel",
         "RBLNDPTForDepthEstimation",
+        "RBLNExaoneForCausalLM",
         "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
+        "RBLNQwen2ForCausalLM",
         "RBLNWav2Vec2ForCTC",
         "RBLNWhisperForConditionalGeneration",
         "RBLNLlamaForCausalLM",
+        "RBLNPhiForCausalLM",
+        "RBLNT5ForConditionalGeneration",
+        "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
-        "RBLNMistralForCausalLM",
         "RBLNXLMRobertaModel",
+        "RBLNMistralForCausalLM",
     ],
 }
 
@@ -48,14 +69,35 @@ if TYPE_CHECKING:
     from .cache_utils import RebelDynamicCache
     from .generation import BatchTextIteratorStreamer
     from .models import (
+        RBLNAutoModel,
+        RBLNAutoModelForAudioClassification,
+        RBLNAutoModelForCausalLM,
+        RBLNAutoModelForCTC,
+        RBLNAutoModelForDepthEstimation,
+        RBLNAutoModelForImageClassification,
+        RBLNAutoModelForMaskedLM,
+        RBLNAutoModelForQuestionAnswering,
+        RBLNAutoModelForSeq2SeqLM,
+        RBLNAutoModelForSequenceClassification,
+        RBLNAutoModelForSpeechSeq2Seq,
+        RBLNAutoModelForVision2Seq,
+        RBLNBartForConditionalGeneration,
+        RBLNBartModel,
+        RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
+        RBLNCLIPVisionModel,
         RBLNDPTForDepthEstimation,
+        RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
+        RBLNLlavaNextForConditionalGeneration,
         RBLNMidmLMHeadModel,
         RBLNMistralForCausalLM,
+        RBLNPhiForCausalLM,
+        RBLNQwen2ForCausalLM,
+        RBLNT5ForConditionalGeneration,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/transformers/models/__init__.py

@@ -21,13 +21,35 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+
+from .auto import (
+    RBLNAutoModel,
+    RBLNAutoModelForAudioClassification,
+    RBLNAutoModelForCausalLM,
+    RBLNAutoModelForCTC,
+    RBLNAutoModelForDepthEstimation,
+    RBLNAutoModelForImageClassification,
+    RBLNAutoModelForMaskedLM,
+    RBLNAutoModelForQuestionAnswering,
+    RBLNAutoModelForSeq2SeqLM,
+    RBLNAutoModelForSequenceClassification,
+    RBLNAutoModelForSpeechSeq2Seq,
+    RBLNAutoModelForVision2Seq,
+)
+from .bart import RBLNBartForConditionalGeneration, RBLNBartModel
+from .bert import RBLNBertModel
+from .clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection, RBLNCLIPVisionModel
 from .dpt import RBLNDPTForDepthEstimation
+from .exaone import RBLNExaoneForCausalLM
 from .gemma import RBLNGemmaForCausalLM
 from .gpt2 import RBLNGPT2LMHeadModel
 from .llama import RBLNLlamaForCausalLM
+from .llava_next import RBLNLlavaNextForConditionalGeneration
 from .midm import RBLNMidmLMHeadModel
 from .mistral import RBLNMistralForCausalLM
+from .phi import RBLNPhiForCausalLM
+from .qwen2 import RBLNQwen2ForCausalLM
+from .t5 import RBLNT5ForConditionalGeneration
 from .wav2vec2 import RBLNWav2Vec2ForCTC
 from .whisper import RBLNWhisperForConditionalGeneration
 from .xlm_roberta import RBLNXLMRobertaModel

optimum/rbln/transformers/models/auto/__init__.py

@@ -0,0 +1,14 @@
+from .modeling_auto import (
+    RBLNAutoModel,
+    RBLNAutoModelForAudioClassification,
+    RBLNAutoModelForCausalLM,
+    RBLNAutoModelForCTC,
+    RBLNAutoModelForDepthEstimation,
+    RBLNAutoModelForImageClassification,
+    RBLNAutoModelForMaskedLM,
+    RBLNAutoModelForQuestionAnswering,
+    RBLNAutoModelForSeq2SeqLM,
+    RBLNAutoModelForSequenceClassification,
+    RBLNAutoModelForSpeechSeq2Seq,
+    RBLNAutoModelForVision2Seq,
+)

optimum/rbln/transformers/models/auto/auto_factory.py

@@ -0,0 +1,84 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+import importlib
+
+from transformers import AutoConfig
+
+
+class _BaseAutoModelClass:
+    # Base class for auto models.
+    _model_mapping = None
+
+    def __init__(self, *args, **kwargs):
+        raise EnvironmentError(
+            f"{self.__class__.__name__} is designed to be instantiated "
+            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or "
+            f"`{self.__class__.__name__}.from_config(config)` methods."
+        )
+
+    @classmethod
+    def get_rbln_cls(
+        cls,
+        model_id,
+        *args,
+        **kwargs,
+    ):
+        # kwargs.update({"return_unused_kwargs": True})
+        config = AutoConfig.from_pretrained(model_id, return_unused_kwargs=True, **kwargs)[0]
+
+        if len(config.architectures) > 1:
+            raise ValueError(
+                f"Model with ID '{model_id}' has multiple architectures defined in the configuration: "
+                f"{config.architectures}. `_BaseAutoModelClass` require exactly one architecture. "
+            )
+
+        architecture_name = config.architectures[0]
+        if architecture_name not in cls._model_mapping.values():
+            raise ValueError(
+                f"The 'RBLN{architecture_name}' architecture is not supported by `{cls.__name__}.from_pretrained()`."
+                "Please use the appropriate class's `from_pretrained()` method to load this model."
+            )
+
+        rbln_class_name = "RBLN" + architecture_name
+        module = importlib.import_module("optimum.rbln")
+
+        try:
+            rbln_cls = getattr(module, rbln_class_name)
+        except AttributeError as e:
+            raise AttributeError(
+                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{model_id}'. "
+                "Ensure that the class name is correctly mapped and available in the 'optimum.rbln' module."
+            ) from e
+
+        return rbln_cls
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id,
+        *args,
+        **kwargs,
+    ):
+        rbln_cls = cls.get_rbln_cls(model_id, *args, **kwargs)
+        return rbln_cls.from_pretrained(model_id, *args, **kwargs)

optimum/rbln/transformers/models/auto/modeling_auto.py

@@ -0,0 +1,95 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
+    MODEL_FOR_CTC_MAPPING_NAMES,
+    MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES,
+    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_MASKED_LM_MAPPING_NAMES,
+    MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES,
+    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
+    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES,
+    MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES,
+    MODEL_MAPPING_NAMES,
+)
+
+from .auto_factory import _BaseAutoModelClass
+
+
+MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.update(
+    {
+        "midm": "MidmLMHeadModel",
+        "exaone": "ExaoneForCausalLM",
+    }
+)
+
+
+class RBLNAutoModel(_BaseAutoModelClass):
+    _model_mapping = MODEL_MAPPING_NAMES
+
+
+class RBLNAutoModelForCTC(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_CTC_MAPPING_NAMES
+
+
+class RBLNAutoModelForCausalLM(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+
+
+class RBLNAutoModelForSeq2SeqLM(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
+
+
+class RBLNAutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
+
+
+class RBLNAutoModelForDepthEstimation(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES
+
+
+class RBLNAutoModelForSequenceClassification(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
+
+
+class RBLNAutoModelForVision2Seq(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
+
+
+class RBLNAutoModelForMaskedLM(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_MASKED_LM_MAPPING_NAMES
+
+
+class RBLNAutoModelForAudioClassification(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
+
+
+class RBLNAutoModelForImageClassification(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
+
+
+class RBLNAutoModelForQuestionAnswering(_BaseAutoModelClass):
+    _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES

optimum/rbln/transformers/models/bart/__init__.py

@@ -22,3 +22,4 @@
 # from Rebellions Inc.
 
 from .bart_architecture import BartDecoderWrapper, BartEncoderWrapper
+from .modeling_bart import RBLNBartForConditionalGeneration, RBLNBartModel