optimum-rbln 0.8.1a5__py3-none-any.whl → 0.8.1a7__py3-none-any.whl

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

@@ -0,0 +1,98 @@
+# 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 typing import Any, Dict, Optional
+
+from diffusers import CosmosVideoToWorldPipeline
+from diffusers.schedulers import EDMEulerScheduler
+from transformers import T5TokenizerFast
+
+from ....transformers.models.t5.modeling_t5 import RBLNT5EncoderModel
+from ....utils.logging import get_logger
+from ...modeling_diffusers import RBLNDiffusionMixin
+from ...models.autoencoders.autoencoder_kl_cosmos import RBLNAutoencoderKLCosmos
+from ...models.transformers.transformer_cosmos import RBLNCosmosTransformer3DModel
+from .cosmos_guardrail import RBLNCosmosSafetyChecker
+
+
+logger = get_logger(__name__)
+
+
+class RBLNCosmosVideoToWorldPipeline(RBLNDiffusionMixin, CosmosVideoToWorldPipeline):
+    """
+    RBLN-accelerated implementation of Cosmos Video to World pipeline for video-to-video generation.
+
+    This pipeline compiles Cosmos Video to World models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating images with distinctive artistic style and enhanced visual quality.
+    """
+
+    original_class = CosmosVideoToWorldPipeline
+    _submodules = ["text_encoder", "transformer", "vae"]
+    _optional_components = ["safety_checker"]
+
+    def __init__(
+        self,
+        text_encoder: RBLNT5EncoderModel,
+        tokenizer: T5TokenizerFast,
+        transformer: RBLNCosmosTransformer3DModel,
+        vae: RBLNAutoencoderKLCosmos,
+        scheduler: EDMEulerScheduler,
+        safety_checker: RBLNCosmosSafetyChecker = None,
+    ):
+        if safety_checker is None:
+            safety_checker = RBLNCosmosSafetyChecker()
+
+        super().__init__(
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            vae=vae,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+        )
+
+    def handle_additional_kwargs(self, **kwargs):
+        if "num_frames" in kwargs and kwargs["num_frames"] != self.transformer.rbln_config.num_frames:
+            logger.warning(
+                f"The transformer in this pipeline is compiled with 'num_frames={self.transformer.rbln_config.num_frames}'. 'num_frames' set by the user will be ignored"
+            )
+            kwargs.pop("num_frames")
+        if (
+            "max_sequence_length" in kwargs
+            and kwargs["max_sequence_length"] != self.transformer.rbln_config.max_seq_len
+        ):
+            logger.warning(
+                f"The transformer in this pipeline is compiled with 'max_seq_len={self.transformer.rbln_config.max_seq_len}'. 'max_sequence_length' set by the user will be ignored"
+            )
+            kwargs.pop("max_sequence_length")
+        return kwargs
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str,
+        *,
+        export: bool = False,
+        safety_checker: Optional[RBLNCosmosSafetyChecker] = None,
+        rbln_config: Dict[str, Any] = {},
+        **kwargs: Dict[str, Any],
+    ):
+        rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
+        if safety_checker is None and export:
+            safety_checker = RBLNCosmosSafetyChecker(rbln_config=rbln_config.safety_checker)
+
+        return super().from_pretrained(
+            model_id, export=export, safety_checker=safety_checker, rbln_config=rbln_config, **kwargs
+        )

optimum/rbln/modeling.py

@@ -64,7 +64,12 @@ class RBLNModel(RBLNBaseModel):
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         model = cls.wrap_model_if_needed(model, rbln_config)
         rbln_compile_config = rbln_config.compile_cfgs[0]
-        compiled_model = cls.compile(model, rbln_compile_config=rbln_compile_config)
+        compiled_model = cls.compile(
+            model,
+            rbln_compile_config=rbln_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
+        )
         return compiled_model
 
     @classmethod
@@ -237,7 +242,38 @@ class RBLNModel(RBLNBaseModel):
             for compiled_model in compiled_models
         ]
 
