optimum-rbln 0.8.2a0__py3-none-any.whl → 0.9.3__py3-none-any.whl

optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -13,13 +13,21 @@
 # limitations under the License.
 
 
+from typing import TYPE_CHECKING, Optional, Union
+
 import torch
-from transformers import AutoModelForMaskedLM, Wav2Vec2ForCTC
+from transformers import AutoModelForCTC, Wav2Vec2Config, Wav2Vec2ForCTC
+from transformers.modeling_outputs import CausalLMOutput
 
-from ...modeling_generic import RBLNModelForMaskedLM
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
 from .configuration_wav2vec2 import RBLNWav2Vec2ForCTCConfig
 
 
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+
 class _Wav2Vec2(torch.nn.Module):
     def __init__(self, model: "Wav2Vec2ForCTC"):
         super().__init__()
@@ -30,22 +38,67 @@ class _Wav2Vec2(torch.nn.Module):
         return self.model.lm_head(output[0])
 
 
-class RBLNWav2Vec2ForCTC(RBLNModelForMaskedLM):
+class RBLNWav2Vec2ForCTC(RBLNModel):
     """
     Wav2Vec2 Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).
 
-    This model inherits from [`RBLNModelForMaskedLM`]. Check the superclass documentation for the generic methods the
-    library implements for all its model.
-
     It implements the methods to convert a pre-trained Wav2Vec2 model into a RBLN Wav2Vec2 model by:
+
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
     """
 
     main_input_name = "input_values"
-    auto_model_class = AutoModelForMaskedLM
+    auto_model_class = AutoModelForCTC
     rbln_dtype = "float32"
 
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNWav2Vec2ForCTCConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNWav2Vec2ForCTCConfig) -> torch.nn.Module:
         return _Wav2Vec2(model).eval()
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "Wav2Vec2Config" = None,
+        rbln_config: Optional[RBLNWav2Vec2ForCTCConfig] = None,
+    ) -> RBLNWav2Vec2ForCTCConfig:
+        if rbln_config.max_seq_len is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_max_length"):
+                    rbln_config.max_seq_len = tokenizer.model_max_length
+                    break
+            if rbln_config.max_seq_len is None:
+                raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "input_values",
+                    [
+                        rbln_config.batch_size,
+                        rbln_config.max_seq_len,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+        return rbln_config
+
+    def forward(
+        self, input_values: torch.Tensor, return_dict: Optional[bool] = None, **kwargs
+    ) -> Union[CausalLMOutput, tuple]:
+        """
+        Forward pass for the RBLN-optimized Wav2Vec2 model for Connectionist Temporal Classification (CTC).
+
+        Args:
+            input_values (torch.FloatTensor of shape (batch_size, sequence_length)): Float values of input raw speech waveform. Values can be obtained by loading a .flac or .wav audio file into an array of type List[float] or a numpy.ndarray, e.g. via the soundfile library (pip install soundfile). To prepare the array into input_values, the AutoProcessor should be used for padding and conversion into a tensor of type torch.FloatTensor.
+            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CausalLMOutput object.
+        """
+        return super().forward(input_values=input_values, return_dict=return_dict, **kwargs)

optimum/rbln/transformers/models/whisper/configuration_whisper.py

@@ -12,9 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict
-
-import rebel
+from typing import Any
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
@@ -38,17 +36,22 @@ class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
         use_attention_mask: bool = None,
         enc_max_seq_len: int = None,
         dec_max_seq_len: int = None,
-        **kwargs: Dict[str, Any],
+        kvcache_num_blocks: int = None,
+        kvcache_block_size: int = None,
+        **kwargs: Any,
     ):
         """
         Args:
             batch_size (int, optional): The batch size for inference. Defaults to 1.
             token_timestamps (bool, optional): Whether to output token timestamps during generation. Defaults to False.
             use_attention_mask (bool, optional): Whether to use attention masks during inference. This is automatically
-                set to True for RBLN-CA02 devices.
             enc_max_seq_len (int, optional): Maximum sequence length for the encoder.
             dec_max_seq_len (int, optional): Maximum sequence length for the decoder.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kvcache_num_blocks (int, optional): The total number of blocks to allocate for the
+                PagedAttention KV cache for the SelfAttention. Defaults to batch_size.
+            kvcache_block_size (int, optional): Sets the size (in number of tokens) of each block
+                in the PagedAttention KV cache for the SelfAttention. Defaults to dec_max_seq_len.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.
@@ -64,10 +67,6 @@ class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
         self.dec_max_seq_len = dec_max_seq_len
 
         self.use_attention_mask = use_attention_mask
-        npu = self.npu or rebel.get_npu_name()
-        if npu == "RBLN-CA02":
-            if self.use_attention_mask is False:
-                logger.warning("Attention mask should be used with RBLN-CA02. Setting use_attention_mask to True.")
-            self.use_attention_mask = True
-        else:
-            self.use_attention_mask = self.use_attention_mask or False
+        self.use_attention_mask = self.use_attention_mask or False
+        self.kvcache_num_blocks = kvcache_num_blocks
+        self.kvcache_block_size = kvcache_block_size

optimum/rbln/transformers/models/whisper/generation_whisper.py

@@ -31,22 +31,73 @@ Generation utilities for Whisper.
 Modified from `transformers.models.whisper.generation_whisper.py`
 """
 
