transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc1__py3-none-any.whl

transformers/models/clip/tokenization_clip.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Tokenization classes for CLIP."""
 
-from typing import Optional
+from typing import Optional, Union
 
 from tokenizers import Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors
 from tokenizers.models import BPE
@@ -37,6 +37,10 @@ class CLIPTokenizer(TokenizersBackend):
     refer to this superclass for more information regarding those methods.
 
     Args:
+        vocab (`str`, `dict` or `list`, *optional*):
+            Vocabulary dict to use for the tokenizer.
+        merges (`str` or `list`, *optional*):
+            Merges list to use for the BPE tokenizer.
         unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
             The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
             token instead.
@@ -46,53 +50,38 @@ class CLIPTokenizer(TokenizersBackend):
             The end of sequence token.
         pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
             The token used for padding, for example when batching sequences of different lengths.
-        vocab (`dict`, *optional*):
-            Vocabulary dict to use for the tokenizer.
-        merges (`list`, *optional*):
-            Merges list to use for the BPE tokenizer.
-        vocab_file (`str`, *optional*):
-            Path to the vocabulary file.
-        merges_file (`str`, *optional*):
-            Path to the merges file.
     """
 
     vocab_files_names = VOCAB_FILES_NAMES
     model_input_names = ["input_ids", "attention_mask"]
-    slow_tokenizer_class = None
+    model = BPE
 
     def __init__(
         self,
+        vocab: Optional[Union[str, dict[str, int]]] = None,
+        merges: Optional[Union[str, list[str]]] = None,
         unk_token: str = "<|endoftext|>",
         bos_token: str = "<|startoftext|>",
         eos_token: str = "<|endoftext|>",
         pad_token: str = "<|endoftext|>",
-        vocab: Optional[dict] = None,
-        merges: Optional[list] = None,
-        vocab_file: Optional[str] = None,
-        merges_file: Optional[str] = None,
         **kwargs,
     ):