-    def forward(self, *args, return_dict: Optional[bool] = None, **kwargs):
+    def forward(self, *args: Any, return_dict: Optional[bool] = None, **kwargs: Dict[str, Any]) -> Any:
+        """
+        Defines the forward pass of the RBLN model, providing a drop-in replacement for HuggingFace PreTrainedModel.
+
+        This method executes the compiled RBLN model on RBLN NPU devices while maintaining full compatibility
+        with HuggingFace transformers and diffusers APIs. The RBLNModel can be used as a direct substitute
+        for any HuggingFace nn.Module/PreTrainedModel, enabling seamless integration into existing workflows.
+
+        Args:
+            *args: Variable length argument list containing model inputs. The format matches the original
+                HuggingFace model's forward method signature (e.g., input_ids, attention_mask for
+                transformers models, or sample, timestep for diffusers models).
+            return_dict:
+                Whether to return outputs as a dictionary-like object or as a tuple. When `None`:
+                - For transformers models: Uses `self.config.use_return_dict` (typically `True`)
+                - For diffusers models: Defaults to `True`
+            **kwargs: Arbitrary keyword arguments containing additional model inputs and parameters,
+                matching the original HuggingFace model's interface.
+
+        Returns:
+            Model outputs in the same format as the original HuggingFace model.
+
+            - If `return_dict=True`: Returns a dictionary-like object (e.g., BaseModelOutput,
+                CausalLMOutput) with named fields such as `logits`, `hidden_states`, etc.
+            - If `return_dict=False`: Returns a tuple containing the raw model outputs.
+
+        Note:
+            - This method maintains the exact same interface as the original HuggingFace model's forward method
+            - The compiled model runs on RBLN NPU hardware for accelerated inference
+            - All HuggingFace model features (generation, attention patterns, etc.) are preserved
+            - Can be used directly in HuggingFace pipelines, transformers.Trainer, and other workflows
+        """
         if self.hf_library_name == "transformers":
             return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         else:

optimum/rbln/modeling_base.py

@@ -27,7 +27,7 @@ from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConf
 from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
 from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
-from .utils.runtime_utils import UnavailableRuntime
+from .utils.runtime_utils import UnavailableRuntime, tp_and_devices_are_ok
 from .utils.save_utils import maybe_load_preprocessors
 from .utils.submodule import SubModulesMixin
 
@@ -374,7 +374,23 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return from_pretrained_method(model_id=model_id, **kwargs, rbln_config=rbln_config)
 
     @classmethod
