optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

optimum/rbln/transformers/models/__init__.py

@@ -18,6 +18,10 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
+    "audio_spectrogram_transformer": [
+        "RBLNASTForAudioClassification",
+        "RBLNASTForAudioClassificationConfig",
+    ],
     "auto": [
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
@@ -65,6 +69,14 @@ _import_structure = {
         "RBLNCLIPVisionModelWithProjection",
         "RBLNCLIPVisionModelWithProjectionConfig",
     ],
+    "colpali": [
+        "RBLNColPaliForRetrieval",
+        "RBLNColPaliForRetrievalConfig",
+    ],
+    "distilbert": [
+        "RBLNDistilBertForQuestionAnswering",
+        "RBLNDistilBertForQuestionAnsweringConfig",
+    ],
     "qwen2_5_vl": [
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
@@ -101,11 +113,18 @@ _import_structure = {
     "mistral": ["RBLNMistralForCausalLM", "RBLNMistralForCausalLMConfig"],
     "phi": ["RBLNPhiForCausalLM", "RBLNPhiForCausalLMConfig"],
     "qwen2": ["RBLNQwen2ForCausalLM", "RBLNQwen2ForCausalLMConfig"],
+    "resnet": ["RBLNResNetForImageClassification", "RBLNResNetForImageClassificationConfig"],
+    "roberta": [
+        "RBLNRobertaForMaskedLM",
+        "RBLNRobertaForMaskedLMConfig",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassificationConfig",
+    ],
     "siglip": [
         "RBLNSiglipVisionModel",
         "RBLNSiglipVisionModelConfig",
     ],
-    "time_series_transformers": [
+    "time_series_transformer": [
         "RBLNTimeSeriesTransformerForPrediction",
         "RBLNTimeSeriesTransformerForPredictionConfig",
     ],
@@ -115,12 +134,22 @@ _import_structure = {
         "RBLNT5EncoderModelConfig",
         "RBLNT5ForConditionalGenerationConfig",
     ],
+    "vit": ["RBLNViTForImageClassification", "RBLNViTForImageClassificationConfig"],
     "wav2vec2": ["RBLNWav2Vec2ForCTC", "RBLNWav2Vec2ForCTCConfig"],
     "whisper": ["RBLNWhisperForConditionalGeneration", "RBLNWhisperForConditionalGenerationConfig"],
-    "xlm_roberta": ["RBLNXLMRobertaModel", "RBLNXLMRobertaModelConfig"],
+    "xlm_roberta": [
+        "RBLNXLMRobertaModel",
+        "RBLNXLMRobertaModelConfig",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNXLMRobertaForSequenceClassificationConfig",
+    ],
 }
 
 if TYPE_CHECKING:
+    from .audio_spectrogram_transformer import (
+        RBLNASTForAudioClassification,
+        RBLNASTForAudioClassificationConfig,
+    )
     from .auto import (
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
@@ -168,10 +197,18 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjection,
         RBLNCLIPVisionModelWithProjectionConfig,
     )
+    from .colpali import (
+        RBLNColPaliForRetrieval,
+        RBLNColPaliForRetrievalConfig,
+    )
     from .decoderonly import (
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
     )
+    from .distilbert import (
+        RBLNDistilBertForQuestionAnswering,
+        RBLNDistilBertForQuestionAnsweringConfig,
+    )
     from .dpt import (
         RBLNDPTForDepthEstimation,
         RBLNDPTForDepthEstimationConfig,
@@ -204,6 +241,13 @@ if TYPE_CHECKING:
         RBLNQwen2_5_VLForConditionalGeneration,
         RBLNQwen2_5_VLForConditionalGenerationConfig,
     )
+    from .resnet import RBLNResNetForImageClassification, RBLNResNetForImageClassificationConfig
+    from .roberta import (
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForMaskedLMConfig,
+        RBLNRobertaForSequenceClassification,
+        RBLNRobertaForSequenceClassificationConfig,
+    )
     from .siglip import RBLNSiglipVisionModel, RBLNSiglipVisionModelConfig
     from .t5 import (
         RBLNT5EncoderModel,
@@ -211,13 +255,19 @@ if TYPE_CHECKING:
         RBLNT5ForConditionalGeneration,
         RBLNT5ForConditionalGenerationConfig,
     )
-    from .time_series_transformers import (
+    from .time_series_transformer import (
         RBLNTimeSeriesTransformerForPrediction,
         RBLNTimeSeriesTransformerForPredictionConfig,
     )
+    from .vit import RBLNViTForImageClassification, RBLNViTForImageClassificationConfig
     from .wav2vec2 import RBLNWav2Vec2ForCTC, RBLNWav2Vec2ForCTCConfig
     from .whisper import RBLNWhisperForConditionalGeneration, RBLNWhisperForConditionalGenerationConfig
-    from .xlm_roberta import RBLNXLMRobertaModel, RBLNXLMRobertaModelConfig
+    from .xlm_roberta import (
+        RBLNXLMRobertaForSequenceClassification,
+        RBLNXLMRobertaForSequenceClassificationConfig,
+        RBLNXLMRobertaModel,
+        RBLNXLMRobertaModelConfig,
+    )
 
 else:
     import sys

optimum/rbln/transformers/models/{wav2vec2/configuration_wav2vec.py → audio_spectrogram_transformer/__init__.py}

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...configuration_generic import RBLNModelForMaskedLMConfig
 
-
-class RBLNWav2Vec2ForCTCConfig(RBLNModelForMaskedLMConfig):
-    rbln_model_input_names = ["input_values"]
+from .configuration_audio_spectrogram_transformer import RBLNASTForAudioClassificationConfig
+from .modeling_audio_spectrogram_transformer import RBLNASTForAudioClassification

optimum/rbln/transformers/models/audio_spectrogram_transformer/configuration_audio_spectrogram_transformer.py

@@ -0,0 +1,21 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# 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.
+
+from ...configuration_generic import RBLNModelForAudioClassificationConfig
+
+
+class RBLNASTForAudioClassificationConfig(RBLNModelForAudioClassificationConfig):
+    """
+    Configuration class for RBLNASTForAudioClassification.
+    """

optimum/rbln/transformers/models/audio_spectrogram_transformer/modeling_audio_spectrogram_transformer.py

@@ -0,0 +1,28 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# 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.
+
+from ...modeling_generic import RBLNModelForAudioClassification
+
+
+class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
+    """
+    Audio Spectrogram Transformer model with an audio classification head on top (a linear layer on top of the pooled output) e.g. for datasets like AudioSet, Speech Commands v2.
+    This model inherits from [`RBLNModelForAudioClassification`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformer-based `ASTForAudioClassification` models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `ASTForAudioClassification` 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.
+    """

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

@@ -11,10 +11,10 @@
 # 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.
-
 import importlib
 import inspect
 import warnings
+from typing import Type
 
 from transformers import AutoConfig, PretrainedConfig
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
@@ -22,7 +22,12 @@ from transformers.models.auto.auto_factory import _get_model_class
 
 from optimum.rbln.configuration_utils import RBLNAutoConfig
 from optimum.rbln.modeling_base import RBLNBaseModel
-from optimum.rbln.utils.model_utils import convert_hf_to_rbln_model_name, convert_rbln_to_hf_model_name
+from optimum.rbln.utils.model_utils import (
+    MODEL_MAPPING,
+    convert_hf_to_rbln_model_name,
+    convert_rbln_to_hf_model_name,
+    get_rbln_model_cls,
+)
 
 
 class _BaseAutoModelClass:
@@ -58,7 +63,7 @@ class _BaseAutoModelClass:
             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)
+            rbln_class_name = cls.get_rbln_model_cls_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(
@@ -68,8 +73,7 @@ class _BaseAutoModelClass:
                 )
 
         try:
-            module = importlib.import_module("optimum.rbln")
-            rbln_cls = getattr(module, rbln_class_name)
+            rbln_cls = get_rbln_model_cls(rbln_class_name)
         except AttributeError as e:
             raise AttributeError(
                 f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
@@ -136,7 +140,7 @@ class _BaseAutoModelClass:
         return model_class
 
     @classmethod
-    def get_rbln_model_class_name(cls, pretrained_model_name_or_path, **kwargs):
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path, **kwargs):
         """
         Retrieve the path to the compiled model directory for a given RBLN model.
 
@@ -159,11 +163,30 @@ class _BaseAutoModelClass:
         return rbln_config.rbln_model_cls_name
 
     @classmethod
-    def from_pretrained(
-        cls,
-        model_id,
-        *args,
-        **kwargs,
-    ):
+    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)
+
+    @classmethod
+    def from_model(cls, model, *args, **kwargs):
+        rbln_cls = get_rbln_model_cls(f"RBLN{model.__class__.__name__}")
+        return rbln_cls.from_model(model, *args, **kwargs)
+
+    @staticmethod
+    def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
+        """
+        Register a new RBLN model class.
+
+        Args:
+            rbln_cls (Type[RBLNBaseModel]): The RBLN model class to register.
+            exist_ok (bool): Whether to allow registering an already registered model.
+        """
+        if not issubclass(rbln_cls, RBLNBaseModel):
+            raise ValueError("`rbln_cls` must be a subclass of RBLNBaseModel.")
+
+        native_cls = getattr(importlib.import_module("optimum.rbln"), rbln_cls.__name__, None)
+        if rbln_cls.__name__ in MODEL_MAPPING or native_cls is not None:
+            if not exist_ok:
+                raise ValueError(f"Model for {rbln_cls.__name__} already registered.")
+
+        MODEL_MAPPING[rbln_cls.__name__] = rbln_cls

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

@@ -22,6 +22,7 @@ from transformers.modeling_attn_mask_utils import (
 from transformers.utils import logging
 
 from ..seq2seq.seq2seq_architecture import (
+    Seq2SeqCrossAttention,
     Seq2SeqDecoder,
     Seq2SeqDecoderLayer,
     Seq2SeqDecoderWrapper,
@@ -45,7 +46,8 @@ class BartDecoderWrapper(Seq2SeqDecoderWrapper):
         new_layers = []
         for layer in model.get_decoder().layers:
             self_attn = BartSelfAttention(layer.self_attn, use_attention_mask=self.use_attention_mask)
-            new_layers.append(BartDecoderLayer(layer, self_attn))
+            cross_attn = BartCrossAttention(layer.encoder_attn)
+            new_layers.append(BartDecoderLayer(layer, self_attn, cross_attn))
 
         decoder_model = BartDecoder(model.get_decoder(), new_layers)
         new_model = BartForConditionalGeneration(model, decoder_model)
@@ -153,3 +155,14 @@ class BartSelfAttention(Seq2SeqSelfAttention):
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
         return query_states, key_states, value_states
+
+
+class BartCrossAttention(Seq2SeqCrossAttention):
+    def __post_init__(self):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.out_proj = self._original_mod.out_proj
+        self.num_heads = self._original_mod.num_heads
+        self.head_dim = self._original_mod.embed_dim // self._original_mod.num_heads
+        self.embed_dim = self._original_mod.embed_dim

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

@@ -17,8 +17,18 @@ from ..seq2seq import RBLNModelForSeq2SeqLMConfig
 
 
 class RBLNBartModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
-    pass
+    """
+    Configuration class for RBLNBartModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BART models for feature extraction tasks.
+    """
 
 
 class RBLNBartForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
-    pass
+    """
+    Configuration class for RBLNBartForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BART models for conditional text generation tasks.
+    """

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable
+from typing import Any, Callable
 
 from transformers import BartForConditionalGeneration, PreTrainedModel
 
@@ -27,19 +27,28 @@ from .configuration_bart import RBLNBartForConditionalGenerationConfig
 logger = get_logger()
 
 
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-
 class RBLNBartModel(RBLNTransformerEncoderForFeatureExtraction):
-    pass
+    """
+    RBLN optimized BART model for feature extraction tasks.
+
+    This class provides hardware-accelerated inference for BART encoder models
+    on RBLN devices, optimized for feature extraction use cases.
+    """
 
 
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    """
+    RBLN optimized BART model for conditional text generation tasks.
+
+    This class provides hardware-accelerated inference for BART models
+    on RBLN devices, supporting sequence-to-sequence generation tasks
+    such as summarization, translation, and text generation.
+    """
+
     support_causal_attn = True
 
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNBartForConditionalGenerationConfig):
+    def wrap_model_if_needed(self, model: PreTrainedModel, rbln_config: RBLNBartForConditionalGenerationConfig):
         return BartWrapper(
             model, enc_max_seq_len=rbln_config.enc_max_seq_len, use_attention_mask=rbln_config.use_attention_mask
         )

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

@@ -20,12 +20,27 @@ from ...configuration_generic import (
 
 
 class RBLNBertModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
-    pass
+    """
+    Configuration class for RBLNBertModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BERT models for feature extraction tasks.
+    """
 
 
 class RBLNBertForMaskedLMConfig(RBLNModelForMaskedLMConfig):
-    pass
+    """
+    Configuration class for RBLNBertForMaskedLM.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BERT models for masked language modeling tasks.
+    """
 
 
 class RBLNBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
-    pass
+    """
+    Configuration class for RBLNBertForQuestionAnswering.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BERT models for question answering tasks.
+    """

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

@@ -24,12 +24,36 @@ logger = get_logger(__name__)
 
 
 class RBLNBertModel(RBLNTransformerEncoderForFeatureExtraction):
+    """
+    RBLN optimized BERT model for feature extraction tasks.
+
+    This class provides hardware-accelerated inference for BERT models
+    on RBLN devices, optimized for extracting contextualized embeddings
+    and features from text sequences.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNBertForMaskedLM(RBLNModelForMaskedLM):
+    """
+    RBLN optimized BERT model for masked language modeling tasks.
+
+    This class provides hardware-accelerated inference for BERT models
+    on RBLN devices, supporting masked language modeling tasks such as
+    token prediction and text completion.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
 
 
 class RBLNBertForQuestionAnswering(RBLNModelForQuestionAnswering):
+    """
+    RBLN optimized BERT model for question answering tasks.
+
+    This class provides hardware-accelerated inference for BERT models
+    on RBLN devices, supporting extractive question answering tasks where
+    the model predicts start and end positions of answers in text.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]

optimum/rbln/transformers/models/blip_2/configuration_blip_2.py

@@ -12,16 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNBlip2VisionModelConfig(RBLNModelConfig):
-    pass
+    """
+    Configuration class for RBLNBlip2VisionModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BLIP-2 vision encoder models for multimodal tasks.
+    """
 
 
 class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNBlip2QFormerModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BLIP-2 Q-Former models that bridge vision and language modalities.
+    """
+
     def __init__(
         self,
         num_query_tokens: Optional[int] = None,
@@ -50,7 +62,7 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         vision_model: Optional[RBLNModelConfig] = None,
         qformer: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/blip_2/modeling_blip_2.py

@@ -65,6 +65,13 @@ class LoopProjector:
 
 
 class RBLNBlip2VisionModel(RBLNModel):
+    """
+    RBLN optimized BLIP-2 vision encoder model.
+
+    This class provides hardware-accelerated inference for BLIP-2 vision encoders
+    on RBLN devices, supporting image encoding for multimodal vision-language tasks.
+    """
+
     def get_input_embeddings(self):
         return self.embeddings
 
@@ -136,6 +143,14 @@ class RBLNBlip2VisionModel(RBLNModel):
 
 
 class RBLNBlip2QFormerModel(RBLNModel):
+    """
+    RBLN optimized BLIP-2 Q-Former model.
+
+    This class provides hardware-accelerated inference for BLIP-2 Q-Former models
+    on RBLN devices, which bridge vision and language modalities through cross-attention
+    mechanisms for multimodal understanding tasks.
+    """
+
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
 
@@ -251,6 +266,38 @@ class RBLNBlip2QFormerModel(RBLNModel):
 
 
 class RBLNBlip2ForConditionalGeneration(RBLNModel):
+    """
+    RBLNBlip2ForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
+    optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
+
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model includes a Large Language Model (LLM) as a submodule. For optimal performance, it is highly recommended to use
+        tensor parallelism for the language model.  This can be achieved by using the `rbln_config` parameter in the
+        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNBlip2ForConditionalGeneration class for details.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNBlip2ForConditionalGeneration
+
+        model = RBLNBlip2ForConditionalGeneration.from_pretrained(
+            "Salesforce/blip2-opt-2.7b",
+            export=True,
+            rbln_config={
+                "language_model": {
+                    "batch_size": 1,
+                    "max_seq_len": 2048,
+                    "tensor_parallel_size": 1,
+                    "use_inputs_embeds": True,
+                },
+            },
+        )
+
+        model.save_pretrained("compiled-blip2-opt-2.7b")
+        ```
+    """
+
     auto_model_class = AutoModelForVisualQuestionAnswering
     _rbln_submodules = [{"name": "vision_model"}, {"name": "qformer"}, {"name": "language_model"}]
 
@@ -275,10 +322,9 @@ class RBLNBlip2ForConditionalGeneration(RBLNModel):
         subfolder: str,
         rbln_config: RBLNModelConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
+
         save_dict = {}
         save_dict["query_tokens"] = model.query_tokens
         torch.save(save_dict, save_dir_path / subfolder / "query_tokens.pth")

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

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNCLIPTextModelConfig(RBLNModelConfig):
-    def __init__(self, batch_size: Optional[int] = None, **kwargs):
+    def __init__(self, batch_size: Optional[int] = None, **kwargs: Dict[str, Any]):
         """
         Args:
             batch_size (Optional[int]): The batch size for text processing. Defaults to 1.
@@ -34,11 +34,16 @@ class RBLNCLIPTextModelConfig(RBLNModelConfig):
 
 
 class RBLNCLIPTextModelWithProjectionConfig(RBLNCLIPTextModelConfig):
-    pass
+    """
+    Configuration class for RBLNCLIPTextModelWithProjection.
+
+    This configuration inherits from RBLNCLIPTextModelConfig and stores
+    configuration parameters for CLIP text models with projection layers.
+    """
 
 
 class RBLNCLIPVisionModelConfig(RBLNModelConfig):
-    def __init__(self, batch_size: Optional[int] = None, image_size: Optional[int] = None, **kwargs):
+    def __init__(self, batch_size: Optional[int] = None, image_size: Optional[int] = None, **kwargs: Dict[str, Any]):
         """
         Args:
             batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
@@ -76,4 +81,9 @@ class RBLNCLIPVisionModelConfig(RBLNModelConfig):
 
 
 class RBLNCLIPVisionModelWithProjectionConfig(RBLNCLIPVisionModelConfig):
-    pass
+    """
+    Configuration class for RBLNCLIPVisionModelWithProjection.
+
+    This configuration inherits from RBLNCLIPVisionModelConfig and stores
+    configuration parameters for CLIP vision models with projection layers.
+    """