optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

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

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Llama models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNLlamaForCausalLM, RBLNLlamaForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNLlamaForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNLlamaForCausalLM.from_pretrained(
+        "meta-llama/Llama-2-7b-hf",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

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

@@ -27,8 +27,57 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers LlamaForCausalLM 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.
+
+    **Configuration:**
+    This model uses [`RBLNLlamaForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNLlamaForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNLlamaForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNLlamaForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNLlamaForCausalLM.from_pretrained(
+            "meta-llama/Llama-2-7b-hf",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNLlamaForCausalLM.from_pretrained(
+            "meta-llama/Llama-2-7b-hf",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNLlamaForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNLlamaForCausalLMConfig
+
+        config = RBLNLlamaForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNLlamaForCausalLM.from_pretrained(
+            "meta-llama/Llama-2-7b-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = LlamaWrapper

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

@@ -12,12 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaNextForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA-Next models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
     submodules = ["vision_tower", "language_model"]
 
     def __init__(
@@ -25,7 +33,7 @@ class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -109,6 +109,36 @@ class LoopProjector:
 
 
 class RBLNLlavaNextForConditionalGeneration(RBLNModel):
+    """
+    RBLNLlavaNextForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
+    optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
+
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model includes a Large Language Model (LLM) as a submodule. For optimal performance, it is highly recommended to use
+        tensor parallelism for the language model. This can be achieved by using the `rbln_config` parameter in the
+        `from_pretrained` method. Refer to the `from_pretrained` documentation and the RBLNLlavaNextForConditionalGenerationConfig class for details.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNLlavaNextForConditionalGeneration
+
+        model = RBLNLlavaNextForConditionalGeneration.from_pretrained(
+            "llava-hf/llava-v1.6-mistral-7b-hf",
+            export=True,
+            rbln_config={
+                "language_model": {
+                    "tensor_parallel_size": 4,
+                    "use_inputs_embeds": True,  # In Llava-Next, language model must use inputs_embeds as input.
+                },
+            },
+        )
+
+        model.save_pretrained("compiled-llava-next-mistral-7b-hf")
+        ```
+    """
+
     auto_model_class = AutoModelForVision2Seq
     _rbln_submodules = [
         {"name": "vision_tower"},
@@ -136,10 +166,8 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         subfolder: str,
         rbln_config: RBLNModelConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
         save_dict = {}
         save_dict["image_newline"] = model.image_newline
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")

optimum/rbln/transformers/models/midm/configuration_midm.py

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNMidmLMHeadModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for MIDM models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNMidmLMHeadModel, RBLNMidmLMHeadModelConfig
+
+    # Create a configuration object
+    config = RBLNMidmLMHeadModelConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNMidmLMHeadModel.from_pretrained(
+        "KT-AI/midm-bitext-S-7B-inst-v1",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/midm/midm_architecture.py

@@ -68,7 +68,12 @@ class MidmLMHeadModelWrapper(DecoderOnlyWrapper):
             )
             new_layer = MidmLayer(layer, new_self_attn)
             new_layers.append(new_layer)
-        new_model = MidmModel(causal_lm.transformer, new_layers, max_seq_len=max_seq_len)
+        new_model = MidmModel(
+            causal_lm.transformer,
+            new_layers,
+            max_seq_len=max_seq_len,
+            sliding_window_layers=self.sliding_window_layers,
+        )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
 

optimum/rbln/transformers/models/midm/modeling_midm.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
+from typing import Any, Callable
+
 from transformers import AutoModelForCausalLM
+from transformers.generation.utils import GenerationMixin
 
 from ....utils import logging
 from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
@@ -24,22 +28,79 @@ logger = logging.get_logger(__name__)
 
 class RBLNMidmLMHeadModel(RBLNDecoderOnlyModelForCausalLM):
     """
-    The Midm Model transformer with a language modeling head on top (linear layer with weights tied to the input
-    embeddings).
+    The MIDM Model transformer with a language modeling head (linear layer) on top.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
-    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the
-    library implements for all its model.
+    A class to convert and run pre-trained transformers based MidmForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers MidmForCausalLM model into a RBLN transformer model by:
 
-    It implements the methods to convert a pre-trained transformers Midm 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.
 
+    **Configuration:**
+    This model uses [`RBLNMidmLMHeadModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNMidmLMHeadModelConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNMidmLMHeadModelConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNMidmLMHeadModel
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNMidmLMHeadModel.from_pretrained(
+            "KT-AI/midm-bitext-S-7B-inst-v1",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNMidmLMHeadModel.from_pretrained(
+            "KT-AI/midm-bitext-S-7B-inst-v1",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNMidmLMHeadModelConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNMidmLMHeadModelConfig
+
+        config = RBLNMidmLMHeadModelConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNMidmLMHeadModel.from_pretrained(
+            "KT-AI/midm-bitext-S-7B-inst-v1",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = MidmLMHeadModelWrapper
     _hf_class = AutoModelForCausalLM
+    _supports_cache_class = True
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         kwargs.setdefault("trust_remote_code", True)
         return super().from_pretrained(*args, **kwargs)
+
+    def __getattr__(self, __name: str) -> Any:
+        def redirect(func):
+            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
+
+        val = getattr(GenerationMixin, __name)
+
+        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
+            return redirect(val)
+        return val

optimum/rbln/transformers/models/mistral/configuration_mistral.py

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNMistralForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Mistral models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNMistralForCausalLM, RBLNMistralForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNMistralForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNMistralForCausalLM.from_pretrained(
+        "mistralai/Mistral-7B-v0.1",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/mistral/modeling_mistral.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from transformers import PretrainedConfig
+
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM, RBLNDecoderOnlyModelForCausalLMConfig
 from .mistral_architecture import MistralForCausalLMWrapper
 
 
@@ -22,13 +24,69 @@ logger = logging.get_logger(__name__)
 
 class RBLNMistralForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
-    The Llama Model transformer with a language modeling head (linear layer) on top.
+    The Mistral Model transformer with a language modeling head (linear layer) on top.
     This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
-    A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
-    It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
+    A class to convert and run pre-trained transformers based MistralForCausalLM model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers MistralForCausalLM 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.
+
+    **Configuration:**
+    This model uses [`RBLNMistralForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNMistralForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNMistralForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNMistralForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNMistralForCausalLM.from_pretrained(
+            "mistralai/Mistral-7B-v0.1",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNMistralForCausalLM.from_pretrained(
+            "mistralai/Mistral-7B-v0.1",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+        # Using a RBLNMistralForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNMistralForCausalLMConfig
+
+        config = RBLNMistralForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNMistralForCausalLM.from_pretrained(
+            "mistralai/Mistral-7B-v0.1",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = MistralForCausalLMWrapper
+
+    @classmethod
+    def _update_sliding_window_config(
+        cls, model_config: PretrainedConfig, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig
+    ):
+        rbln_config.cache_impl = "sliding_window"
+        rbln_config.sliding_window = model_config.sliding_window
+        rbln_config.sliding_window_layers = list(range(model_config.num_hidden_layers))
+
+        return rbln_config

optimum/rbln/transformers/models/opt/configuration_opt.py

@@ -16,4 +16,7 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNOPTForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for OPT causal language model.
+    Inherits from RBLNDecoderOnlyModelForCausalLMConfig with no additional parameters.
+    """

optimum/rbln/transformers/models/opt/modeling_opt.py

@@ -45,8 +45,15 @@ class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     A class to convert and run pre-trained transformers based OPTForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers OPTForCausalLM 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.
+
+    **Configuration:**
+    This model uses [`RBLNOPTForCausalLM`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNOPTForCausalLM`] or a dictionary conforming to its structure.
+
+    See the [`RBLNOPTForCausalLM`] class for all available configuration options.
     """
 
     _decoder_wrapper_cls = OPTWrapper
@@ -72,6 +79,9 @@ class RBLNOPTForCausalLM(RBLNDecoderOnlyModelForCausalLM):
             "use_attention_mask": rbln_config.use_attention_mask,
             "use_position_ids": rbln_config.use_position_ids,
             "use_inputs_embeds": rbln_config.use_inputs_embeds,
+            "cache_impl": rbln_config.cache_impl,
+            "sliding_window": rbln_config.sliding_window,
+            "sliding_window_layers": rbln_config.sliding_window_layers,
         }
 
         for i in range(len(model.model.decoder.layers)):

optimum/rbln/transformers/models/opt/opt_architecture.py

@@ -45,7 +45,13 @@ class OPTWrapper(DecoderOnlyWrapper):
             )
             new_layer = OPTDecoderLayer(layer, new_self_attn)
             new_layers.append(new_layer)
-        new_model = OPTModel(causal_lm.model.decoder, new_layers, max_seq_len=max_seq_len, use_learned_pos_emb=True)
+        new_model = OPTModel(
+            causal_lm.model.decoder,
+            new_layers,
+            max_seq_len=max_seq_len,
+            use_learned_pos_emb=True,
+            sliding_window_layers=self.sliding_window_layers,
+        )
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
 

optimum/rbln/transformers/models/phi/configuration_phi.py

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNPhiForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Phi models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNPhiForCausalLM, RBLNPhiForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNPhiForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNPhiForCausalLM.from_pretrained(
+        "microsoft/phi-2",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """

optimum/rbln/transformers/models/phi/modeling_phi.py

@@ -27,8 +27,57 @@ class RBLNPhiForCausalLM(RBLNDecoderOnlyModelForCausalLM):
 
     A class to convert and run pre-trained transformers based PhiForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers PhiForCausalLM 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.
+
+    **Configuration:**
+    This model uses [`RBLNPhiForCausalLMConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNPhiForCausalLMConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNPhiForCausalLMConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNPhiForCausalLM
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_len` is automatically inferred from the model config
+        model = RBLNPhiForCausalLM.from_pretrained(
+            "microsoft/phi-2",
+            export=True,
+            rbln_batch_size=1,
+            rbln_tensor_parallel_size=4,
+        )
+
+
+        # Using a config dictionary
+        rbln_config = {
+            "batch_size": 1,
+            "max_seq_len": 4096,
+            "tensor_parallel_size": 4,
+        }
+        model = RBLNPhiForCausalLM.from_pretrained(
+            "microsoft/phi-2",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+
+        # Using a RBLNPhiForCausalLMConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNPhiForCausalLMConfig
+
+        config = RBLNPhiForCausalLMConfig(
+            batch_size=1,
+            max_seq_len=4096,
+            tensor_parallel_size=4
+        )
+        model = RBLNPhiForCausalLM.from_pretrained(
+            "microsoft/phi-2",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     _decoder_wrapper_cls = PhiWrapper

optimum/rbln/transformers/models/phi/phi_architecture.py

@@ -48,7 +48,7 @@ class PhiWrapper(DecoderOnlyWrapper):
                 raise NotImplementedError(f"Unknwon attn : {self.attn_impl}")
             new_layer = PhiLayer(layer, new_self_attn)
             new_layers.append(new_layer)
-        new_model = PhiModel(causal_lm.model, new_layers)
+        new_model = PhiModel(causal_lm.model, new_layers, sliding_window_layers=self.sliding_window_layers)
         new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
         return new_causal_lm
 

optimum/rbln/transformers/models/qwen2/configuration_qwen2.py

@@ -16,4 +16,27 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNQwen2ForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
-    pass
+    """
+    Configuration class for RBLN Qwen2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNQwen2ForCausalLM, RBLNQwen2ForCausalLMConfig
+
+    # Create a configuration object
+    config = RBLNQwen2ForCausalLMConfig(
+        batch_size=1,
+        max_seq_len=4096,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNQwen2ForCausalLM.from_pretrained(
+        "Qwen/Qwen2-7B",
+        export=True,
+        rbln_config=config
+    )
+    ```
+    """