|
| 1 | +"""Local workarounds for known Transformers bugs, applied by the runner before contracts run. |
| 2 | +
|
| 3 | +A framework's fleet lock lists the patches it needs under `local_patches`, with |
| 4 | +the reason. The runner applies them in its own process and in every isolated |
| 5 | +contract process, and records their names in each result's environment, so a |
| 6 | +baseline recorded with a workaround says so. Remove a patch, and re-record the |
| 7 | +affected baselines, once the fix is in the pinned Transformers build. |
| 8 | +""" |
| 9 | + |
| 10 | +import importlib |
| 11 | +import inspect |
| 12 | + |
| 13 | +# Audio models whose `masked_spec_embed` (SpecAugment's mask vector) has no |
| 14 | +# `_init_weights` branch. Checkpoints that lack the tensor get uninitialized |
| 15 | +# memory on load; HuBERT, SEW-D, SpeechT5, and Wav2Vec2-BERT already handle it. |
| 16 | +MASKED_SPEC_EMBED_MODELS = { |
| 17 | + "wav2vec2": "Wav2Vec2PreTrainedModel", |
| 18 | + "wav2vec2_conformer": "Wav2Vec2ConformerPreTrainedModel", |
| 19 | + "data2vec": "Data2VecAudioPreTrainedModel", |
| 20 | + "wavlm": "WavLMPreTrainedModel", |
| 21 | + "unispeech": "UniSpeechPreTrainedModel", |
| 22 | + "unispeech_sat": "UniSpeechSatPreTrainedModel", |
| 23 | + "sew": "SEWPreTrainedModel", |
| 24 | +} |
| 25 | + |
| 26 | + |
| 27 | +def masked_spec_embed_init(): |
| 28 | + """Initialize masked_spec_embed like HuBERT does: init.uniform_ in _init_weights.""" |
| 29 | + import torch |
| 30 | + from transformers import initialization as init |
| 31 | + |
| 32 | + patched = [] |
| 33 | + for model_type, class_name in MASKED_SPEC_EMBED_MODELS.items(): |
| 34 | + module_name = f"transformers.models.{model_type}.modeling_{'data2vec_audio' if model_type == 'data2vec' else model_type}" |
| 35 | + try: |
| 36 | + cls = getattr(importlib.import_module(module_name), class_name) |
| 37 | + except (ImportError, AttributeError): |
| 38 | + continue |
| 39 | + original = cls._init_weights |
| 40 | + if getattr(original, "_mic_patched", False): |
| 41 | + continue |
| 42 | + if "masked_spec_embed" in inspect.getsource(original): |
| 43 | + continue # fixed upstream (#49386): a second init.uniform_ would draw different values |
| 44 | + |
| 45 | + def _init_weights(self, module, original=original): |
| 46 | + original(self, module) |
| 47 | + if isinstance(getattr(module, "masked_spec_embed", None), torch.nn.Parameter): |
| 48 | + init.uniform_(module.masked_spec_embed) |
| 49 | + |
| 50 | + _init_weights._mic_patched = True |
| 51 | + cls._init_weights = _init_weights |
| 52 | + patched.append(class_name) |
| 53 | + return patched |
| 54 | + |
| 55 | + |
| 56 | +def whisper_decode_tensor_batch(): |
| 57 | + """Decode a 2-D tensor with timestamps per sequence, as decode already does for a list of lists (#49403). |
| 58 | +
|
| 59 | + WhisperTokenizer.decode only treats its input as a batch when it is a list of |
| 60 | + lists, so batch_decode(model.generate(...), decode_with_timestamps=True) merges |
| 61 | + the rows into one string (one token per row) or raises (longer rows). |
| 62 | + """ |
| 63 | + try: |
| 64 | + from transformers.models.whisper.tokenization_whisper import WhisperTokenizer |
| 65 | + except ImportError: |
| 66 | + return [] |
| 67 | + original = WhisperTokenizer.decode |
| 68 | + if getattr(original, "_mic_patched", False): |
| 69 | + return [] |
| 70 | + |
| 71 | + def decode(self, token_ids, *args, original=original, **kwargs): |
| 72 | + if kwargs.get("decode_with_timestamps") and getattr(token_ids, "ndim", 0) == 2: |
| 73 | + token_ids = token_ids.tolist() |
| 74 | + return original(self, token_ids, *args, **kwargs) |
| 75 | + |
| 76 | + decode._mic_patched = True |
| 77 | + WhisperTokenizer.decode = decode |
| 78 | + return ["WhisperTokenizer"] |
| 79 | + |
| 80 | + |
| 81 | +def config_field_generation_params(): |
| 82 | + """Do not count a config's own declared fields as user-set generation parameters (#49456, from #48282). |
| 83 | +
|
| 84 | + generate() raises when model.config holds generation parameters. It used to |
| 85 | + skip every field of the config class; since #48282 it skips only fields |
| 86 | + whose default is not None, so WhisperConfig.suppress_tokens (default None, |
| 87 | + set in the OpenAI checkpoints' config.json) makes every generate() raise. |
| 88 | + Restore the skip for every declared field. |
| 89 | + """ |
| 90 | + import dataclasses |
| 91 | + |
| 92 | + from transformers import PreTrainedConfig |
| 93 | + |
| 94 | + original = PreTrainedConfig._get_generation_parameters |
| 95 | + if getattr(original, "_mic_patched", False) or not hasattr(PreTrainedConfig, "default_config_fields"): |
| 96 | + return [] # already patched, or a build before the regression |
| 97 | + |
| 98 | + def _get_generation_parameters(self, original=original): |
| 99 | + params = original(self) |
| 100 | + if dataclasses.is_dataclass(self): |
| 101 | + declared = {f.name for f in dataclasses.fields(self)} |
| 102 | + params = {k: v for k, v in params.items() if k not in declared} |
| 103 | + return params |
| 104 | + |
| 105 | + _get_generation_parameters._mic_patched = True |
| 106 | + PreTrainedConfig._get_generation_parameters = _get_generation_parameters |
| 107 | + return ["PreTrainedConfig"] |
| 108 | + |
| 109 | +PATCHES = { |
| 110 | + "masked-spec-embed-init": masked_spec_embed_init, |
| 111 | + "whisper-decode-tensor-batch": whisper_decode_tensor_batch, |
| 112 | + "config-field-generation-params": config_field_generation_params, |
| 113 | +} |
| 114 | + |
| 115 | + |
| 116 | +def apply(names): |
| 117 | + """Apply the named patches; return {name: what was patched}.""" |
| 118 | + unknown = sorted(set(names) - set(PATCHES)) |
| 119 | + if unknown: |
| 120 | + raise SystemExit(f"unknown local patches: {unknown}") |
| 121 | + return {name: PATCHES[name]() for name in names} |
0 commit comments