yobx.torch.in_transformers.patches#

yobx.torch.in_transformers.patches.enable_transformers_onnx_export_flags(model: Module | None = None, verbose: int = 0) → Generator[Dict[str, Any], None, None][source]#

Context manager that enables the ONNX-export flags that are already implemented inside the transformers library and restores them on exit.

Two kinds of flags are handled:

  • config.onnx_export (and any sub-config that exposes it, such as LongformerConfig or LEDConfig) is set to True.

  • Every submodule that implements prepare_for_onnx_export_ (for example ProphetNetNgramSelfAttention) has the method called on it, and its onnx_trace attribute is restored on exit.

Calling this with model=None is a no-op and yields an empty report.

Parameters:

  • model – the model whose flags should be toggled

  • verbose – prints out a line every time a flag is toggled

Returns:

yields a dict with two keys configs and modules listing the configs and submodules that were touched.

yobx.torch.in_transformers.patches.get_patches_for(model: Module | None = None) → List[PatchInfo][source]#

Returns the list of patches for a specific model. if model is None, patches everything it can.

Note

The function detects that RotaryEmbedding.forward is wrapped by checking if can find substring transformers/modeling_rope_utils.py in RotaryEmbedding.forward.__wrapped__. It does not seem to be the case with Python 3.10.