-        self.vocab_file = vocab_file
-        self.merges_file = merges_file
-
-        if vocab is not None:
-            _vocab = {token: idx for idx, (token, _score) in enumerate(vocab)} if isinstance(vocab, list) else vocab
-        else:
-            _vocab = {
+        _vocab = (
+            vocab
+            if vocab is not None
+            else {
                 str(bos_token): 0,
                 str(eos_token): 1,
                 str(pad_token): 2,
             }
+        )
 
-        if merges is not None:
-            _merges = [tuple(merge) if isinstance(merge, list) else merge for merge in merges]
-        else:
-            _merges = []
+        self._merges = merges or []
 
         self._tokenizer = Tokenizer(
             BPE(
                 vocab=_vocab,
-                merges=_merges,
+                merges=self._merges,
                 dropout=None,
                 continuing_subword_prefix="",
                 end_of_word_suffix="</w>",
@@ -120,20 +109,7 @@ class CLIPTokenizer(TokenizersBackend):
 
         self._tokenizer.decoder = decoders.ByteLevel()
 
-        bos_token_id = _vocab.get(str(bos_token), 0)
-        eos_token_id = _vocab.get(str(eos_token), 1)
-
-        self._tokenizer.post_processor = processors.RobertaProcessing(
-            sep=(str(eos_token), eos_token_id),
-            cls=(str(bos_token), bos_token_id),
-            add_prefix_space=False,
-            trim_offsets=False,
-        )
-
-        tokenizer_object = self._tokenizer
-
         super().__init__(
-            tokenizer_object=tokenizer_object,
             unk_token=unk_token,
             bos_token=bos_token,
             eos_token=eos_token,
@@ -141,14 +117,16 @@ class CLIPTokenizer(TokenizersBackend):
             **kwargs,
         )
 
-        if hasattr(self, "_post_init"):
-            self._post_init()
+        self._tokenizer.post_processor = processors.RobertaProcessing(
+            sep=(str(eos_token), self.eos_token_id),
+            cls=(str(bos_token), self.bos_token_id),
+            add_prefix_space=False,
+            trim_offsets=False,
+        )
 
-    def _post_init(self):
-        super()._post_init()
+        # Very ugly hack to enable padding to have a correct decoding see https://github.com/huggingface/tokenizers/issues/872
         self._wrap_decode_method_backend_tokenizer()
 
-    # Very ugly hack to enable padding to have a correct decoding see https://github.com/huggingface/tokenizers/issues/872
     def _wrap_decode_method_backend_tokenizer(self):
         orig_decode_method = self.backend_tokenizer.decode
 

transformers/models/clipseg/modeling_clipseg.py

@@ -676,6 +676,7 @@ class CLIPSegTextModel(CLIPSegPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, BaseModelOutputWithPooling]:
         r"""
         Examples:
@@ -776,6 +777,7 @@ class CLIPSegVisionModel(CLIPSegPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         interpolate_pos_encoding: Optional[bool] = True,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, BaseModelOutputWithPooling]:
         r"""
         Examples:
@@ -933,6 +935,7 @@ class CLIPSegModel(CLIPSegPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         interpolate_pos_encoding: bool = True,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, CLIPSegOutput]:
         r"""
         return_loss (`bool`, *optional*):
@@ -1125,6 +1128,7 @@ class CLIPSegDecoder(CLIPSegPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = True,
+        **kwargs,
     ):
         all_hidden_states = () if output_hidden_states else None
         all_attentions = () if output_attentions else None
@@ -1239,6 +1243,7 @@ class CLIPSegForImageSegmentation(CLIPSegPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         interpolate_pos_encoding: bool = True,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, CLIPSegOutput]:
         r"""
         conditional_pixel_values (`torch.FloatTensor`, *optional*):

transformers/models/clvp/modeling_clvp.py

@@ -861,6 +861,7 @@ class ClvpEncoder(ClvpPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, BaseModelOutput]:
         r"""
         Args:
@@ -1020,6 +1021,7 @@ class ClvpDecoder(ClvpPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
     ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -1170,6 +1172,7 @@ class ClvpModel(ClvpPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
     ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -1339,6 +1342,7 @@ class ClvpForCausalLM(ClvpPreTrainedModel, GenerationMixin):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
     ) -> Union[tuple, CausalLMOutputWithCrossAttentions]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
@@ -1635,6 +1639,7 @@ class ClvpModelForConditionalGeneration(ClvpPreTrainedModel, GenerationMixin):
         output_attentions: Optional[bool] = False,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
     ) -> Union[tuple, ClvpOutput]:
         r"""
         conditioning_encoder_inputs_embeds (`torch.FloatTensor`, *optional*):

transformers/models/clvp/tokenization_clvp.py

@@ -15,9 +15,7 @@
 """Tokenization class for CLVP."""
 
 import json
-import os
 from functools import lru_cache
-from typing import Optional
 
 import regex as re
 
@@ -123,10 +121,6 @@ class ClvpTokenizer(PreTrainedTokenizer):
         add_prefix_space (`bool`, *optional*, defaults to `False`):
             Whether or not to add an initial space to the input. This allows to treat the leading word just as any
             other word. (CLVP tokenizer detect beginning of words by the preceding space).
-        add_bos_token (`bool`, *optional*, defaults to `False`):
-            Whether to add `bos_token` in front of the sequence when add_special_tokens=True.
-        add_eos_token (`bool`, *optional*, defaults to `False`):
-            Whether to add `eos_token` in end of the sequence when add_special_tokens=True.
     """
 
     vocab_files_names = VOCAB_FILES_NAMES
@@ -145,8 +139,6 @@ class ClvpTokenizer(PreTrainedTokenizer):
         eos_token="[STOP]",
         pad_token="[STOP]",
         add_prefix_space=False,
-        add_bos_token=False,
-        add_eos_token=False,
         **kwargs,
     ):
         bos_token = AddedToken(bos_token, special=True) if isinstance(bos_token, str) else bos_token
@@ -154,20 +146,7 @@ class ClvpTokenizer(PreTrainedTokenizer):
         unk_token = AddedToken(unk_token, special=True) if isinstance(unk_token, str) else unk_token
         pad_token = AddedToken(pad_token, special=True) if isinstance(pad_token, str) else pad_token
 
-        self.add_bos_token = add_bos_token
-        self.add_eos_token = add_eos_token
         self._normalizer = None
-
-        # Set special_tokens_pattern based on add_bos_token and add_eos_token flags
-        if add_bos_token and add_eos_token:
-            kwargs["special_tokens_pattern"] = "bos_eos"
-        elif add_bos_token:
-            kwargs["special_tokens_pattern"] = "bos"
-        elif add_eos_token:
-            kwargs["special_tokens_pattern"] = "eos"
-        else:
-            kwargs["special_tokens_pattern"] = "none"
-
         with open(vocab_file, encoding="utf-8") as vocab_handle:
             self.encoder = json.load(vocab_handle)
         self.decoder = {v: k for k, v in self.encoder.items()}
@@ -191,8 +170,7 @@ class ClvpTokenizer(PreTrainedTokenizer):
             eos_token=eos_token,
             pad_token=pad_token,
             add_prefix_space=add_prefix_space,
-            add_bos_token=add_bos_token,
-            add_eos_token=add_eos_token,
+            special_tokens_pattern="none",
             **kwargs,
         )
 
@@ -251,17 +229,6 @@ class ClvpTokenizer(PreTrainedTokenizer):
         self.cache[token] = word
         return word
 
-    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
-        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
-        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
-        output = bos_token_id + token_ids_0 + eos_token_id
-
-        if token_ids_1 is not None:
-            output = output + bos_token_id + token_ids_1 + eos_token_id
-
-        return output
-
     def _tokenize(self, text):
         """Tokenize a string."""
         bpe_tokens = []
@@ -303,34 +270,5 @@ class ClvpTokenizer(PreTrainedTokenizer):
         text = text.replace(self.unk_token, "").replace("   ", " ").replace("  ", " ")
         return text
 
-    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
-        if not os.path.isdir(save_directory):
-            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
-            return
-        vocab_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
-        )
-        merge_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
-        )
-
-        with open(vocab_file, "w", encoding="utf-8") as f:
-            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
-        index = 0
-        with open(merge_file, "w", encoding="utf-8") as writer:
-            writer.write("#version: 0.2\n")
-            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
-                if index != token_index:
-                    logger.warning(
-                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
-                        " Please check that the tokenizer is not corrupted!"
-                    )
-                    index = token_index
-                writer.write(" ".join(bpe_tokens) + "\n")
-                index += 1
-
-        return vocab_file, merge_file
-
 
 __all__ = ["ClvpTokenizer"]

transformers/models/code_llama/tokenization_code_llama.py

@@ -14,10 +14,11 @@
 # limitations under the License.
 
 
-from tokenizers import AddedToken, Tokenizer, decoders, normalizers, pre_tokenizers, processors
+from typing import Optional, Union
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
 from tokenizers.models import BPE
 
-from ...tokenization_utils_base import _get_prepend_scheme, generate_merges
 from ...tokenization_utils_tokenizers import TokenizersBackend
 from ...utils import logging
 
@@ -97,9 +98,9 @@ class CodeLlamaTokenizer(TokenizersBackend):
         add_prefix_space (`bool`, *optional*):
             Whether or not to add an initial space to the input. This allows to treat the leading word just as any
             other word.
-        vocab (`dict`, *optional*):
+        vocab (`str`, `dict` or `list`, *optional*):
             Custom vocabulary dictionary. If not provided, vocabulary is loaded from vocab_file.
-        merges (`list`, *optional*):
+        merges (`str` or `list`, *optional*):
             Custom merges list. If not provided, merges are loaded from merges_file.
         vocab_file (`str`, *optional*):
             [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .model extension) that
@@ -109,9 +110,12 @@ class CodeLlamaTokenizer(TokenizersBackend):
     vocab_files_names = VOCAB_FILES_NAMES
     padding_side = "left"
     model_input_names = ["input_ids", "attention_mask"]
+    model = BPE
 
     def __init__(
         self,
+        vocab: Optional[Union[str, dict[str, int]]] = None,
+        merges: Optional[Union[str, list[str]]] = None,
         clean_up_tokenization_spaces=False,
         unk_token="<unk>",
         bos_token="<s>",
@@ -122,37 +126,28 @@ class CodeLlamaTokenizer(TokenizersBackend):
         eot_token="▁<EOT>",
         fill_token="<FILL_ME>",
         additional_special_tokens=None,
-        add_bos_token=True,
-        add_eos_token=False,
-        use_default_system_prompt=False,
-        add_prefix_space=None,
-        vocab=None,
-        merges=None,
-        vocab_file=None,
+        use_default_system_prompt: bool = False,
+        add_prefix_space: Optional[bool] = True,
+        add_bos_token: bool = True,
         **kwargs,
     ):
         self.add_prefix_space = add_prefix_space if add_prefix_space is not None else True
         self.use_default_system_prompt = use_default_system_prompt
-
         additional_special_tokens = additional_special_tokens or []
         for token in [prefix_token, middle_token, suffix_token, eot_token, fill_token]:
             additional_special_tokens += [token] if token is not None else []
 
-        if vocab is not None:
-            self._vocab = (
-                {token: idx for idx, (token, _score) in enumerate(vocab)} if isinstance(vocab, list) else vocab
-            )
-        else:
-            self._vocab = {
+        self._vocab = (
+            vocab
+            if vocab is not None
+            else {
                 str(unk_token): 0,
                 str(bos_token): 1,
                 str(eos_token): 2,
             }
+        )
 
-        filtered_vocab = {
-            t: i for t, i in self._vocab.items() if t not in {str(eos_token), str(bos_token), str(unk_token)}
-        }
-        self._merges = merges if merges is not None else generate_merges(filtered_vocab)
+        self._merges = merges or []
         self._tokenizer = Tokenizer(
             BPE(
                 vocab=self._vocab,
@@ -163,8 +158,9 @@ class CodeLlamaTokenizer(TokenizersBackend):
                 unk_token=str(unk_token),
             )
         )
+        prepend_scheme = "first" if self.add_prefix_space else "none"
         self._tokenizer.pre_tokenizer = pre_tokenizers.Metaspace(
-            replacement="▁", prepend_scheme=_get_prepend_scheme(self.add_prefix_space, self), split=False
+            replacement="▁", prepend_scheme=prepend_scheme, split=False
         )
 
         self._tokenizer.decoder = decoders.Sequence(
@@ -172,13 +168,10 @@ class CodeLlamaTokenizer(TokenizersBackend):
         )
 
         super().__init__(
-            tokenizer_object=self._tokenizer,
             clean_up_tokenization_spaces=clean_up_tokenization_spaces,
             unk_token=unk_token,
             bos_token=bos_token,
             eos_token=eos_token,
-            add_bos_token=add_bos_token,
-            add_eos_token=add_eos_token,
             use_default_system_prompt=use_default_system_prompt,
             add_prefix_space=add_prefix_space,
             prefix_token=prefix_token,
@@ -186,32 +179,16 @@ class CodeLlamaTokenizer(TokenizersBackend):
             suffix_token=suffix_token,
             eot_token=eot_token,
             fill_token=fill_token,
+            add_bos_token=add_bos_token,
             additional_special_tokens=additional_special_tokens,
             **kwargs,
         )
-
-        self._add_bos_token = add_bos_token
-        self._add_eos_token = add_eos_token
-        self.vocab_file = vocab_file
-
         self._prefix_token = prefix_token
         self._middle_token = middle_token
         self._suffix_token = suffix_token
         self._eot_token = eot_token
         self.fill_token = fill_token
 
-        self._post_init()
-
-    def _post_init(self):
-        self._tokenizer.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="first", split=False)
-        self._tokenizer.normalizer = None
-
-        # This matches LlamaTokenizer's behavior and is needed when loading from vocab/merges
-        self.add_tokens([AddedToken(token, special=True) for token in self.all_special_tokens])
-
-        self.update_post_processor()
-        super()._post_init()
-
     @property
     def prefix_token(self):
         return self._prefix_token

transformers/models/codegen/tokenization_codegen.py

@@ -67,6 +67,10 @@ class CodeGenTokenizer(TokenizersBackend):
     refer to this superclass for more information regarding those methods.
 
     Args:
+        vocab (`str` or `dict[str, int]`, *optional*):
+            Custom vocabulary dictionary. If not provided, vocabulary is loaded from `vocab_file`.
+        merges (`str` or `list[str]`, *optional*):
+            Custom merges list. If not provided, merges are loaded from `merges_file`.
         unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
             The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
             token instead.
@@ -79,31 +83,24 @@ class CodeGenTokenizer(TokenizersBackend):
         add_prefix_space (`bool`, *optional*, defaults to `False`):
             Whether or not to add an initial space to the input. This allows to treat the leading word just as any
             other word. (CodeGen tokenizer detect beginning of words by the preceding space).
-        add_bos_token (`bool`, *optional*, defaults to `False`):
-            Whether or not to add an initial beginning of sentence token to the input.
         return_token_type_ids (`bool`, *optional*, defaults to `False`):
             Whether to return token type IDs.
-        vocab (`dict`, *optional*):
-            Custom vocabulary dictionary. If not provided, vocabulary is loaded from vocab_file.
-        merges (`list`, *optional*):
-            Custom merges list. If not provided, merges are loaded from merges_file.
     """
 
     vocab_files_names = VOCAB_FILES_NAMES
     model_input_names = ["input_ids", "attention_mask"]
-    slow_tokenizer_class = None
+    model = BPE
 
     def __init__(
         self,
-        unk_token="<|endoftext|>",
-        bos_token="<|endoftext|>",
-        eos_token="<|endoftext|>",
+        vocab: Optional[Union[str, dict[str, int]]] = None,
+        merges: Optional[Union[str, list[str]]] = None,
+        unk_token: str = "<|endoftext|>",
+        bos_token: str = "<|endoftext|>",
+        eos_token: str = "<|endoftext|>",
         pad_token=None,
-        add_prefix_space=False,
-        add_bos_token=False,
-        return_token_type_ids=False,
-        vocab: Optional[dict] = None,
-        merges: Optional[list] = None,
+        add_prefix_space: bool = False,
+        return_token_type_ids: bool = False,
         **kwargs,
     ):
         self.return_token_type_ids = return_token_type_ids
@@ -112,17 +109,8 @@ class CodeGenTokenizer(TokenizersBackend):
 
         self.add_prefix_space = add_prefix_space
 
-        if vocab is not None:
-            self._vocab = (
-                {token: idx for idx, (token, _score) in enumerate(vocab)} if isinstance(vocab, list) else vocab
-            )
-        else:
-            self._vocab = {}
-
-        if merges is not None:
-            self._merges = merges
-        else:
-            self._merges = []
+        self._vocab = vocab if vocab is not None else {}
+        self._merges = merges or []
 
         self._tokenizer = Tokenizer(
             BPE(
@@ -141,33 +129,16 @@ class CodeGenTokenizer(TokenizersBackend):
             add_prefix_space=True, use_regex=True, trim_offsets=False
         )
 
-        tokenizer_object = self._tokenizer
-
-        # Set these before calling super().__init__() so the base class _post_init() can use them
-        self._add_bos_token = add_bos_token
-        self._add_eos_token = False
-
         super().__init__(
-            tokenizer_object=tokenizer_object,
             unk_token=unk_token,
             bos_token=bos_token,
             eos_token=eos_token,
             pad_token=pad_token,
             add_prefix_space=add_prefix_space,
-            add_bos_token=add_bos_token,
             return_token_type_ids=return_token_type_ids,
             **kwargs,
         )
 
-        self._post_init()
-
-    def _post_init(self):
-        self._tokenizer.post_processor = processors.ByteLevel(
-            add_prefix_space=True, use_regex=True, trim_offsets=False
-        )
-        # Ensure base class post-init runs to register special/extra tokens, etc.
-        super()._post_init()
-
     def decode(
         self,
         token_ids: Union[int, list[int], np.ndarray, "torch.Tensor"],

transformers/models/cohere/modeling_cohere.py

@@ -36,6 +36,7 @@ from torch import nn
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
+from ...integrations import use_kernelized_func
 from ...masking_utils import create_causal_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
@@ -44,7 +45,7 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
-from ...utils.generic import check_model_inputs
+from ...utils.generic import check_model_inputs, maybe_autocast
 from .configuration_cohere import CohereConfig
 
 
@@ -121,7 +122,7 @@ class CohereRotaryEmbedding(nn.Module):
         position_ids_expanded = position_ids[:, None, :].float()
 
         device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
             freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             emb = torch.repeat_interleave(freqs, 2, dim=-1)  # diff from Llama: we interleave() instead of cat()
             cos = emb.cos() * self.attention_scaling
@@ -222,6 +223,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     return q_embed.to(dtype=dtype), k_embed.to(dtype=dtype)
 
 
+@use_kernelized_func(apply_rotary_pos_emb)
 class CohereAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
@@ -247,7 +249,6 @@ class CohereAttention(nn.Module):
         self.o_proj = nn.Linear(
             config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
         )
-        self.rotary_fn = apply_rotary_pos_emb
         self.use_qk_norm = config.use_qk_norm
         if self.use_qk_norm:
             # When sharding the model using Tensor Parallelism, need to be careful to use n_local_heads

transformers/models/cohere/modular_cohere.py

@@ -36,6 +36,7 @@ from ...modeling_rope_utils import dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, logging
+from ...utils.generic import maybe_autocast
 from ..llama.modeling_llama import (
     LlamaAttention,
     LlamaForCausalLM,
@@ -75,7 +76,7 @@ class CohereRotaryEmbedding(LlamaRotaryEmbedding):
         position_ids_expanded = position_ids[:, None, :].float()
 
         device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
             freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             emb = torch.repeat_interleave(freqs, 2, dim=-1)  # diff from Llama: we interleave() instead of cat()
             cos = emb.cos() * self.attention_scaling