optimum-rbln 0.8.2a0__py3-none-any.whl → 0.9.3__py3-none-any.whl

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

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPVisionConfig, CLIPVisionModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
@@ -50,8 +51,10 @@ class RBLNCLIPTextModel(RBLNModel):
     on RBLN devices, supporting text encoding for multimodal tasks.
     """
 
+    _tp_support = False
+
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
 
     @classmethod
@@ -82,7 +85,18 @@ class RBLNCLIPTextModel(RBLNModel):
         rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config
 
-    def forward(self, input_ids: torch.LongTensor, return_dict: bool = None, **kwargs) -> torch.FloatTensor:
+    def forward(self, input_ids: torch.LongTensor, return_dict: Optional[bool] = None, **kwargs) -> torch.FloatTensor:
+        """
+        Forward pass for the RBLN-optimized CLIP text encoder model.
+
+        Args:
+            input_ids (torch.LongTensor): The input ids to the model.
+            return_dict (Optional[bool]): Whether to return a dictionary of outputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CLIPTextModelOutput object.
+        """
+
         # To ignore using attention_mask, we override forward method.
         output = super().forward(input_ids, return_dict=return_dict)
         return output
@@ -111,12 +125,27 @@ class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
 
 
 class _VisionEncoder(torch.nn.Module):
-    def __init__(self, enc: CLIPVisionModel):
+    def __init__(
+        self,
+        enc: CLIPVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.enc = enc
+        self.interpolate_pos_encoding = interpolate_pos_encoding
+        self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
-        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
+        enc_out = self.enc(
+            inp,
+            output_hidden_states=self.output_hidden_states,
+            interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
+            return_dict=False,
+        )
         return enc_out
 
 
@@ -128,9 +157,16 @@ class RBLNCLIPVisionModel(RBLNModel):
     on RBLN devices, supporting image encoding for multimodal tasks.
     """
 
+    _tp_support = False
+
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
-        return _VisionEncoder(model).eval()
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
+        wrapper_cfg = {
+            "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
+            "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
+        }
+        return _VisionEncoder(model, **wrapper_cfg).eval()
 
     @classmethod
     def update_rbln_config_using_pipe(
@@ -155,6 +191,12 @@ class RBLNCLIPVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
                 (
@@ -175,28 +217,91 @@ class RBLNCLIPVisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        return_dict: bool = None,
+        pixel_values: torch.FloatTensor,
+        return_dict: bool = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        interpolate_pos_encoding: bool = False,
         **kwargs,
-    ) -> Union[Tuple, CLIPVisionModelOutput]:
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPooling object.
+        """
+
         if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
             logger.warning(
                 f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
             )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
+        if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
+            raise ValueError(
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
+                f"Please compile again with the correct argument."
+            )
+
         output = super().forward(pixel_values, return_dict=return_dict)
         return output
 
     def _prepare_output(self, output, return_dict):
         # Prepare model output based on return_dict flag.
         # This method can be overridden by subclasses to provide task-specific output handling.
+        last_hidden_state = output.pop(0)
+        pooler_output = output.pop(0)
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
 
         if not return_dict:
-            return (output,) if not isinstance(output, (tuple, list)) else output
+            return tuple(
+                item for item in (last_hidden_state, pooler_output, hidden_states, attentions) if item is not None
+            )
         else:
-            return CLIPVisionModelOutput(
-                image_embeds=output[0],
-                last_hidden_state=output[1],
-                hidden_states=output[2:],
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
             )
 
 
@@ -210,19 +315,70 @@ class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values: torch.FloatTensor,
+        return_dict: bool = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        interpolate_pos_encoding: bool = False,
         **kwargs,
     ) -> Union[Tuple, CLIPVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
-
-        output = super().forward(pixel_values)
-        image_embeds = output[0]
-        last_hidden_state = output[1]
-        hidden_states = output[2:]
-
-        return CLIPVisionModelOutput(
-            image_embeds=image_embeds,
-            last_hidden_state=last_hidden_state,
-            hidden_states=hidden_states,
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model with projection.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CLIPVisionModelOutput object.
+        """
+
+        return super().forward(
+            pixel_values=pixel_values,
+            return_dict=return_dict,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+            **kwargs,
         )
+
+    def _prepare_output(self, output, return_dict):
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
+        image_embeds = output.pop(0) if isinstance(output, (tuple, list)) else output
+        last_hidden_state = output.pop(0)
+
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
+
+        if not return_dict:
+            return tuple(
+                item for item in (image_embeds, last_hidden_state, hidden_states, attentions) if item is not None
+            )
+
+        else:
+            return CLIPVisionModelOutput(
+                image_embeds=image_embeds,
+                last_hidden_state=last_hidden_state,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )

optimum/rbln/transformers/models/colpali/colpali_architecture.py

@@ -4,10 +4,7 @@ import torch
 from torch import nn
 from transformers import GemmaForCausalLM, GemmaModel
 
-from ..decoderonly.decoderonly_architecture import (
-    RotaryEmbedding,
-    apply_rotary_pos_emb,
-)
+from ..decoderonly.decoderonly_architecture import RotaryEmbedding, apply_rotary_pos_emb
 
 
 def slice_and_unsqueeze_cos_sin(cos, sin, position_ids):
@@ -27,11 +24,11 @@ class RBLNColPaliForRetrievalWrapper(nn.Module):
         output_hidden_states: bool = False,
     ):
         super().__init__()
