Retrieve alignment_heads of a fine-tuned model

[Ar770]

Hi,
I'm looking for a method to retrieve the alignment_heads of a fine-tuned Whisper model.
I need it as I understand that the word-level timestamps depend on them.
Thanks

[Ar770]

@jongwook can you please help with this?

[jongwook]

I think you can keep using the same alignment heads as the base model. If you're curious, save the QKs tensor in timing.py:

whisper/whisper/timing.py

Line 205 in f6f01c5

weights = torch.stack([QKs[_l][_h] for _l, _h in model.alignment_heads.indices().T])

and manually inspect the heads (that's how I made the masks at least...).