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

optimum/rbln/transformers/models/__init__.py

@@ -21,35 +21,84 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+from typing import TYPE_CHECKING
 
-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
+from transformers.utils import _LazyModule
+
+
+_import_structure = {
+    "auto": [
+        "RBLNAutoModel",
+        "RBLNAutoModelForAudioClassification",
+        "RBLNAutoModelForCausalLM",
+        "RBLNAutoModelForCTC",
+        "RBLNAutoModelForDepthEstimation",
+        "RBLNAutoModelForImageClassification",
+        "RBLNAutoModelForMaskedLM",
+        "RBLNAutoModelForQuestionAnswering",
+        "RBLNAutoModelForSeq2SeqLM",
+        "RBLNAutoModelForSequenceClassification",
+        "RBLNAutoModelForSpeechSeq2Seq",
+        "RBLNAutoModelForVision2Seq",
+    ],
+    "bart": ["RBLNBartForConditionalGeneration", "RBLNBartModel"],
+    "bert": ["RBLNBertModel"],
+    "clip": ["RBLNCLIPTextModel", "RBLNCLIPTextModelWithProjection", "RBLNCLIPVisionModel"],
+    "dpt": ["RBLNDPTForDepthEstimation"],
+    "exaone": ["RBLNExaoneForCausalLM"],
+    "gemma": ["RBLNGemmaForCausalLM"],
+    "gpt2": ["RBLNGPT2LMHeadModel"],
+    "llama": ["RBLNLlamaForCausalLM"],
+    "llava_next": ["RBLNLlavaNextForConditionalGeneration"],
+    "midm": ["RBLNMidmLMHeadModel"],
+    "mistral": ["RBLNMistralForCausalLM"],
+    "phi": ["RBLNPhiForCausalLM"],
+    "qwen2": ["RBLNQwen2ForCausalLM"],
+    "t5": ["RBLNT5EncoderModel", "RBLNT5ForConditionalGeneration"],
+    "wav2vec2": ["RBLNWav2Vec2ForCTC"],
+    "whisper": ["RBLNWhisperForConditionalGeneration"],
+    "xlm_roberta": ["RBLNXLMRobertaModel"],
+}
+
+if TYPE_CHECKING:
+    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 RBLNT5EncoderModel, RBLNT5ForConditionalGeneration
+    from .wav2vec2 import RBLNWav2Vec2ForCTC
+    from .whisper import RBLNWhisperForConditionalGeneration
+    from .xlm_roberta import RBLNXLMRobertaModel
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

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

@@ -22,8 +22,16 @@
 # from Rebellions Inc.
 
 import importlib
+import inspect
+import warnings
 
-from transformers import AutoConfig
+from transformers import AutoConfig, PretrainedConfig
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+from transformers.models.auto.auto_factory import _get_model_class
+
+from optimum.rbln.modeling_base import RBLNBaseModel
+from optimum.rbln.modeling_config import RBLNConfig
+from optimum.rbln.utils.model_utils import convert_hf_to_rbln_model_name, convert_rbln_to_hf_model_name
 
 
 class _BaseAutoModelClass:
@@ -33,46 +41,132 @@ class _BaseAutoModelClass:
     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."
+            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)`"
         )
 
     @classmethod
     def get_rbln_cls(
         cls,
-        model_id,
+        pretrained_model_name_or_path,
         *args,
+        export=True,
         **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")
+        """
+        Determine the appropriate RBLN model class based on the given model ID and configuration.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier or path to the pretrained model.
+            export (bool): Whether to infer the class based on Hugging Face (HF) architecture.
+            kwargs: Additional arguments for configuration and loading.
+
+        Returns:
+            RBLNBaseModel: The corresponding RBLN model class.
+        """
+        if export:
+            hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
+            rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
+        else:
+            rbln_class_name = cls.get_rbln_model_class_name(pretrained_model_name_or_path, **kwargs)
+
+            if convert_rbln_to_hf_model_name(rbln_class_name) not in cls._model_mapping_names.values():
+                raise ValueError(
+                    f"The architecture '{rbln_class_name}' is not supported by the `{cls.__name__}.from_pretrained()` method. "
+                    "Please use the `from_pretrained()` method of the appropriate class to load this model, "
+                    f"or directly use '{rbln_class_name}.from_pretrained()`."
+                )
 
         try:
+            module = importlib.import_module("optimum.rbln")
             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}'. "
+                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
                 "Ensure that the class name is correctly mapped and available in the 'optimum.rbln' module."
             ) from e
 
         return rbln_cls
 
+    @classmethod
+    def infer_hf_model_class(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        **kwargs,
+    ):
+        """
+        Infer the Hugging Face model class based on the configuration or model name.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier or path to the pretrained model.
+            kwargs: Additional arguments for configuration and loading.
+
+        Returns:
+            PretrainedModel: The inferred Hugging Face model class.
+        """
+
+        # Try to load configuration if provided or retrieve it from the model ID
+        config = kwargs.pop("config", None)
+        kwargs.update({"trust_remote_code": True})
+        kwargs["_from_auto"] = True
+
+        # Load configuration if not already provided
+        if not isinstance(config, PretrainedConfig):
+            config, kwargs = AutoConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                return_unused_kwargs=True,
+                **kwargs,
+            )
+
+        # Get hf_model_class from Config
+        has_remote_code = (
+            hasattr(config, "auto_map") and convert_rbln_to_hf_model_name(cls.__name__) in config.auto_map
+        )
+        if has_remote_code:
+            class_ref = config.auto_map[convert_rbln_to_hf_model_name(cls.__name__)]
+            model_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+        elif type(config) in cls._model_mapping.keys():
+            model_class = _get_model_class(config, cls._model_mapping)
+        else:
+            raise ValueError(
+                f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
+                f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
+            )
+
+        if model_class.__name__ != config.architectures[0]:
+            warnings.warn(
+                f"`{cls.__name__}.from_pretrained()` is invoking `{convert_hf_to_rbln_model_name(model_class.__name__)}.from_pretrained()`, which does not match the "
+                f"expected architecture `RBLN{config.architectures[0]}` from config. This mismatch could cause some operations to not be properly loaded "
+                f"from the checkpoint, leading to potential unintended behavior. If this is not intentional, consider calling the "
+                f"`from_pretrained()` method directly from the `RBLN{config.architectures[0]}` class instead.",
+                UserWarning,
+            )
+
+        return model_class
+
+    @classmethod
+    def get_rbln_model_class_name(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Retrieve the path to the compiled model directory for a given RBLN model.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier of the model.
+            kwargs: Additional arguments that match the parameters of `_load_compiled_model_dir`.
+
+        Returns:
+            str: Path to the compiled model directory.
+        """
+        sig = inspect.signature(RBLNBaseModel._load_compiled_model_dir)
+        valid_params = sig.parameters.keys()
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+
+        model_path_subfolder = RBLNBaseModel._load_compiled_model_dir(
+            model_id=pretrained_model_name_or_path, **filtered_kwargs
+        )
+        rbln_config = RBLNConfig.load(model_path_subfolder)
+
+        return rbln_config.meta["cls"]
+
     @classmethod
     def from_pretrained(
         cls,

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

@@ -21,18 +21,31 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+
 from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING,
     MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_CAUSAL_LM_MAPPING,
     MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
+    MODEL_FOR_CTC_MAPPING,
     MODEL_FOR_CTC_MAPPING_NAMES,
+    MODEL_FOR_DEPTH_ESTIMATION_MAPPING,
     MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES,
+    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
     MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_MASKED_LM_MAPPING,
     MODEL_FOR_MASKED_LM_MAPPING_NAMES,
+    MODEL_FOR_QUESTION_ANSWERING_MAPPING,
     MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES,
+    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
     MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
+    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
     MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING,
     MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES,
+    MODEL_FOR_VISION_2_SEQ_MAPPING,
     MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES,
+    MODEL_MAPPING,
     MODEL_MAPPING_NAMES,
 )
 
@@ -48,48 +61,60 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.update(
 
 
 class RBLNAutoModel(_BaseAutoModelClass):
-    _model_mapping = MODEL_MAPPING_NAMES
+    _model_mapping = MODEL_MAPPING
+    _model_mapping_names = MODEL_MAPPING_NAMES
 
 
 class RBLNAutoModelForCTC(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_CTC_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_CTC_MAPPING
+    _model_mapping_names = MODEL_FOR_CTC_MAPPING_NAMES
 
 
 class RBLNAutoModelForCausalLM(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING
+    _model_mapping_names = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForSeq2SeqLM(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
+    _model_mapping_names = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING
+    _model_mapping_names = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
 
 
 class RBLNAutoModelForDepthEstimation(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING
+    _model_mapping_names = MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForSequenceClassification(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
+    _model_mapping_names = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForVision2Seq(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING
+    _model_mapping_names = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
 
 
 class RBLNAutoModelForMaskedLM(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_MASKED_LM_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_MASKED_LM_MAPPING
+    _model_mapping_names = MODEL_FOR_MASKED_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForAudioClassification(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
+    _model_mapping_names = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForImageClassification(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
+    _model_mapping_names = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForQuestionAnswering(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING
+    _model_mapping_names = MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES

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

@@ -24,9 +24,9 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
-from transformers import BartConfig, BartForConditionalGeneration, BartModel, PretrainedConfig
+from transformers import BartForConditionalGeneration, PretrainedConfig
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
@@ -41,9 +41,6 @@ if TYPE_CHECKING:
 
 
 class RBLNBartModel(RBLNModel):
-    original_model_class = BartModel
-    original_config_class = BartConfig
-
     @classmethod
     def _get_rbln_config(
         cls,
@@ -82,7 +79,7 @@ class RBLNBartModel(RBLNModel):
             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:
-                input_names_order = inspect.signature(cls.original_model_class.forward).parameters.keys()
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
                 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 BartModel forward() arguments like ({list(input_names_order)})"

optimum/rbln/transformers/models/bert/modeling_bert.py

@@ -25,9 +25,9 @@ import inspect
 import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-from transformers import BertConfig, BertModel, PretrainedConfig
+from transformers import PretrainedConfig
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
@@ -38,9 +38,6 @@ if TYPE_CHECKING:
 
 
 class RBLNBertModel(RBLNModel):
-    original_model_class = BertModel
-    original_config_class = BertConfig
-
     @classmethod
     def _get_rbln_config(
         cls,
@@ -75,7 +72,7 @@ class RBLNBertModel(RBLNModel):
             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:
-                input_names_order = inspect.signature(cls.original_model_class.forward).parameters.keys()
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
                 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 BertModel forward() arguments like ({list(input_names_order)})"

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -26,19 +26,17 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from transformers import (
-    AutoConfig,
-    AutoModel,
     CLIPTextConfig,
     CLIPTextModel,
-    CLIPTextModelWithProjection,
     CLIPVisionConfig,
     CLIPVisionModel,
 )
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....modeling_diffusers import RBLNDiffusionMixin
 
 
 logger = logging.getLogger(__name__)
@@ -58,24 +56,14 @@ class _TextEncoder(torch.nn.Module):
 
 
 class RBLNCLIPTextModel(RBLNModel):
-    original_model_class = CLIPTextModel
-    original_config_class = CLIPTextConfig
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        AutoConfig.from_pretrained = cls.original_config_class.from_pretrained
-        AutoModel.from_pretrained = cls.original_model_class.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        return rt
-
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
 
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        return rbln_config
+
     @classmethod
     def _get_rbln_config(
         cls,
@@ -119,7 +107,7 @@ class RBLNCLIPTextModel(RBLNModel):
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
-    original_model_class = CLIPTextModelWithProjection
+    pass
 
 
 class _VisionEncoder(torch.nn.Module):
@@ -133,20 +121,6 @@ class _VisionEncoder(torch.nn.Module):
 
 
 class RBLNCLIPVisionModel(RBLNModel):
-    original_model_class = CLIPVisionModel
-    original_config_class = CLIPVisionConfig
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        AutoConfig.from_pretrained = cls.original_config_class.from_pretrained
-        AutoModel.from_pretrained = cls.original_model_class.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        return rt
-
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
@@ -155,7 +129,7 @@ class RBLNCLIPVisionModel(RBLNModel):
     def _get_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "CLIPTextConfig",
+        model_config: "CLIPVisionConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
         rbln_batch_size = rbln_kwargs.get("batch_size", 1)

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

@@ -22,12 +22,7 @@
 # from Rebellions Inc.
 
 from .decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyDecoderLayer,
-    DecoderOnlyModel,
     DecoderOnlyWrapper,
-    DynamicNTKScalingRotaryEmbedding,
-    LinearScalingRotaryEmbedding,
     RotaryEmbedding,
     apply_rotary_pos_emb,
     rotate_half,