Providing Models in SafeTensors format

[ohensf]

Hi,

We are currently utilizing models in our project stored in pickle format. However, after exploring the advantages of SafeTensors in terms of improved security we believe that it will provide us with an extra layer of security and protection against potential vulnerabilities.
model link: base, base.en

We wish to inquire about the plans or timeline, if any, regarding the update of models from pickle to SafeTensors:

• Are there any plans to update the models to SafeTensors?
• If yes, is there a projected timeline for this migration?

[phineas-pta]

u can do it yourself, not too much difficult

simple code

convert to SafeTensors:

import whisper
import safetensors.torch

model = whisper.load_model("large")
safetensors.torch.save_model(model, "large.safetensors")

load from SafeTensors:

import whisper
import safetensors.torch

model = whisper.model.Whisper(whisper.model.ModelDimensions(
    n_mels        = 80,
    n_vocab       = 51865,
    n_audio_ctx   = 1500,
    n_audio_state = 1280,
    n_audio_head  = 20,
    n_audio_layer = 32,
    n_text_ctx    = 448,
    n_text_state  = 1280,
    n_text_head   = 20,
    n_text_layer  = 32
))
# model.to("cuda")

safetensors.torch.load_model(model, "large.safetensors")
model.set_alignment_heads(b"ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj")

# model.transcribe(...)

models info

value for alignment heads can be found in https://github.com/openai/whisper/blob/main/whisper/__init__.py

_ALIGNMENT_HEADS = {
    "tiny.en"  : b"ABzY8J1N>@0{>%R00Bk>$p{7v037`oCl~+#00",
    "tiny"     : b"ABzY8bu8Lr0{>%RKn9Fp%m@SkK7Kt=7ytkO",
    "base.en"  : b"ABzY8;40c<0{>%RzzG;p*o+Vo09|#PsxSZm00",
    "base"     : b"ABzY8KQ!870{>%RzyTQH3`Q^yNP!>##QT-<FaQ7m",
    "small.en" : b"ABzY8>?_)10{>%RpeA61k&I|OI3I$65C{;;pbCHh0B{qLQ;+}v00",
    "small"    : b"ABzY8DmU6=0{>%Rpa?J`kvJ6qF(V^F86#Xh7JUGMK}P<N0000",
    "medium.en": b"ABzY8usPae0{>%R7<zz_OvQ{)4kMa0BMw6u5rT}kRKX;$NfYBv00*Hl@qhsU00",
    "medium"   : b"ABzY8B0Jh+0{>%R7}kK1fFL7w6%<-Pf*t^=N)Qr&0RR9",
    "large-v1" : b"ABzY8r9j$a0{>%R7#4sLmoOs{s)o3~84-RPdcFk!JR<kSfC2yj",
    "large-v2" : b"ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj",
    "large"    : b"ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj",
}

value for dimensions can be found using torch.load("PATH to checkpoint like ~/.cache/whisper/large-v2.pt")["dims"]

here the values:

	tiny	base	small	medium	large
n_mels	80
n_vocab	51865 if multilingual models 51864 if english-only models
n_audio_ctx	1500
n_audio_state	384	512	768	1024	1280
n_audio_head	6	8	12	16	20
n_audio_layer	4	6	12	24	32
n_text_ctx	448
n_text_state	same as n_audio_state
n_text_head	same as n_audio_head
n_text_layer	same as n_audio_layer

[ohensf]

Great, it is really helpful. 🙏 thanks!

[ohensf]

@phineas-pta One thing that we notice is that when running the pickle version vs. the SafeTensors version, we are sometimes getting different results (one word, sometimes a few...).
Do you know if the SafeTensors conversion should impact the quality of the results? Or is it just a matter of configuration?

[phineas-pta]

#81

[srezasm]

How should I set the whisper.model.ModelDimensions parameters for small model?

[phineas-pta]

huggingface also created safetensor version of whisper, but to be used with their pipeline interface