-    def compile(cls, model, rbln_compile_config: Optional[RBLNCompileConfig] = None, **kwargs):
+    def compile(
+        cls,
+        model,
+        rbln_compile_config: RBLNCompileConfig,
+        create_runtimes: bool,
+        device: Union[int, List[int]],
+        **kwargs,
+    ):
+        if create_runtimes:
+            runtime_cannot_be_created = tp_and_devices_are_ok(
+                tensor_parallel_size=rbln_compile_config.tensor_parallel_size,
+                device=device,
+                npu=rbln_compile_config.npu,
+            )
+            if runtime_cannot_be_created:
+                raise ValueError(runtime_cannot_be_created)
+
         compiled_model = rebel.compile_from_torch(
             model,
             input_info=rbln_compile_config.input_info,

optimum/rbln/transformers/modeling_generic.py

@@ -139,7 +139,7 @@ class RBLNTransformerEncoder(RBLNModel):
         return rbln_config
 
 
-class _RBLNImageModel(RBLNModel):
+class RBLNImageModel(RBLNModel):
     auto_model_class = AutoModel
     main_input_name = "pixel_values"
     output_class = BaseModelOutput
@@ -233,11 +233,11 @@ class RBLNTransformerEncoderForFeatureExtraction(RBLNTransformerEncoder):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNModelForImageClassification(_RBLNImageModel):
+class RBLNModelForImageClassification(RBLNImageModel):
     auto_model_class = AutoModelForImageClassification
 
 
-class RBLNModelForDepthEstimation(_RBLNImageModel):
+class RBLNModelForDepthEstimation(RBLNImageModel):
     auto_model_class = AutoModelForDepthEstimation
 
 

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

@@ -18,10 +18,22 @@ 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,

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
 

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

@@ -34,7 +34,12 @@ 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):
@@ -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.
+    """

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

@@ -43,6 +43,13 @@ class _TextEncoder(torch.nn.Module):
 
 
 class RBLNCLIPTextModel(RBLNModel):
+    """
+    RBLN optimized CLIP text encoder model.
+
+    This class provides hardware-accelerated inference for CLIP text encoders
+    on RBLN devices, supporting text encoding for multimodal tasks.
+    """
+
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
@@ -95,7 +102,12 @@ class RBLNCLIPTextModel(RBLNModel):
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
-    pass
+    """
+    RBLN optimized CLIP text encoder model with projection layer.
+
+    This class extends RBLNCLIPTextModel with a projection layer for
+    multimodal embedding alignment tasks.
+    """
 
 
 class _VisionEncoder(torch.nn.Module):
@@ -109,6 +121,13 @@ class _VisionEncoder(torch.nn.Module):
 
 
 class RBLNCLIPVisionModel(RBLNModel):
+    """
+    RBLN optimized CLIP vision encoder model.
+
+    This class provides hardware-accelerated inference for CLIP vision encoders
+    on RBLN devices, supporting image encoding for multimodal tasks.
+    """
+
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
@@ -182,6 +201,13 @@ class RBLNCLIPVisionModel(RBLNModel):
 
 
 class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
+    """
+    RBLN optimized CLIP vision encoder model with projection layer.
+
+    This class extends RBLNCLIPVisionModel with a projection layer for
+    multimodal embedding alignment tasks.
+    """
+
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,

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

@@ -78,7 +78,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 torch.ones(1, 1, self.rbln_config.prefill_chunk_size, self.rbln_config.prefill_chunk_size), diagonal=1
             )
 
-    def get_block_tables(self, cache_position: torch.Tensor, batch_idx: int = None):
+    def get_block_tables(self, cache_position: torch.Tensor, batch_idx: int = None) -> torch.Tensor:
         """
         Manages and returns the KV cache block tables.
         Updates the block tables based on the given cache_position, allocating new blocks or reusing existing ones as needed.
@@ -88,7 +88,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             batch_idx (int, optional): Specific batch index, used when phase is 'prefill'.
 
         Returns:
