optimum-rbln 0.1.9__py3-none-any.whl → 0.1.12__py3-none-any.whl

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

@@ -22,14 +22,23 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
-from transformers import AutoConfig, AutoModel, CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    CLIPTextConfig,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPVisionConfig,
+    CLIPVisionModel,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput
 
 from ....modeling_base import RBLNModel
-from ....modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 logger = logging.getLogger(__name__)
@@ -41,23 +50,17 @@ if TYPE_CHECKING:
 class _TextEncoder(torch.nn.Module):
     def __init__(self, enc: "CLIPTextModel"):
         super().__init__()
-        enc.config.return_dict = False
-        enc.config.output_hidden_states = True
         self.enc = enc
 
     def forward(self, inp):
-        enc_out = self.enc(inp)
+        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
         return enc_out
 
 
 class RBLNCLIPTextModel(RBLNModel):
-    auto_model_class = AutoModel  # feature extraction
     original_model_class = CLIPTextModel
     original_config_class = CLIPTextConfig
 
-    def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
-
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         configtmp = AutoConfig.from_pretrained
@@ -78,28 +81,32 @@ class RBLNCLIPTextModel(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "CLIPTextConfig",
+        rbln_kwargs: Dict[str, Any] = {},
         rbln_batch_size: Optional[int] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
     ) -> RBLNConfig:
-        model_config.return_dict = False
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        rbln_runtime_config = RBLNRuntimeConfig(
-            input_info=[
-                (
-                    "input_ids",
-                    [
-                        rbln_batch_size,
-                        model_config.max_position_embeddings,
-                    ],
-                    "int64",
-                ),
-            ],
-        )
+        model_config.return_dict = False
 
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config])
+        input_info = [
+            (
+                "input_ids",
+                [
+                    rbln_batch_size,
+                    model_config.max_position_embeddings,
+                ],
+                "int64",
+            ),
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
         return rbln_config
 
     def forward(self, input_ids: "torch.Tensor", **kwargs):
@@ -113,3 +120,97 @@ class RBLNCLIPTextModel(RBLNModel):
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
     original_model_class = CLIPTextModelWithProjection
+
+
+class _VisionEncoder(torch.nn.Module):
+    def __init__(self, enc: CLIPVisionModel):
+        super().__init__()
+        self.enc = enc
+
+    def forward(self, inp):
+        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
+        return enc_out
+
+
+class RBLNCLIPVisionModel(RBLNModel):
+    original_model_class = CLIPVisionModel
+    original_config_class = CLIPVisionConfig
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        configtmp = AutoConfig.from_pretrained
+        modeltmp = AutoModel.from_pretrained
+        AutoConfig.from_pretrained = cls.original_config_class.from_pretrained
+        AutoModel.from_pretrained = cls.original_model_class.from_pretrained
+        rt = super().from_pretrained(*args, **kwargs)
+        AutoConfig.from_pretrained = configtmp
+        AutoModel.from_pretrained = modeltmp
+        return rt
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+        return _VisionEncoder(model).eval()
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model_config: "CLIPTextConfig",
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_batch_size = rbln_kwargs.get("batch_size", 1)
+        rbln_image_size = rbln_kwargs.get("image_size", None)
+
+        if rbln_image_size is None:
+            rbln_image_size = getattr(model_config, "image_size", None)
+
+        if isinstance(rbln_image_size, int):
+            rbln_image_size = (rbln_image_size, rbln_image_size)
+
+        if rbln_image_size is None:
+            raise ValueError("`rbln_image_size` should be specified!")
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "pixel_values",
+                    [
+                        rbln_batch_size,
+                        3,
+                        rbln_image_size[0],
+                        rbln_image_size[1],
+                    ],
+                    "float32",
+                )
+            ]
+        )
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "batch_size": rbln_batch_size,
+                "image_size": rbln_image_size,
+            }
+        )
+
+        return rbln_config
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        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)
+        return BaseModelOutputWithPooling(
+            last_hidden_state=output[0],
+            pooler_output=output[1],
+            hidden_states=output[2:],
+        )

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

@@ -93,17 +93,29 @@ class DecoderOnlyWrapper(torch.nn.Module):
 
     def forward(
         self,
-        input_ids,
+        input_ids_or_inputs_embeds,
         attention_mask,
         cache_position,
         batch_position,
+        query_idx,
         *past_key_values,
     ):
-        if input_ids.shape[1] == 1:
+        if input_ids_or_inputs_embeds.shape[1] == 1:
             rbln_batch_position = None
         else:
             rbln_batch_position = batch_position
 
+        if input_ids_or_inputs_embeds.ndim == 2:
+            # input_ids
+            input_ids = input_ids_or_inputs_embeds
+            inputs_embeds = None
+        elif input_ids_or_inputs_embeds.ndim == 3:
+            # inputs_embeds
+            input_ids = None
+            inputs_embeds = input_ids_or_inputs_embeds
+        else:
+            raise NotImplementedError(f"Unknown ndim of input : {input_ids_or_inputs_embeds.ndim}")
+
         # Formatting list of past_kv to DynamicCache class.
         past_key_values = RebelDynamicCache.from_input_format(
             cache_position,
@@ -115,6 +127,7 @@ class DecoderOnlyWrapper(torch.nn.Module):
         outputs = forward_dict["wrapper"](
             self.model,
             input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             position_ids=cache_position,
             past_key_values=past_key_values,
@@ -124,11 +137,14 @@ class DecoderOnlyWrapper(torch.nn.Module):
         )
 
         hidden_states = outputs[0]
+        if batch_position >= 0:
+            hidden_states = hidden_states[:, query_idx].unsqueeze(1)
+
         logits = self.lm_head(hidden_states)
 
         output = (logits,) + outputs[1:]
 
-        return output, batch_position
+        return output, batch_position + query_idx
 
 
 class DecoderOnlyAttention:
@@ -323,8 +339,16 @@ class DecoderOnlyModel:
         forward_dict: Optional[Dict[str, classmethod]] = None,
         rotary_pos_emb=None,
     ) -> BaseModelOutputWithPast:
+        # retrieve input_ids and inputs_embeds
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
         # embed positions
-        inputs_embeds = self.embed_tokens(input_ids)
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
         hidden_states = inputs_embeds
         attention_mask = (1 - attention_mask) * torch.finfo(torch.float16).min