search

vade / OpenAI-Whisper-CoreML

OpenAI's Whisper ported to CoreML
144 stars 3 forks source link

Set max token length during CoreML export #9

Open rsomani95 opened 1 year ago

rsomani95 commented 1 year ago

Currently, there is no upper bound being set to the number of tokens: https://github.com/vade/OpenAI-Whisper-CoreML/blob/2c60f5cb67dee05adbfaf1f31a12595dfbe75789/whisper_to_cml.py#L34-L35

There is an internal limit to the number of tokens that can be passed in. I believe it is 448 for all variants, but to be safer, you can fetch it directly from the decoder itself, like so:

max_tokens = model.positional_embedding.shape[0]
token_flexible_shape = ct.Shape((
    1,
    ct.RangeDim(lower_bound=1, upper_bound=max_tokens, default=1),
))