-        self.text_config = causal_lm.config
+        self.text_config = causal_lm.config.text_config
         self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
 
         self.output_hidden_states = output_hidden_states
-        self.language_model = self.convert_to_rbln_language_model(causal_lm.model, max_seq_len)
+        self.language_model = self.convert_to_rbln_language_model(causal_lm.model.language_model, max_seq_len)
 
         self.num_hidden_layers = getattr(self.text_config, "num_hidden_layers", None)
         self.embedding_proj_layer = embedding_proj_layer

optimum/rbln/transformers/models/colpali/configuration_colpali.py

@@ -11,9 +11,13 @@
 # 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, List, Optional, Union
+from typing import Any, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
@@ -24,45 +28,57 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
     including vision tower settings and multi-sequence length support.
 
     Example usage:
-    ```python
-    from optimum.rbln import RBLNColPaliForRetrieval, RBLNColPaliForRetrievalConfig
-
-    # Create a configuration object
-    config = RBLNColPaliForRetrievalConfig(
-        max_seq_lens=1152,
-        output_hidden_states=False,
-        tensor_parallel_size=4
-    )
-
-    # Use the configuration with from_pretrained
-    model = RBLNColPaliForRetrieval.from_pretrained(
-        "vidore/colpali-v1.3-hf",
-        export=True,
-        rbln_config=config
-    )
-    ```
+        ```python
+        from optimum.rbln import RBLNColPaliForRetrieval, RBLNColPaliForRetrievalConfig
+
+        # Create a configuration object
+        config = RBLNColPaliForRetrievalConfig(
+            max_seq_lens=1152,
+            output_hidden_states=False,
+            tensor_parallel_size=4
+        )
+
+        # Use the configuration with from_pretrained
+        model = RBLNColPaliForRetrieval.from_pretrained(
+            "vidore/colpali-v1.3-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     submodules = ["vision_tower"]
 
     def __init__(
         self,
+        batch_size: Optional[int] = None,
         max_seq_lens: Union[int, List[int]] = None,
         output_hidden_states: Optional[bool] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
+            batch_size (Optional[int]): The batch size for the model.
             vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
             max_seq_lens (Union[int, List[int]]): The maximum sequence lengths for the language model.
                 This can be multiple values, and the model will be compiled for each max_seq_len, allowing selection of the most appropriate max_seq_len at inference time.
             output_hidden_states (Optional[bool]): Whether to output the hidden states of the language model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
-        self.vision_tower = vision_tower
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for ColPali vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
         self.max_seq_lens = max_seq_lens
         self.output_hidden_states = output_hidden_states