optimum-rbln 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

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

@@ -28,16 +28,15 @@ import torch
 from transformers import (
     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 ....diffusers.modeling_diffusers import RBLNDiffusionMixin
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....utils.context import override_auto_classes
 
 
 logger = logging.getLogger(__name__)
@@ -57,20 +56,14 @@ class _TextEncoder(torch.nn.Module):
 
 
 class RBLNCLIPTextModel(RBLNModel):
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(
-            config_func=CLIPTextConfig.from_pretrained,
-            model_func=CLIPTextModel.from_pretrained,
-            skip_taskmanager=False,
-        ):
-            rt = super().from_pretrained(*args, **kwargs)
-        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,
@@ -114,7 +107,7 @@ class RBLNCLIPTextModel(RBLNModel):
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
-    original_model_class = CLIPTextModelWithProjection
+    pass
 
 
 class _VisionEncoder(torch.nn.Module):
@@ -128,16 +121,6 @@ class _VisionEncoder(torch.nn.Module):
 
 
 class RBLNCLIPVisionModel(RBLNModel):
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(
-            config_func=CLIPVisionConfig.from_pretrained,
-            model_func=CLIPVisionModel.from_pretrained,
-            skip_taskmanager=False,
-        ):
-            rt = super().from_pretrained(*args, **kwargs)
-        return rt
-
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
@@ -146,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

@@ -21,14 +21,4 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyDecoderLayer,
-    DecoderOnlyModel,
-    DecoderOnlyWrapper,
-    RotaryEmbedding,
-    apply_rotary_pos_emb,
-    rotate_half,
-    slice_and_unsqueeze_cos_sin,
-)
 from .modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM