code: def _inject_drugged_attention(
      cls, model: PreTrainedModel, **kwargs
  ) -> Optional[int]:
      model_type = model.config.model_type
      from drugs.models import (
          llama_drugged_sdpa_forward,
          llama_drugged_flash2_attention_forward,
          llama_drugged_attention_forward,
          mistral_drugged_attention_forward,
      )
  
      LLAMA_ATTENTION_FORWARDS = {
          "LlamaAttention": llama_drugged_attention_forward,
          "LlamaFlashAttention2": llama_drugged_flash2_attention_forward,
          "LlamaSdpaAttention": llama_drugged_sdpa_forward,
      }
      MISTRAL_ATTENTION_FORWARDS = {
          "MistralAttention": mistral_drugged_attention_forward,
          # "MistralFlashAttention2" : mistral_drugged_flash2_attention_forward,
          # "MistralSdpaAttention": mistral_drugged_sdpa_forward
      }
      ATTENTION_NAME_MAPPINGS = {
          "llama": LLAMA_ATTENTION_FORWARDS,
          "mistral": MISTRAL_ATTENTION_FORWARDS,
      }
  
      # Not all models require updated attention forwards
      if (
          model_type not in MODEL_NAME_MAPPING
          or MODEL_NAME_MAPPING[model_type] is None
      ):
          return
  
      def overwrite_forward(module, new_func, **kwargs) -> None:
          module.forward = types.MethodType(new_func, module, **kwargs)
  
      return cls._call_modules_by_name(
          model,
          ATTENTION_NAME_MAPPINGS[model_type],
          overwrite_forward,
          **kwargs,
      )
  
file_path: theta-layer-drugs/drugs/inject_mixin.py
imports: import types
  from typing import Callable, Optional, List
  import torch
  from transformers import PreTrainedModel
  from transformers.utils import logging
  from drugs.generation.utils import _update_model_kwargs_for_generation
  from drugs.models import *
label_desc: Use image node_id 1 for calling the InjectDrugsMixin obj's underlying member method code with example usage: obj._inject_drugged_attention(cls, model) and returns: cls