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

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

@@ -22,13 +22,13 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union
 
 from transformers import AutoModelForDepthEstimation
 from transformers.modeling_outputs import DepthEstimatorOutput
 
 from ....modeling_base import RBLNModel
-from ....modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 logger = logging.getLogger(__name__)
@@ -47,9 +47,11 @@ class RBLNDPTForDepthEstimation(RBLNModel):
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_image_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
+        rbln_image_size = rbln_kwargs.get("image_size", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
@@ -79,10 +81,22 @@ class RBLNDPTForDepthEstimation(RBLNModel):
 
         input_info = [("pixel_values", [rbln_batch_size, 3, rbln_image_size[0], rbln_image_size[1]], "float32")]
 
-        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
-        meta = {"rbln_image_size": rbln_image_size}
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "image_size": rbln_image_size,
+                "batch_size": rbln_batch_size,
+            }
+        )
 
-        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        return rbln_config
 
     def forward(self, *args, **kwargs):
         predicted_depth = super().forward(*args, **kwargs)

optimum/rbln/transformers/models/gemma/modeling_gemma.py

@@ -52,7 +52,7 @@ class RBLNGemmaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        rbln_max_seq_len = rbln_config.meta["rbln_max_seq_len"]
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
         return GemmaWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:

optimum/rbln/transformers/models/gpt2/gpt2_architecture.py

@@ -53,6 +53,7 @@ class GPT2LMHeadModelWrapper(torch.nn.Module):
         attention_mask,
         cache_position,
         batch_position,
+        query_idx,
         *past_key_values,
     ):
         if input_ids.shape[1] == 1:
@@ -79,11 +80,13 @@ class GPT2LMHeadModelWrapper(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 _GPT2Model:

optimum/rbln/transformers/models/gpt2/modeling_gpt2.py

@@ -53,7 +53,7 @@ class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        rbln_max_seq_len = rbln_config.meta["rbln_max_seq_len"]
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
         return GPT2LMHeadModelWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:

optimum/rbln/transformers/models/llama/modeling_llama.py

@@ -52,7 +52,7 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        rbln_max_seq_len = rbln_config.meta["rbln_max_seq_len"]
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
         return LlamaWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:

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

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .modeling_llava_next import RBLNLlavaNextForConditionalGeneration