+from typing import Any, Dict, Optional, Union
+
 import torch
 import transformers
 from packaging import version
 from transformers import GenerationMixin
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.modeling_outputs import ModelOutput
 from transformers.models.whisper.generation_whisper import WhisperGenerationMixin
 
 
 class RBLNWhisperGenerationMixin(WhisperGenerationMixin, GenerationMixin):
-    """
-    This class is based on transformers version 4.44.2.
-    It uses the same generate() method, so it's crucial to maintain the inheritance order.
-    Ensure WhisperGenerationMixin is listed before GenerationMixin.
-    """
+    def generate(
+        self,
+        input_features: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        return_segments: Optional[bool] = None,
+        return_timestamps: Optional[bool] = None,
+        return_token_timestamps: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[ModelOutput, Dict[str, Any], torch.LongTensor]:
+        """
+        The generate function is utilized in its standard form as in the HuggingFace transformers library. User can use this function to generate text from the model.
+        Check the [HuggingFace transformers documentation](https://huggingface.co/docs/transformers/v4.57.1/en/model_doc/whisper#transformers.WhisperForConditionalGeneration.generate) for more details.
+
+        Args:
+            input_features(torch.Tensor, optional): The input features to the model.
+            attention_mask(torch.Tensor, optional): Attention mask needs to be passed when doing long-form transcription using a batch size > 1.
+            generation_config(GenerationConfig, optional): The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them.
+                If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration.
+                Please note that unspecified parameters will inherit [GenerationConfig](https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/text_generation#transformers.GenerationConfig)’s default values.
+            return_segments(bool, optional): Whether to return segments.
+            return_timestamps(bool, optional): Whether to return the timestamps with the text. For audios longer than 30 seconds, it is necessary to set return_timestamps=True.
+            return_token_timestamps(bool, optional): Whether to return token timestamps.
+            kwargs(dict[str, Any], optional): Additional arguments passed to the generate function.
+
+        Returns:
+            Transcribes or translates log-mel input features to a sequence of auto-regressively generated token ids.
+        """
+        if kwargs.get("num_beams", None) is not None:
+            if kwargs.get("num_beams") != 1:
+                raise ValueError(
+                    "Beam search is not supported in RBLNWhisperGenerationMixin. "
+                    "Received num_beams={num_beams}, but only num_beams=1 is allowed. "
+                    "Please set num_beams=1 for greedy search or adjust your configuration."
+                )
+
+        return super().generate(
+            input_features,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            return_segments=return_segments,
+            return_timestamps=return_timestamps,
+            return_token_timestamps=return_token_timestamps,
+            **kwargs,
+        )
 
     def _postprocess_outputs(
-        self, seek_outputs, decoder_input_ids, return_token_timestamps, generation_config, *args, **kwargs
+        self,
+        seek_outputs,
+        decoder_input_ids,
+        return_token_timestamps,
+        generation_config,
+        is_shortform,
+        seek,
+        batch_idx_map,
     ):
         # remove all previously passed decoder input ids
         # should happen only if it is the first generated segment
@@ -64,6 +115,11 @@ class RBLNWhisperGenerationMixin(WhisperGenerationMixin, GenerationMixin):
 
         if return_token_timestamps and hasattr(generation_config, "alignment_heads"):
             num_frames = getattr(generation_config, "num_frames", None)
+
+            if num_frames is not None:
+                num_frames = num_frames - seek
+                num_frames = num_frames[batch_idx_map]
+
             if version.parse(transformers.__version__) >= version.parse("4.46.0"):
                 seek_outputs["token_timestamps"] = self._extract_token_timestamps(
                     seek_outputs,

optimum/rbln/transformers/models/whisper/modeling_whisper.py

@@ -46,7 +46,7 @@ if TYPE_CHECKING:
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+    def forward(self, *args: List[torch.Tensor], **kwargs: torch.Tensor):
         output = super().forward(*args, **kwargs)
         return BaseModelOutput(last_hidden_state=output)
 
@@ -73,6 +73,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         decoder_input_ids: torch.Tensor = None,
         decoder_attention_mask: torch.Tensor = None,
         cache_position: torch.Tensor = None,
+        block_tables: torch.Tensor = None,
     ):
         inputs_bsz = decoder_input_ids.shape[0]
         padded_bsz = self.batch_size - inputs_bsz
@@ -89,11 +90,14 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
                     )
                 decoder_attention_mask[b_idx, : decoding_step + 1] = 1
 
+        if block_tables is None:
+            block_tables = self.default_block_tables
+
         outputs = super().forward(
             decoder_input_ids,
             decoder_attention_mask if self.use_attention_mask else None,
             cache_position,
-            block_tables=self.default_block_tables,
+            block_tables=block_tables,
         )
 
         if isinstance(outputs, torch.Tensor):
