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

optimum/rbln/transformers/models/exaone/modeling_exaone.py

@@ -0,0 +1,78 @@
+# 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.
+
+import inspect
+import logging
+from typing import TYPE_CHECKING, Any, Callable
+
+from ....modeling_config import RBLNConfig
+from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .exaone_architecture import ExaoneForCausalLMWrapper
+from .hf_hub_cached.modeling_exaone import ExaoneForCausalLM
+
+
+logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from transformers import (
+        PreTrainedModel,
+    )
+
+
+class RBLNExaoneForCausalLM(RBLNDecoderOnlyModelForCausalLM):
+    """
+    The Exaone Model transformer with a language modeling head on top (linear layer with weights tied to the input
+    embeddings).
+
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the
+    library implements for all its model.
+
+    It implements the methods to convert a pre-trained transformers Exaone model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    """
+
+    @classmethod
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
+        return ExaoneForCausalLMWrapper(model, rbln_max_seq_len).eval()
+
+    def __getattr__(self, __name: str) -> Any:
+        """This is the key method to implement RBLN-Exaone.
+
+        Returns:
+            Any: Exaone's corresponding method
+        """
+
+        def redirect(func):
+            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
+
+        val = getattr(ExaoneForCausalLM, __name)
+        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
+            return redirect(val)
+        return val
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        kwargs.setdefault("trust_remote_code", True)
+        return super().from_pretrained(*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