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

optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -326,7 +326,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         attention_mask: torch.Tensor,
         position_ids: torch.Tensor,
         token_type_ids: Optional[torch.Tensor] = None,
-    ):
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.Tensor]:
         """
         Pads inputs, attention_mask, and position_ids so image token groups (256 tokens with token_type_ids == 1)
         start at multiples of prefill_chunk_size (256). Returns padded tensors and total padded length.
@@ -338,7 +338,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             token_type_ids: (1, seq_len) tensor, 0 for text, 1 for image.
 
         Returns:
-            Tuple: (inputs_padded, attention_mask_padded, position_ids_padded, padded_len, token_type_ids_padded).
+            (inputs_padded, attention_mask_padded, position_ids_padded, padded_len, token_type_ids_padded).
         """
 
         if token_type_ids is None:
@@ -816,9 +816,11 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
                     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,
                 )

optimum/rbln/transformers/models/llava_next/configuration_llava_next.py

@@ -18,6 +18,14 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaNextForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA-Next models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
     submodules = ["vision_tower", "language_model"]
 
     def __init__(

optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py

@@ -19,6 +19,14 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLNQwen2_5_VLForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Qwen2.5-VL models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
     submodules = ["visual"]
 
     def __init__(
@@ -37,6 +45,14 @@ class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausal
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNQwen2_5_VisionTransformerPretrainedModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Qwen2.5-VL vision transformer models with window-based attention
+    mechanisms for processing images and videos.
+    """
+
     def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
         """
         Args:

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -54,6 +54,14 @@ if TYPE_CHECKING:
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
+    """
+    RBLN optimized Qwen2.5-VL vision transformer model.
+
+    This class provides hardware-accelerated inference for Qwen2.5-VL vision transformers
+    on RBLN devices, supporting image and video encoding for multimodal vision-language tasks
+    with window-based attention mechanisms.
+    """
+
     auto_model_class = None
 
     def __post_init__(self, **kwargs):

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -17,4 +17,9 @@ from ...configuration_generic import RBLNModelForImageClassificationConfig
 
 
 class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    ""
+    """
+    Configuration class for RBLNResNetForImageClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized ResNet models for image classification tasks.
+    """

optimum/rbln/transformers/models/resnet/modeling_resnet.py

@@ -18,5 +18,9 @@ from ...modeling_generic import RBLNModelForImageClassification
 
 class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     """
-    ResNet model for image classification tasks on RBLN NPU.
+    RBLN optimized ResNet model for image classification tasks.
+
+    This class provides hardware-accelerated inference for ResNet models
+    on RBLN devices, supporting image classification with convolutional neural networks
+    designed for computer vision tasks.
     """

optimum/rbln/transformers/models/roberta/configuration_roberta.py

@@ -16,8 +16,18 @@ from ...configuration_generic import RBLNModelForMaskedLMConfig, RBLNModelForSeq
 
 
 class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
-    ""
+    """
+    Configuration class for RBLNRobertaForMaskedLM.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for masked language modeling tasks.
+    """
 
 
 class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
-    ""
+    """
+    Configuration class for RBLNRobertaForSequenceClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for sequence classification tasks.
+    """

optimum/rbln/transformers/models/roberta/modeling_roberta.py

@@ -16,8 +16,24 @@ from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassi
 
 
 class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
+    """
+    RBLN optimized RoBERTa model for masked language modeling tasks.
+
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting masked language modeling tasks such as
+    token prediction and text completion.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
+    """
+    RBLN optimized RoBERTa model for sequence classification tasks.
+
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting text classification tasks such as sentiment analysis,
+    topic classification, and other sequence-level prediction tasks.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask"]

optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -161,16 +161,20 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        compiled_encoder = super().compile(
+        compiled_encoder = cls.compile(
             wrapped_model.encoder,
             enc_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=enc_example_inputs,
             compile_context=context,
         )
 
-        compiled_decoder = super().compile(
+        compiled_decoder = cls.compile(
             wrapped_model.decoder,
             dec_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=dec_example_inputs,
             compile_context=context,
         )

optimum/rbln/transformers/models/siglip/configuration_siglip.py

@@ -18,6 +18,13 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNSiglipVisionModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNSiglipVisionModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized SigLIP vision models for image encoding in multimodal tasks.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,

optimum/rbln/transformers/models/siglip/modeling_siglip.py

@@ -58,6 +58,13 @@ class _SiglipVisionModel(torch.nn.Module):
 
 
 class RBLNSiglipVisionModel(RBLNModel):
+    """
+    RBLN optimized SigLIP vision model.
+
+    This class provides hardware-accelerated inference for SigLIP vision models
+    on RBLN devices, supporting image encoding for multimodal vision-language tasks.
+    """
+
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig) -> torch.nn.Module:
         wrapper_cfg = {

optimum/rbln/transformers/models/t5/configuration_t5.py

@@ -17,8 +17,18 @@ from ..seq2seq import RBLNModelForSeq2SeqLMConfig
 
 
 class RBLNT5EncoderModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
-    pass
+    """
+    Configuration class for RBLNT5EncoderModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized T5 encoder models for feature extraction tasks.
+    """
 
 
 class RBLNT5ForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
-    pass
+    """
+    Configuration class for RBLNT5ForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized T5 models for conditional text generation tasks.
+    """

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable
 
 import torch
 from transformers import AutoModelForTextEncoding, T5EncoderModel, T5ForConditionalGeneration
+from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
 
 from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
@@ -64,7 +65,7 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     """
 
     auto_model_class = AutoModelForTextEncoding
-    rbln_model_input_names = ["input_ids", "attention_mask"]
+    output_class = BaseModelOutputWithPastAndCrossAttentions
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
@@ -74,11 +75,16 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     def update_rbln_config_using_pipe(
         cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
-        submodule_config = getattr(rbln_config, submodule_name)
-        submodule_config.max_seq_len = rbln_config.max_seq_len or 256
-        submodule_config.model_input_names = ["input_ids"]
         return rbln_config
 
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        input_dict = {"input_ids": input_ids.long()}
+        if attention_mask is not None:
+            input_dict["attention_mask"] = attention_mask.long()
+
+        output = super().forward(**input_dict, **kwargs)
+        return output
+
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
     """

optimum/rbln/transformers/models/time_series_transformer/configuration_time_series_transformer.py

@@ -4,6 +4,13 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNTimeSeriesTransformerForPredictionConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNTimeSeriesTransformerForPrediction.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Time Series Transformer models for time series forecasting tasks.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,

optimum/rbln/transformers/models/time_series_transformer/modeling_time_series_transformer.py

@@ -194,15 +194,19 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        compiled_decoder = super().compile(
+        compiled_decoder = cls.compile(
             wrapped_model.decoder,
             dec_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=dec_example_inputs,
             compile_context=context,
         )
-        compiled_encoder = super().compile(
+        compiled_encoder = cls.compile(
             wrapped_model.encoder,
             enc_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=enc_example_inputs,
             compile_context=context,
         )

optimum/rbln/transformers/models/vit/configuration_vit.py

@@ -16,4 +16,9 @@ from ...configuration_generic import RBLNModelForImageClassificationConfig
 
 
 class RBLNViTForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    ""
+    """
+    Configuration class for RBLNViTForImageClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Vision Transformer (ViT) models for image classification tasks.
+    """

optimum/rbln/transformers/models/vit/modeling_vit.py

@@ -16,4 +16,10 @@ from ...modeling_generic import RBLNModelForImageClassification
 
 
 class RBLNViTForImageClassification(RBLNModelForImageClassification):
-    ""
+    """
+    RBLN optimized Vision Transformer (ViT) model for image classification tasks.
+
+    This class provides hardware-accelerated inference for Vision Transformer models
+    on RBLN devices, supporting image classification with transformer-based architectures
+    that process images as sequences of patches.
+    """

optimum/rbln/transformers/models/wav2vec2/configuration_wav2vec2.py

@@ -16,4 +16,11 @@ from ...configuration_generic import RBLNModelForMaskedLMConfig
 
 
 class RBLNWav2Vec2ForCTCConfig(RBLNModelForMaskedLMConfig):
+    """
+    Configuration class for RBLNWav2Vec2ForCTC.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Wav2Vec2 models for Connectionist Temporal Classification (CTC) tasks.
+    """
+
     rbln_model_input_names = ["input_values"]

optimum/rbln/transformers/models/whisper/configuration_whisper.py

@@ -24,6 +24,13 @@ logger = get_logger()
 
 
 class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNWhisperForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Whisper models for speech recognition and transcription tasks.
+    """
+
     def __init__(
         self,
         batch_size: int = None,

optimum/rbln/transformers/models/whisper/modeling_whisper.py

@@ -230,15 +230,19 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        compiled_encoder = super().compile(
+        compiled_encoder = cls.compile(
             wrapped_model.encoder,
             enc_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=enc_example_inputs,
             compile_context=context,
         )
-        compiled_decoder = super().compile(
+        compiled_decoder = cls.compile(
             wrapped_model.decoder,
             dec_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=dec_example_inputs,
             compile_context=context,
         )

optimum/rbln/utils/runtime_utils.py

@@ -13,12 +13,57 @@
 # limitations under the License.
 
 import threading
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Union
 
 import rebel
 import torch
 
 
+def tp_and_devices_are_ok(
+    tensor_parallel_size: Optional[int] = None,
+    device: Optional[Union[int, List[int]]] = None,
+    npu: Optional[str] = None,
+) -> Optional[str]:
+    if tensor_parallel_size is None:
+        tensor_parallel_size = 1
+
+    if rebel.device_count() < tensor_parallel_size:
+        return (
+            f"Tensor parallel size {tensor_parallel_size} is greater than "
+            f"the number of available devices {rebel.device_count()}."
+        )
+
+    if device is None:
+        device = list(range(tensor_parallel_size))
+    elif isinstance(device, int):
+        device = [device]
+    elif isinstance(device, list):
+        if any(not isinstance(d, int) for d in device):
+            return "Device must be a(n) (list of) integer(s)."
+        if len(device) != tensor_parallel_size:
+            return (
+                f"The number of devices ({len(device)}) does not match tensor parallel size ({tensor_parallel_size})."
+            )
+    else:
+        return f"Invalid device: {device}"
+
+    for device_id in device:
+        if device_id < 0:  # if any device is dummy device, skip it
+            return None
+        if rebel.get_npu_name(device_id) is None:
+            return (
+                f"Device {device_id} is not a valid NPU device. Please check your NPU status with 'rbln-stat' command."
+            )
+
+    if npu is not None:
+        for device_id in device:
+            npu_name = rebel.get_npu_name(device_id)
+            if npu_name != npu:
+                return f"Device {device_id} ({npu_name}) is not on the same NPU as {npu}."
+
+    return None
+
+
 class RBLNPytorchRuntime:
     mandatory_members = []
 
@@ -43,6 +88,9 @@ class RBLNPytorchRuntime:
     def __repr__(self) -> str:
         return repr(self.runtime)
 
+    def parameters(self):
+        yield torch.tensor([1.0], dtype=torch.float32, device=torch.device("cpu"))
+
 
 class UnavailableRuntime:
     """

{optimum_rbln-0.8.1a5.dist-info → optimum_rbln-0.8.1a7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a5
+Version: 0.8.1a7
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai