optimum-rbln 0.9.1__py3-none-any.whl → 0.9.2__py3-none-any.whl

optimum/rbln/transformers/__init__.py

@@ -54,6 +54,8 @@ _import_structure = {
         "RBLNBlip2VisionModelConfig",
         "RBLNColPaliForRetrieval",
         "RBLNColPaliForRetrievalConfig",
+        "RBLNColQwen2ForRetrieval",
+        "RBLNColQwen2ForRetrievalConfig",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelConfig",
         "RBLNCLIPTextModelWithProjection",
@@ -110,6 +112,8 @@ _import_structure = {
         "RBLNPegasusModelConfig",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNLlavaNextForConditionalGenerationConfig",
+        "RBLNLoRAAdapterConfig",
+        "RBLNLoRAConfig",
         "RBLNMidmLMHeadModel",
         "RBLNMidmLMHeadModelConfig",
         "RBLNMistralForCausalLM",
@@ -216,6 +220,8 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjectionConfig,
         RBLNColPaliForRetrieval,
         RBLNColPaliForRetrievalConfig,
+        RBLNColQwen2ForRetrieval,
+        RBLNColQwen2ForRetrievalConfig,
         RBLNDecoderOnlyModel,
         RBLNDecoderOnlyModelConfig,
         RBLNDecoderOnlyModelForCausalLM,
@@ -258,6 +264,8 @@ if TYPE_CHECKING:
         RBLNLlavaForConditionalGenerationConfig,
         RBLNLlavaNextForConditionalGeneration,
         RBLNLlavaNextForConditionalGenerationConfig,
+        RBLNLoRAAdapterConfig,
+        RBLNLoRAConfig,
         RBLNMidmLMHeadModel,
         RBLNMidmLMHeadModelConfig,
         RBLNMistralForCausalLM,

optimum/rbln/transformers/modeling_generic.py

@@ -23,6 +23,7 @@ different model architectures.
 import inspect
 from typing import TYPE_CHECKING, Optional, Union
 
+from torch import nn
 from transformers import (
     AutoModel,
     AutoModelForAudioClassification,
@@ -57,6 +58,28 @@ class RBLNTransformerEncoder(RBLNModel):
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
     rbln_dtype = "int64"
 
+    @classmethod
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNTransformerEncoderConfig) -> nn.Module:
+        class TransformerEncoderWrapper(nn.Module):
+            # Parameters to disable for RBLN compilation
+            DISABLED_PARAMS = {"return_dict", "use_cache"}
+
+            def __init__(self, model: "PreTrainedModel", rbln_config: RBLNTransformerEncoderConfig):
+                super().__init__()
+                self.model = model
+                self.rbln_config = rbln_config
+                self._forward_signature = inspect.signature(model.forward)
+
+            def forward(self, *args, **kwargs):
+                # Disable parameters that are not compatible with RBLN compilation
+                for param_name in self.DISABLED_PARAMS:
+                    if param_name in self._forward_signature.parameters:
+                        kwargs[param_name] = False
+
+                return self.model(*args, **kwargs)
+
+        return TransformerEncoderWrapper(model, rbln_config).eval()
+
     @classmethod
     def _update_rbln_config(
         cls,
@@ -208,7 +231,6 @@ class RBLNModelForQuestionAnswering(RBLNTransformerEncoder):
 
     def _prepare_output(self, output, return_dict):
         # Prepare QuestionAnswering specific output format.
-
         start_logits, end_logits = output
 
         if not return_dict:
@@ -245,6 +267,20 @@ class RBLNModelForImageClassification(RBLNImageModel):
 class RBLNModelForDepthEstimation(RBLNImageModel):
     auto_model_class = AutoModelForDepthEstimation
 
+    @classmethod
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
+        class ImageModelWrapper(nn.Module):
+            def __init__(self, model: "PreTrainedModel", rbln_config: RBLNImageModelConfig):
+                super().__init__()
+                self.model = model
+                self.rbln_config = rbln_config
+
+            def forward(self, *args, **kwargs):
+                output = self.model(*args, return_dict=True, **kwargs)
+                return output.predicted_depth
+
+        return ImageModelWrapper(model, rbln_config).eval()
+
 
 class RBLNModelForAudioClassification(RBLNModel):
     """

optimum/rbln/transformers/models/__init__.py

@@ -75,6 +75,10 @@ _import_structure = {
         "RBLNColPaliForRetrieval",
         "RBLNColPaliForRetrievalConfig",
     ],
+    "colqwen2": [
+        "RBLNColQwen2ForRetrieval",
+        "RBLNColQwen2ForRetrievalConfig",
+    ],
     "distilbert": [
         "RBLNDistilBertForQuestionAnswering",
         "RBLNDistilBertForQuestionAnsweringConfig",
@@ -96,6 +100,8 @@ _import_structure = {
         "RBLNDecoderOnlyModel",
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
+        "RBLNLoRAAdapterConfig",
+        "RBLNLoRAConfig",
     ],
     "depth_anything": ["RBLNDepthAnythingForDepthEstimationConfig", "RBLNDepthAnythingForDepthEstimation"],
     "dpt": [
@@ -234,11 +240,14 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjectionConfig,
     )
     from .colpali import RBLNColPaliForRetrieval, RBLNColPaliForRetrievalConfig
+    from .colqwen2 import RBLNColQwen2ForRetrieval, RBLNColQwen2ForRetrievalConfig
     from .decoderonly import (
         RBLNDecoderOnlyModel,
         RBLNDecoderOnlyModelConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
+        RBLNLoRAAdapterConfig,
+        RBLNLoRAConfig,
     )
     from .depth_anything import RBLNDepthAnythingForDepthEstimation, RBLNDepthAnythingForDepthEstimationConfig
     from .distilbert import RBLNDistilBertForQuestionAnswering, RBLNDistilBertForQuestionAnsweringConfig

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

@@ -15,6 +15,10 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNBlip2VisionModelConfig(RBLNModelConfig):
@@ -25,6 +29,16 @@ class RBLNBlip2VisionModelConfig(RBLNModelConfig):
     RBLN-optimized BLIP-2 vision encoder models for multimodal tasks.
     """
 
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        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}")
+
 
 class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
     """
@@ -36,6 +50,7 @@ class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
 
     def __init__(
         self,
+        batch_size: Optional[int] = None,
         num_query_tokens: Optional[int] = None,
         image_text_hidden_size: Optional[int] = None,
         **kwargs,
@@ -47,11 +62,22 @@ class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
             kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
+        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}")
+
         self.num_query_tokens = num_query_tokens
         self.image_text_hidden_size = image_text_hidden_size
 
 
 class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNBlip2ForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BLIP-2 models for conditional generation tasks that involve both image and text inputs.
+    """
+
     submodules = ["vision_model", "qformer", "language_model"]
 
     def __init__(
@@ -78,6 +104,12 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         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}")
 
-        self.vision_model = self.init_submodule_config(RBLNBlip2VisionModelConfig, vision_model)
-        self.language_model = language_model
-        self.qformer = self.init_submodule_config(RBLNBlip2QFormerModelConfig, qformer)
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Blip2 vision model. It will be set to 1.")
+            logger.warning("Ignore batch_size for Blip2 qformer. It will be set to 1.")
+
+        self.vision_model = self.initialize_submodule_config(
+            submodule_config=vision_model, batch_size=1, force_kwargs=True
+        )
+        self.qformer = self.initialize_submodule_config(submodule_config=qformer, batch_size=1, force_kwargs=True)
+        self.language_model = self.initialize_submodule_config(submodule_config=language_model)

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

@@ -30,34 +30,31 @@ from transformers.utils import logging
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
+from ...utils.rbln_runtime_wrapper import LoopProcessor
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 logger = logging.get_logger(__name__)
 
 if TYPE_CHECKING:
+    import rebel
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
-class LoopProjector:
-    def __init__(self, language_projection) -> None:
-        self.language_projection = language_projection
+class LoopProjector(LoopProcessor):
+    def __init__(self, language_projection: Union[RBLNModel, "rebel.Runtime"]):
+        super().__init__(model=language_projection)
 
-    def forward(self, *args, **kwargs):
-        query_output = args[0]
+    def _get_batch_size(self, query_output, **kwargs):
+        return query_output.shape[0]
 
-        batch_size = query_output.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.language_projection(query_output[i : i + 1]))
-
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
-
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
+    def _prepare_inputs_for_iteration(self, index, common_inputs, query_output, **kwargs):
+        query_output_item = query_output[index : index + 1]
+        return ([query_output_item], {})
 
-    def __repr__(self) -> str:
-        return repr(self.language_projection)
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = torch.cat(outputs, dim=0)
+        return output
 
 
 class RBLNBlip2VisionModel(RBLNModel):
@@ -68,6 +65,8 @@ class RBLNBlip2VisionModel(RBLNModel):
     on RBLN devices, supporting image encoding for multimodal vision-language tasks.
     """
 
+    _tp_support = False
+
     def get_input_embeddings(self):
         return self.embeddings
 
@@ -96,8 +95,7 @@ class RBLNBlip2VisionModel(RBLNModel):
             (
                 "pixel_values",
                 [
-                    # support for vllm CB (prefill)
-                    1,
+                    rbln_config.batch_size,
                     model_config.num_channels,
                     model_config.image_size,
                     model_config.image_size,
@@ -147,6 +145,8 @@ class RBLNBlip2QFormerModel(RBLNModel):
     mechanisms for multimodal understanding tasks.
     """
 
+    _tp_support = False
+
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
 
@@ -200,7 +200,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "query_embeds",
                 [
-                    1,
+                    rbln_config.batch_size,
                     rbln_config.num_query_tokens,
                     model_config.hidden_size,
                 ],
@@ -209,7 +209,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "encoder_hidden_states",
                 [
-                    1,
+                    rbln_config.batch_size,
                     # image_text_hidden_size + cls token
                     rbln_config.image_text_hidden_size + 1,
                     model_config.encoder_hidden_size,
@@ -219,7 +219,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "encoder_attention_mask",
                 # image_text_hidden_size + cls token
-                [1, rbln_config.image_text_hidden_size + 1],
+                [rbln_config.batch_size, rbln_config.image_text_hidden_size + 1],
                 "int64",
             ),
         ]
