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

optimum/rbln/modeling_base.py

@@ -21,7 +21,6 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-import copy
 import importlib
 import inspect
 import logging
@@ -50,9 +49,9 @@ from transformers import (
     PretrainedConfig,
 )
 
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig, use_rbln_config
 from .utils.runtime_utils import UnavailableRuntime
-from .utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
+from .utils.save_utils import maybe_load_preprocessors
 
 
 if TYPE_CHECKING:
@@ -86,13 +85,11 @@ class SubModulesMixin:
             setattr(self, submodule_meta["name"], submodule)
 
     @classmethod
-    def _from_model(
+    def _export_submodules_from_model(
         cls,
         model: "PreTrainedModel",
         model_save_dir: str,
-        rbln_sub_configs_dict: Dict[str, Any],
         rbln_kwargs: Dict[str, Any],
-        subfolder=None,  # warning: will be ignored
         **kwargs,
     ) -> List["RBLNBaseModel"]:
         rbln_submodules = []
@@ -100,18 +97,15 @@ class SubModulesMixin:
             submodule_name = submodule["name"]
             torch_submodule: "PreTrainedModel" = getattr(model, submodule["name"])
             cls_name = torch_submodule.__class__.__name__
-            submodule_cls: "RBLNBaseModel" = getattr(importlib.import_module("optimum.rbln"), f"RBLN{cls_name}")
+            submodule_cls: "RBLNModel" = getattr(importlib.import_module("optimum.rbln"), f"RBLN{cls_name}")
 
-            if submodule_name in rbln_sub_configs_dict:
-                kwargs["rbln_config"] = rbln_sub_configs_dict[submodule_name]
+            if submodule_name in rbln_kwargs:
+                kwargs["rbln_config"] = rbln_kwargs[submodule_name]
 
-            rbln_submodule = submodule_cls._export(
-                model_id=None,
-                config=torch_submodule.config,
+            rbln_submodule = submodule_cls.from_model(
+                model=torch_submodule,
                 subfolder=submodule_name,
                 model_save_dir=model_save_dir,
-                model=torch_submodule,
-                **rbln_kwargs,
                 **kwargs,
             )
 
@@ -120,26 +114,29 @@ class SubModulesMixin:
         return rbln_submodules
 
     @classmethod
-    def _submodule_from_compiled_model(
-        cls, model_save_dir: str, rbln_sub_configs_dict: Dict[str, Any], rbln_kwargs: Dict[str, Any], **kwargs
+    def _load_submodules_from_compiled_models(
+        cls,
+        model_save_dir: str,
+        rbln_kwargs: Dict[str, Any],
+        **kwargs,
     ):
         rbln_submodules = []
         for submodule in cls._rbln_submodules:
             submodule_name = submodule["name"]
-            rbln_submodule_config_dict = rbln_sub_configs_dict.get(submodule_name, None)
+
+            if submodule_name in rbln_kwargs:
+                kwargs["rbln_config"] = rbln_kwargs[submodule_name]
 
             # Get cls name for call the constructor of the rbln class
             submodule_rbln_config = RBLNConfig.load(Path(model_save_dir) / submodule_name)
             submodule_cls_name = submodule_rbln_config.meta["cls"]
             submodule_cls: "RBLNBaseModel" = getattr(importlib.import_module("optimum.rbln"), submodule_cls_name)
 
-            config = OptimizedModel._load_config(Path(model_save_dir) / submodule_name, **kwargs)
+            config = OptimizedModel._load_config(Path(model_save_dir) / submodule_name)
             rbln_submodule = submodule_cls._from_pretrained(
                 model_id=model_save_dir,
                 config=config,
                 subfolder=submodule_name,
-                rbln_config=rbln_submodule_config_dict,
-                **rbln_kwargs,
                 **kwargs,
             )
             rbln_submodules.append(rbln_submodule)
@@ -149,27 +146,24 @@ class SubModulesMixin:
     def _load_submodules(
         cls,
         model_save_dir,
-        rbln_sub_configs_dict,
         rbln_kwargs,
         model=None,
         **kwargs,
     ):
-        # Two way :
+        # Two ways :
         # 1. Compile from pytorch object
         # 2. Load from compiled file
         if model is not None:
-            return cls._from_model(
+            return cls._export_submodules_from_model(
                 model=model,
                 model_save_dir=model_save_dir,
-                rbln_sub_configs_dict=rbln_sub_configs_dict,
                 rbln_kwargs=rbln_kwargs,
                 **kwargs,
             )
 
         else:
-            return cls._submodule_from_compiled_model(
+            return cls._load_submodules_from_compiled_models(
                 model_save_dir=model_save_dir,
-                rbln_sub_configs_dict=rbln_sub_configs_dict,
                 rbln_kwargs=rbln_kwargs,
                 **kwargs,
             )
@@ -370,45 +364,7 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         return rbln_compiled_models
 
     @classmethod
-    def _split_submodule_config(cls, rbln_config_dict: Dict[str, Any] = {}) -> Dict[str, Any]:
-        # {"language_model" : {"rbln_tensor_parallel_size":4}}
-        rbln_sub_configs_dict: Dict[str, Dict[str, Any]] = {}
-
-        # Remove submodule-configs from rbln_config
-        if len(cls._rbln_submodules) > 0:
-            keys = list(rbln_config_dict.keys())
-            submodule_names = [m["name"] for m in cls._rbln_submodules]
-            for key in keys:
-                if key in submodule_names:
-                    rbln_sub_configs_dict[key] = rbln_config_dict.pop(key)
-
-        return rbln_sub_configs_dict
-
-    @classmethod
-    def resolve_rbln_config(cls, rbln_config: Union[RBLNConfig, Dict[str, Any]], kwargs):
-        if isinstance(rbln_config, RBLNConfig):
-            # Already resolved
-            return rbln_config, None
-
-        else:
-            if rbln_config is None:
-                rbln_config_dict = {}
-            else:
-                rbln_config_dict = rbln_config
-
-            rbln_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-            rbln_sub_configs_dict = cls._split_submodule_config(rbln_config_dict)
-
-            for key in rbln_config_dict:
-                if key in rbln_kwargs:
-                    raise KeyError(f"duplicate key in both `rbln_config` and {key}")
-
-            merged_rbln_kwargs = copy.deepcopy(rbln_kwargs)
-            merged_rbln_kwargs.update(rbln_config_dict)
-
-            return (merged_rbln_kwargs, rbln_sub_configs_dict)
-
-    @classmethod
+    @use_rbln_config
     def _from_pretrained(
         cls,
         model_id: Union[str, Path],
@@ -430,7 +386,7 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
 
         if not from_export_method:
             # from compiled dir
-            rbln_kwargs, rbln_sub_configs_dict = cls.resolve_rbln_config(rbln_config, kwargs)
+            rbln_kwargs = rbln_config or {}
 
             model_path_subfolder = cls._load_compiled_model_dir(
                 model_id=model_id,
@@ -450,7 +406,6 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
             if len(cls._rbln_submodules) > 0:
                 rbln_submodules = cls._load_submodules(
                     model_save_dir=model_id,
-                    rbln_sub_configs_dict=rbln_sub_configs_dict,
                     rbln_kwargs=rbln_kwargs,
                     **kwargs,
                 )
@@ -477,9 +432,9 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         cls,
         rbln_compiled_models: Dict[str, rebel.RBLNCompiledModel],
         rbln_config: RBLNConfig,
-        config,
-        model_save_dir: str,
-        subfolder: str,
+        config: "PretrainedConfig",
+        model_save_dir: Union[Path, str],
+        subfolder: Union[Path, str],
         rbln_submodules: List["RBLNBaseModel"] = [],
         **kwargs,
     ):
@@ -538,25 +493,10 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         rbln_config = cls._get_rbln_config(**others, rbln_kwargs=rbln_kwargs)
         return rbln_config
 
-    @staticmethod
-    def pop_rbln_kwargs_from_kwargs(kwargs: Dict[str, Any], runtime_only=False):
-        keys = list(kwargs.keys())
-        rbln_kwargs = {key[5:]: kwargs.pop(key) for key in keys if key.startswith("rbln_")}
-
-        if runtime_only:
-            rbln_kwargs = {
-                key: value
-                for key, value in rbln_kwargs.items()
-                if key in {"create_runtimes", "optimize_host_memory", "device", "device_map"}
-            }
-
-        return rbln_kwargs
-
     def can_generate(self):
         return False
 
     def to(self, *args, **kwargs):
-        # Do nothing
         return self
 
     def __call__(self, *args, **kwargs):
@@ -599,12 +539,49 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
 
     @classmethod
     @abstractmethod
-    def _export(cls, *args, **kwargs):
-        """
-        Exports a vanilla Transformers model into a rbln-compiled Module.
-        """
+    def get_pytorch_model(cls, *args, **kwargs):
         pass
 
+    @classmethod
+    @abstractmethod
+    @use_rbln_config
+    def from_model(
+        cls,
+        model: "PreTrainedModel",
+        rbln_config: Dict[str, Any] = {},
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        subfolder: str = "",
+        **kwargs,
+    ):
+        pass
+
+    @classmethod
+    @use_rbln_config
+    def _export(
+        cls,
+        model_id: Union[str, Path],
+        config: "PretrainedConfig",  # FIXME : optimum passes config, but we ignore it.
+        rbln_config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> "RBLNModel":
+        subfolder = kwargs.get("subfolder", "")
+        model_save_dir = kwargs.pop("model_save_dir", None)
+
+        rbln_kwargs = rbln_config
+        model: "PreTrainedModel" = cls.get_pytorch_model(
+            model_id=model_id,
+            rbln_kwargs=rbln_kwargs,
+            **kwargs,
+        )
+        preprocessors = maybe_load_preprocessors(model_id, subfolder=subfolder)
+        return cls.from_model(
+            model,
+            rbln_config=rbln_config,
+            preprocessors=preprocessors,
+            model_save_dir=model_save_dir,
+            **kwargs,
+        )
+
 
 class RBLNModel(RBLNBaseModel):
     """
@@ -695,25 +672,19 @@ class RBLNModel(RBLNBaseModel):
         return compiled_model
 
     @classmethod
-    @torch.no_grad()
-    def _export(
+    @use_rbln_config
+    def from_model(
         cls,
-        model_id: Union[str, Path],
-        config: "PretrainedConfig",
-        use_auth_token: Optional[Union[bool, str]] = None,
-        revision: Optional[str] = None,
-        force_download: bool = False,
-        cache_dir: Optional[str] = None,
-        subfolder: str = "",
-        local_files_only: bool = False,
-        trust_remote_code: bool = False,
+        model: "PreTrainedModel",
+        rbln_config: Dict[str, Any] = {},
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        model: "PreTrainedModel" = None,
-        rbln_config: Optional[Dict[str, Any]] = None,
+        subfolder: str = "",
         **kwargs,
-    ) -> "RBLNModel":
-        rbln_kwargs, rbln_sub_configs_dict = cls.resolve_rbln_config(rbln_config, kwargs)
+    ):
+        preprocessors = kwargs.pop("preprocessors", [])
+        rbln_kwargs = rbln_config
 
+        # Directory to save compile artifacts(.rbln) and original configs
         if model_save_dir is None:
             save_dir = TemporaryDirectory()
             save_dir_path = Path(save_dir.name)
@@ -725,43 +696,31 @@ class RBLNModel(RBLNBaseModel):
                 save_dir_path = Path(model_save_dir)
                 save_dir_path.mkdir(exist_ok=True)
 
-        # Load pytorch model if needed.
-        if model is None:
-            model: "PreTrainedModel" = cls.get_pytorch_model(
-                model_id=model_id,
-                subfolder=subfolder,
-                revision=revision,
-                cache_dir=cache_dir,
-                use_auth_token=use_auth_token,
-                local_files_only=local_files_only,
-                force_download=force_download,
-                trust_remote_code=trust_remote_code,
-                rbln_kwargs=rbln_kwargs,
-                **kwargs,
-            )
-            preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
-        else:
-            preprocessors = []
+        # (Optional) Save preprocessors (tokenizer, image preprocessors, etc)
+        for preprocessor in preprocessors:
+            preprocessor.save_pretrained(save_dir_path)
 
-        # FIXME :: optimum passes AutoConfig.
+        # Save configs
+        # FIXME :: optimum passes AutoConfig. But here we ignore it.
         config = model.config
         if hasattr(model, "can_generate") and model.can_generate():
             generation_config = model.generation_config
             generation_config.save_pretrained(save_dir_path / subfolder)
-
         if not isinstance(config, PretrainedConfig):  # diffusers config
             config = PretrainedConfig(**config)
         config.save_pretrained(save_dir_path / subfolder)
 
-        # Get compilation arguments
+        # Get compilation arguments (e.g. input_info)
         rbln_config: RBLNConfig = cls.get_rbln_config(
             preprocessors=preprocessors, model_config=config, rbln_kwargs=rbln_kwargs
         )
+        # rbln_config.update_runtime_cfg(rbln_kwargs) # This is done in get_rbln_config
+
         compiled_model: Union[rebel.RBLNCompiledModel, Dict[str, rebel.RBLNCompiledModel]] = cls.get_compiled_model(
             model, rbln_config=rbln_config
         )
 
-        # Save compiled models
+        # Save compiled models (.rbln)
         (save_dir_path / subfolder).mkdir(exist_ok=True)
         if not isinstance(compiled_model, dict):
             compiled_models = {DEFAULT_COMPILED_MODEL_NAME: compiled_model}
@@ -771,6 +730,7 @@ class RBLNModel(RBLNBaseModel):
             cm.save(save_dir_path / subfolder / f"{compiled_model_name}.rbln")
         rbln_config.save(save_dir_path / subfolder)
 
+        # Save torch artifacts (e.g. embedding matrix if needed.)
         cls.save_torch_artifacts(model, save_dir_path=save_dir_path, subfolder=subfolder, rbln_config=rbln_config)
 
         # Load submodules
@@ -778,7 +738,6 @@ class RBLNModel(RBLNBaseModel):
             rbln_submodules = cls._load_submodules(
                 model=model,
                 model_save_dir=save_dir,
-                rbln_sub_configs_dict=rbln_sub_configs_dict,
                 rbln_kwargs=rbln_kwargs,
                 **kwargs,
             )
@@ -790,12 +749,7 @@ class RBLNModel(RBLNBaseModel):
             model_id=save_dir_path,
             config=config,
             model_save_dir=save_dir,
-            use_auth_token=use_auth_token,
-            revision=revision,
-            force_download=force_download,
-            cache_dir=cache_dir,
             subfolder=subfolder,
-            local_files_only=local_files_only,
             rbln_config=rbln_config,
             rbln_compiled_models=compiled_models,
             rbln_submodules=rbln_submodules,
@@ -817,7 +771,6 @@ class RBLNModel(RBLNBaseModel):
 
 
 class RBLNModelForQuestionAnswering(RBLNModel):
-    model_type = "rbln_model"
     auto_model_class = AutoModelForQuestionAnswering
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
 
@@ -878,7 +831,6 @@ 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
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForImageClassification
 
     @classmethod
@@ -894,8 +846,17 @@ class RBLNModelForImageClassification(RBLNModel):
         if rbln_image_size is None:
             for processor in preprocessors:
                 if hasattr(processor, "size"):
-                    rbln_image_size = processor.size["shortest_edge"]
+                    if all(required_key in processor.size.keys() for required_key in ["height", "width"]):
+                        rbln_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"])
+                    elif "longest_edge" in processor.size.keys():
+                        rbln_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!")
 
@@ -903,12 +864,20 @@ class RBLNModelForImageClassification(RBLNModel):
             rbln_batch_size = 1
 
         if isinstance(rbln_image_size, int):
-            rbln_image_size = rbln_image_size, rbln_image_size
+            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"
+            )
 
         input_info = [
             (
                 "pixel_values",
-                [rbln_batch_size, 3, rbln_image_size[0], rbln_image_size[1]],
+                [rbln_batch_size, 3, rbln_image_height, rbln_image_width],
                 "float32",
             )
         ]
@@ -930,7 +899,6 @@ class RBLNModelForAudioClassification(RBLNModel):
     Currently, this model class only supports the 'AST' model from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForAudioClassification
 
     @classmethod
@@ -1005,7 +973,6 @@ class RBLNModelForSequenceClassification(RBLNModel):
     Currently, this model class supports the 'XLMRoberta' and 'Roberta' model from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForSequenceClassification
 
     @classmethod
@@ -1070,7 +1037,6 @@ class RBLNModelForSequenceClassification(RBLNModel):
 
 
 class RBLNModelForMaskedLM(RBLNModel):
-    model_type = "rbln_model"
     auto_model_class = AutoModelForMaskedLM
 
     @classmethod

optimum/rbln/modeling_config.py

@@ -242,3 +242,54 @@ class RBLNConfig:
                 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 = {}
+
+        else:
+            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

@@ -42,18 +42,22 @@ _import_structure = {
         "RBLNAutoModelForSequenceClassification",
         "RBLNAutoModelForSpeechSeq2Seq",
         "RBLNAutoModelForVision2Seq",
+        "RBLNBartForConditionalGeneration",
         "RBLNBartModel",
         "RBLNBertModel",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
         "RBLNCLIPVisionModel",
         "RBLNDPTForDepthEstimation",
+        "RBLNExaoneForCausalLM",
         "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
+        "RBLNQwen2ForCausalLM",
         "RBLNWav2Vec2ForCTC",
         "RBLNWhisperForConditionalGeneration",
         "RBLNLlamaForCausalLM",
         "RBLNPhiForCausalLM",
+        "RBLNT5ForConditionalGeneration",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
         "RBLNXLMRobertaModel",
@@ -77,12 +81,14 @@ if TYPE_CHECKING:
         RBLNAutoModelForSequenceClassification,
         RBLNAutoModelForSpeechSeq2Seq,
         RBLNAutoModelForVision2Seq,
+        RBLNBartForConditionalGeneration,
         RBLNBartModel,
         RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNCLIPVisionModel,
         RBLNDPTForDepthEstimation,
+        RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
@@ -90,6 +96,8 @@ if TYPE_CHECKING:
         RBLNMidmLMHeadModel,
         RBLNMistralForCausalLM,
         RBLNPhiForCausalLM,
+        RBLNQwen2ForCausalLM,
+        RBLNT5ForConditionalGeneration,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/transformers/models/__init__.py

@@ -36,10 +36,11 @@ from .auto import (
     RBLNAutoModelForSpeechSeq2Seq,
     RBLNAutoModelForVision2Seq,
 )
-from .bart import RBLNBartModel
+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
@@ -47,6 +48,8 @@ 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/modeling_auto.py

@@ -42,6 +42,7 @@ from .auto_factory import _BaseAutoModelClass
 MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.update(
     {
         "midm": "MidmLMHeadModel",
+        "exaone": "ExaoneForCausalLM",
     }
 )
 

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

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