-            torch.Tensor: Updated block tables.
+            Updated block tables.
         """
 
         NO_BLOCKS_ERROR = (
@@ -458,6 +458,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     This class serves as the foundation for various decoder-only architectures like GPT, LLaMA, etc.
 
     The class provides core functionality for:
+
     1. Converting pre-trained transformer models to RBLN-optimized format
     2. Handling the compilation process for RBLN devices
     3. Managing inference operations for causal language modeling
@@ -532,7 +533,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     @classmethod
     def save_torch_artifacts(
         cls,
-        model: "PreTrainedModel",
+        model: PreTrainedModel,
         save_dir_path: Path,
         subfolder: str,
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
@@ -566,7 +567,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     def get_quantized_model(
         cls,
         model_id: str,
-        config: Optional["PretrainedConfig"] = None,
+        config: Optional[PretrainedConfig] = None,
         use_auth_token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
@@ -605,16 +606,15 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         return model
 
     def __getattr__(self, __name: str) -> Any:
-        """
-        Special method to delegate attribute access to the original Huggingface LM class.
-        This method is called when an attribute is not found in the current instance's dictionary.
-        It enables transparent access to the original model's attributes and methods while maintaining
-        proper method binding.
-
-        The method implements a delegation pattern that:
-        1. For methods: Creates a wrapper that properly binds 'self' to method calls
-        2. For other attributes: Returns them directly from the original class
-        """
+        # Special method to delegate attribute access to the original Huggingface LM class.
+        # This method is called when an attribute is not found in the current instance's dictionary.
+        # It enables transparent access to the original model's attributes and methods while maintaining
+        # proper method binding.
+
+        # The method implements a delegation pattern that:
+
+        # 1. For methods: Creates a wrapper that properly binds 'self' to method calls
+        # 2. For other attributes: Returns them directly from the original class
 
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
@@ -627,7 +627,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     @classmethod
     def get_pytorch_model(
         cls, *args, rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None, **kwargs
-    ) -> "PreTrainedModel":
+    ) -> PreTrainedModel:
         if rbln_config and rbln_config.quantization:
             model = cls.get_quantized_model(*args, **kwargs)
         else:
@@ -636,7 +636,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         return model
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: "RBLNDecoderOnlyModelForCausalLMConfig"):
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: "RBLNDecoderOnlyModelForCausalLMConfig"):
         wrapper_cfg = {
             "max_seq_len": rbln_config.max_seq_len,
             "attn_impl": rbln_config.attn_impl,
@@ -654,7 +654,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
     @classmethod
     @torch.inference_mode()
-    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+    def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
         rbln_compile_configs = rbln_config.compile_cfgs
@@ -679,9 +679,11 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                     quantization.maybe_set_quantization_env()
                 original_linear = torch.nn.functional.linear
                 torch.nn.functional.linear = torch.ops.rbln_custom_ops.linear
-                compiled_model = RBLNModel.compile(
+                compiled_model = cls.compile(
                     wrapped_model,
                     compile_config,
+                    create_runtimes=rbln_config.create_runtimes,
+                    device=rbln_config.device,
                     example_inputs=example_inputs,
                     compile_context=compile_context,
                 )
@@ -973,8 +975,8 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     def _update_rbln_config(
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
-        model: Optional["PreTrainedModel"] = None,
-        model_config: Optional["PretrainedConfig"] = None,
+        model: Optional[PreTrainedModel] = None,
+        model_config: Optional[PretrainedConfig] = None,
         rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None,
     ) -> RBLNDecoderOnlyModelForCausalLMConfig:
         if rbln_config.max_seq_len is None:

optimum/rbln/transformers/models/distilbert/configuration_distilbert.py

@@ -16,4 +16,9 @@ from ...configuration_generic import RBLNModelForQuestionAnsweringConfig
 
 
 class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
-    ""
+    """
+    Configuration class for RBLNDistilBertForQuestionAnswering.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized DistilBERT models for question answering tasks.
+    """

optimum/rbln/transformers/models/distilbert/modeling_distilbert.py

@@ -16,4 +16,12 @@ from ...modeling_generic import RBLNModelForQuestionAnswering
 
 
 class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
+    """
+    RBLN optimized DistilBERT model for question answering tasks.
+
+    This class provides hardware-accelerated inference for DistilBERT 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"]

optimum/rbln/transformers/models/dpt/configuration_dpt.py

@@ -16,4 +16,9 @@ from ...configuration_generic import RBLNModelForDepthEstimationConfig
 
 
 class RBLNDPTForDepthEstimationConfig(RBLNModelForDepthEstimationConfig):
-    pass
+    """
+    Configuration class for RBLNDPTForDepthEstimation.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized DPT (Dense Prediction Transformer) models for depth estimation tasks.
+    """

optimum/rbln/transformers/models/dpt/modeling_dpt.py

@@ -17,4 +17,9 @@ from ...modeling_generic import RBLNModelForDepthEstimation
 
 
 class RBLNDPTForDepthEstimation(RBLNModelForDepthEstimation):
-    pass
+    """
+    RBLN optimized DPT model for depth estimation tasks.
+
+    This class provides hardware-accelerated inference for DPT (Dense Prediction Transformer)
+    models on RBLN devices, supporting monocular depth estimation from single images.
+    """