@@ -108,6 +112,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
 
     This model inherits from [`RBLNModel`]. It implements the methods to convert and run
     pre-trained transformers based Whisper model on RBLN devices by:
+
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
 
@@ -145,7 +150,8 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
     """
 
     auto_model_class = AutoModelForSpeechSeq2Seq
-    main_input_name = "input_ids"
+    main_input_name = "input_features"
+    _is_stateful = False
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
@@ -197,7 +203,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         raise NotImplementedError
 
     @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNWhisperForConditionalGenerationConfig):
+    def _wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNWhisperForConditionalGenerationConfig):
         return WhisperWrapper(
             model,
             use_attention_mask=rbln_config.use_attention_mask,
@@ -207,7 +213,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model, rbln_config: RBLNWhisperForConditionalGenerationConfig):
-        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
+        wrapped_model = cls._wrap_model_if_needed(model, rbln_config)
 
         enc_compile_config = rbln_config.compile_cfgs[0]
         dec_compile_config = rbln_config.compile_cfgs[1]
@@ -249,6 +255,23 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
 
         return {"encoder": compiled_encoder, "decoder": compiled_decoder}
 
+    @classmethod
+    def _update_paged_attention_config(
+        cls, model_config: "PretrainedConfig", rbln_config: RBLNWhisperForConditionalGenerationConfig
+    ):
+        rbln_config.kvcache_num_blocks = rbln_config.kvcache_num_blocks or rbln_config.batch_size
+        rbln_config.kvcache_block_size = rbln_config.kvcache_block_size or rbln_config.dec_max_seq_len
+
+        if rbln_config.kvcache_num_blocks != rbln_config.batch_size:
+            raise NotImplementedError(
+                f"kvcache_num_blocks ({rbln_config.kvcache_num_blocks}) must be equal to batch_size ({rbln_config.batch_size}) as flash attention is not supported yet."
+            )
+
+        if rbln_config.kvcache_block_size != rbln_config.dec_max_seq_len:
+            raise NotImplementedError(
+                f"kvcache_block_size ({rbln_config.kvcache_block_size}) must be equal to dec_max_seq_len ({rbln_config.dec_max_seq_len}) as flash attention is not supported yet."
+            )
+
     @classmethod
     def _update_rbln_config(
         cls,
@@ -266,6 +289,8 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         if rbln_config.dec_max_seq_len is None:
             rbln_config.dec_max_seq_len = model_config.max_length
 
+        cls._update_paged_attention_config(model_config, rbln_config)
+
         enc_input_info = [
             ("input_features", [1, num_mel_bins, expected_seq_len], "float32"),
             ("block_tables", [1], "int16"),
@@ -345,12 +370,14 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
                 tensor_type="pt",
                 device=rbln_config.device_map["encoder"],
                 activate_profiler=rbln_config.activate_profiler,
+                timeout=rbln_config.timeout,
             ),
             rebel.Runtime(
                 compiled_models[1],
                 tensor_type="pt",
                 device=rbln_config.device_map["decoder"],
                 activate_profiler=rbln_config.activate_profiler,
+                timeout=rbln_config.timeout,
             ),
         ]
 

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

@@ -12,14 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_xlm_roberta import (
-    RBLNXLMRobertaForSequenceClassificationConfig,
-    RBLNXLMRobertaModelConfig,
-)
-from .modeling_xlm_roberta import (
-    RBLNXLMRobertaForSequenceClassification,
-    RBLNXLMRobertaModel,
-)
+from .configuration_xlm_roberta import RBLNXLMRobertaForSequenceClassificationConfig, RBLNXLMRobertaModelConfig
+from .modeling_xlm_roberta import RBLNXLMRobertaForSequenceClassification, RBLNXLMRobertaModel
 
 
 __all__ = [

optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional, Union
+
+import torch
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions, SequenceClassifierOutput
+
 from ...modeling_generic import RBLNModelForSequenceClassification, RBLNTransformerEncoderForFeatureExtraction
 
 
@@ -20,6 +25,25 @@ class RBLNXLMRobertaModel(RBLNTransformerEncoderForFeatureExtraction):
     XLM-RoBERTa base model optimized for RBLN NPU.
     """
 
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[BaseModelOutputWithPoolingAndCrossAttentions, tuple]:
+        """
+        Forward pass for the RBLN-optimized XLM-RoBERTa base model.
+
+        Args:
+            input_ids (torch.Tensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPoolingAndCrossAttentions object.
+        """
+
+        return super().forward(input_ids, attention_mask, **kwargs)
+
 
 class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     """
@@ -27,3 +51,22 @@ class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification
     """
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[SequenceClassifierOutput, tuple]:
+        """
+        Forward pass for the RBLN-optimized XLM-RoBERTa model for sequence classification.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional): Indices of input sequence tokens in the vocabulary.
+            attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional): Mask to avoid performing attention on padding token indices.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a SequenceClassifierOutput object.
+        """
+
+        return super().forward(input_ids, attention_mask, **kwargs)