@@ -266,7 +266,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
         )
 
 
-class RBLNBlip2ForConditionalGeneration(RBLNModel):
+class RBLNBlip2ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     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.
@@ -434,3 +434,66 @@ class RBLNBlip2ForConditionalGeneration(RBLNModel):
             )
 
         return inputs_embeds
+
+    @torch.no_grad()
+    def generate(
+        self,
+        pixel_values: torch.FloatTensor,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        interpolate_pos_encoding: bool = False,
+        **generate_kwargs,
+    ) -> torch.LongTensor:
+        batch_size = pixel_values.shape[0]
+        image_embeds = self.vision_model(
+            pixel_values,
+            return_dict=True,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+        ).last_hidden_state
+        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
+
+        query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_outputs = self.qformer(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_attention_mask,
+            return_dict=True,
+        )
+        query_output = query_outputs.last_hidden_state
+
+        if query_output.dtype != image_embeds.dtype:
+            query_output = query_output.to(image_embeds.dtype)
+
+        language_model_inputs = self.language_projection(query_output)
+
+        if inputs_embeds is None:
+            if input_ids is None:
+                image_tokens = [self.config.image_token_index] * self.config.num_query_tokens
+                start_tokens = image_tokens + [self.config.text_config.bos_token_id]
+                input_ids = torch.tensor([start_tokens], dtype=torch.long, device=image_embeds.device)
+                input_ids = input_ids.repeat(batch_size, 1)
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_image_mask = special_image_mask.all(-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+
+        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        language_model_inputs = language_model_inputs.to(inputs_embeds.device, inputs_embeds.dtype)
+        inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, language_model_inputs)
+
+        inputs = {"inputs_embeds": inputs_embeds, "attention_mask": attention_mask}
+        if not self.language_model.config.is_encoder_decoder:
+            inputs["input_ids"] = input_ids
+
+        outputs = self.language_model.generate(**inputs, **generate_kwargs)
+
+        return outputs

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

@@ -51,6 +51,8 @@ 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:
         return _TextEncoder(model).eval()
@@ -152,6 +154,8 @@ 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:
         wrapper_cfg = {

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

@@ -24,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

@@ -14,6 +14,10 @@
 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,29 +28,30 @@ 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,
@@ -54,6 +59,8 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
     ):
         """
         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.
@@ -63,6 +70,15 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
             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