repo_description: 🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training. 
repo_topics: audio deep-learning deepseek gemma glm hacktoberfest llm machine-learning model-hub natural-language-processing nlp pretrained-models python pytorch pytorch-transformers qwen speech-recognition transformer vlm
code: # Copyright 2026 The HuggingFace Inc. team.
  #
  # 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.
  
  import re
  import sys
  import threading
  from contextlib import contextmanager
  
  from .utils import is_torch_available, logging
  from .utils.output_capturing import OutputRecorder
  
  
  if is_torch_available():
      import torch.nn as nn
  
  logger = logging.get_logger(__name__)
  
  _monkey_patch_mapping_cache: dict[str, type[nn.Module]] = {}
  _compiled_patterns_cache: dict[str, re.Pattern] = {}
  _monkey_patch_lock = threading.Lock()
  
  
  def _compile_pattern(pattern: str) -> re.Pattern | None:
      """
      Compile a regex pattern and cache it. Returns None if pattern is invalid.
  
      Args:
          pattern: The regex pattern string to compile
  
      Returns:
          Compiled regex pattern or None if invalid
      """
      if pattern in _compiled_patterns_cache:
          return _compiled_patterns_cache[pattern]
  
      try:
          compiled = re.compile(pattern)
          _compiled_patterns_cache[pattern] = compiled
          return compiled
      except re.error as e:
          logger.warning(f"Invalid regex pattern '{pattern}': {e}. Treating as non-pattern.")
          return None
  
  
  def _find_replacement_class(class_name: str, mapping: dict[str, type[nn.Module]]) -> type[nn.Module] | None:
      """
      Find replacement class for a given class name, checking exact matches first, then regex patterns.
  
      Args:
          class_name: The class name to find a replacement for
          mapping: Dictionary of patterns/names to replacement classes
  
      Returns:
          The replacement class if found, None otherwise
      """
      # First check for exact match (highest priority)
      if class_name in mapping:
          return mapping[class_name]
  
      # Then check regex patterns
      for pattern, replacement_class in mapping.items():
          # Skip if already matched as exact
          if pattern == class_name:
              continue
  
          # Try to compile and match as regex
          compiled_pattern = _compile_pattern(pattern)
          if compiled_pattern is not None and compiled_pattern.search(class_name):
              return replacement_class
  
      return None
  
  
  def register_patch_mapping(mapping: dict[str, type[nn.Module]], overwrite: bool = False) -> None:
      """
      Register patch mappings to enable automatic patching during model creation using `from_pretrained`,
      `from_config` or within the `apply_patches` context manager.
  
      Use this to register class replacements that will be automatically applied when loading any model.
      This is useful for quantization library compatibility, structural optimizations, and architectural
      experimentation. The mapping is global, can grow with multiple calls, and can be cleared entirely.
  
      Args:
          mapping (`Dict[str, type[nn.Module]]`):
              Mapping from original class names (or regex patterns) to replacement classes. Supports:
              - Exact class names: `"Qwen2MoeExperts"` → `CustomExperts`
              - Regex patterns: `".*Attention"` matches `LlamaAttention`, `MistralAttention`, etc.,
              or `"^Llama\\d+Attention$"` matches `Llama2Attention`, `Llama3Attention`, etc.
  
              Exact matches take precedence over patterns. Patterns are matched using `re.search()`,
              so they can match anywhere in the class name unless you use anchors (`^` for start, `$` for end).
          overwrite (`bool`, *optional*, defaults to `False`):
              Whether to overwrite existing mappings for class names that are already registered.
  
      Example:
          ```python
          from transformers import AutoModelForCausalLM
          from transformers.monkey_patching import register_patch_mapping
  
          # Define custom expert implementation
          class SequentialExperts(nn.Module):
              ...
  
          # Register exact class name
          register_patch_mapping(
              mapping={"Qwen2MoeExperts": SequentialExperts}
          )
  
          # Register with regex pattern to match multiple classes
          register_patch_mapping(
              mapping={".*Attention": CustomAttention}  # Matches LlamaAttention, MistralAttention, etc.
          )
  
          # Match specific model versions
          register_patch_mapping(
              mapping={"^Llama\\d+Attention$": CustomLlamaAttention}  # Matches Llama2Attention, Llama3Attention
          )
  
          # The patch will be automatically applied during loading
          model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
          ```
  
      Note:
          For weight conversions, use [`~transformers.register_checkpoint_conversion_mapping`] instead.
      """
      global _monkey_patch_mapping_cache
      with _monkey_patch_lock:
          for class_name, replacement_class in mapping.items():
              # Validate that replacement_class is actually a class and is a subclass of nn.Module
              if not isinstance(replacement_class, type):
                  raise TypeError(
                      f"Replacement for '{class_name}' must be a class, got {type(replacement_class).__name__}"
                  )
              if not issubclass(replacement_class, nn.Module):
                  raise TypeError(
                      f"Replacement class for '{class_name}' must be a subclass of nn.Module, "
                      f"got {replacement_class.__name__} which inherits from {[c.__name__ for c in replacement_class.__mro__[1:]]}"
                  )
  
              if class_name in _monkey_patch_mapping_cache and not overwrite:
                  raise ValueError(
                      f"Class '{class_name}' already has a patch mapping registered. Use overwrite=True to replace it."
                  )
              _monkey_patch_mapping_cache[class_name] = replacement_class
  
  
  def unregister_patch_mapping(keys: list[str]) -> None:
      """
      Unregister patch mappings to disable automatic patching.
  
      This removes specified mappings from the global registry, preventing them from being applied
      during model loading. You must provide the exact same name or pattern that was used during registration.
  
      Args:
          keys (`List[str]`):
              List of mapping keys (class names or regex patterns) to remove from the patch mapping
              (e.g., `["Qwen2MoeExperts"]` or `[".*Attention"]`).
  
      Example:
          ```python
          from transformers import AutoModelForCausalLM
          from transformers.monkey_patching import register_patch_mapping, unregister_patch_mapping
  
          # Register a patch
          register_patch_mapping(
              mapping={"Qwen2MoeExperts": CustomExperts}
          )
  
          # Unregister the patch
          unregister_patch_mapping(["Qwen2MoeExperts"])
  
          # The patch will no longer be applied during loading
          model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-MoE-A2.7B")
          ```
      """
      global _monkey_patch_mapping_cache
      with _monkey_patch_lock:
          for key in keys:
              if key not in _monkey_patch_mapping_cache:
                  raise ValueError(
                      f"Class or pattern '{key}' not found in monkey patch mapping cache. "
                      f"Cannot unregister a class that is not registered."
                  )
              del _monkey_patch_mapping_cache[key]
  
  
  def get_patch_mapping() -> dict[str, type[nn.Module]]:
      """
      Get all registered patch mappings.
  
      Returns:
          `Dict[str, type[nn.Module]]`: Dictionary mapping class names or patterns to replacement classes.
      """
      with _monkey_patch_lock:
          return _monkey_patch_mapping_cache.copy()
  
  
  def clear_patch_mapping() -> None:
      """
      Clear all registered patch mappings.
  
      This removes all registered mappings from the global registry.
  
      Example:
          ```python
          from transformers.monkey_patching import register_patch_mapping, clear_patch_mapping
  
          # Register some patches
          register_patch_mapping(
              mapping={"Qwen2MoeExperts": CustomExperts}
          )
  
          # Clear all patches
          clear_patch_mapping()
          ```
      """
      global _monkey_patch_mapping_cache
      with _monkey_patch_lock:
          _monkey_patch_mapping_cache.clear()
  
  
  @contextmanager
  def apply_patches():
      """
      Context manager to apply registered monkey patches within a block of code.
  
      This temporarily replaces original classes with their registered replacements during the execution of the block, and restores the original classes afterward.
  
      Example:
          ```python
          from transformers import Qwen2MoeModel, Qwen2MoeConfig
          from transformers.monkey_patching import register_patch_mapping, apply_patches
  
          # Register a patch
          register_patch_mapping(
              mapping={"Qwen2MoeExperts": CustomExperts}
          )
  
          # Apply patches within the context
          with apply_patches():
              # The model will use CustomExperts instead of Qwen2MoeExperts
              model = Qwen2MoeModel(Qwen2MoeConfig())
  
          # Outside the context, original classes are restored
          # The model will use Qwen2MoeExperts again
          model = Qwen2MoeModel(Qwen2MoeConfig())
          ```
      """
      mapping = get_patch_mapping()
      if not mapping:
          yield
          return
  
      original_classes = {}
  
      # Create list to avoid dict changed during iteration
      for module in list(sys.modules.values()):
          if module is None or not hasattr(module, "__name__"):
              continue
          if not module.__name__.startswith("transformers"):
              continue
  
          # Iterate through all attributes in transformers modules
          for attr_name in dir(module):
              # Check if this attribute name matches any pattern before accessing it
              replacement_class = _find_replacement_class(attr_name, mapping)
              if replacement_class is None:
                  continue
  
              try:
                  attr = getattr(module, attr_name)
                  # Check if it's a class
                  if not isinstance(attr, type):
                      continue
  
                  original_classes[(module.__name__, attr_name)] = attr
                  setattr(module, attr_name, replacement_class)
              except (AttributeError, TypeError, ImportError):
                  # Skip attributes that can't be accessed or modules that can't be imported
                  continue
  
      yield
  
      for (module_name, class_name), original_class in original_classes.items():
          module = sys.modules[module_name]
          setattr(module, class_name, original_class)
  
  
  # _can_record_outputs is a class attribute so patching and unpatching it in the class won't work
  # since the model instance will still reference the original class's _can_record_outputs.
  def patch_output_recorders(model: nn.Module) -> None:
      """
      Patch the model instance's output recorders to use the registered replacement classes.
  
      This function updates output recorders in a model's submodules to use monkey-patched replacement
      classes. Output recorders are used by the transformers library to track intermediate outputs during
      forward passes (via the `_can_record_outputs` attribute). When classes are monkey-patched, these
      recorders need to be updated to reference the new classes.
  
      This is automatically called during model initialization when loading with `from_pretrained` or
      `from_config`. You typically don't need to call this manually unless you're constructing models
      in custom ways.
  
      Note:
          The `_can_record_outputs` attribute is a class-level attribute that maps output names to either:
          - `OutputRecorder` instances that have a `target_class` attribute
          - Class types directly
  
          This function patches both cases to use the replacement classes from the monkey patch registry.
  
      Args:
          model (`nn.Module`):
              The model instance whose output recorders should be patched. All submodules will be
              traversed to find and patch their `_can_record_outputs` attributes.
  
      Example:
          ```python
          from transformers import AutoModelForCausalLM
          from transformers.monkey_patching import register_patch_mapping, patch_output_recorders
  
          # Register a patch
          register_patch_mapping(mapping={"Qwen2MoeExperts": CustomExperts})
  
          # If you construct a model manually (without from_pretrained), patch recorders
          model = Qwen2MoeModel(config)
          patch_output_recorders(model)  # Updates output recorders to use CustomExperts
          ```
      """
  
      mapping = get_patch_mapping()
      if not mapping:
          return
  
      for submodule in model.modules():
          if hasattr(submodule, "_can_record_outputs") and submodule._can_record_outputs is not None:
              for output, recorder in submodule._can_record_outputs.items():
                  if isinstance(recorder, OutputRecorder):
                      # Check if target class matches any registered pattern or exact name
                      replacement_class = _find_replacement_class(recorder.target_class.__name__, mapping)
                      if replacement_class is not None:
                          recorder.target_class = replacement_class
                  elif isinstance(recorder, type):
                      # Check if class type matches any registered pattern or exact name
                      replacement_class = _find_replacement_class(recorder.__name__, mapping)
                      if replacement_class is not None:
                          submodule._can_record_outputs[output] = replacement_class
  