mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-06 20:43:46 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
333 lines
17 KiB
Python
333 lines
17 KiB
Python
# thanks to https://github.com/openai/whisper for a good chunk of MIT licensed code
|
|
|
|
import sys, base64, multiprocessing, itertools, collections
|
|
from typing import Optional, Union, Literal, List
|
|
|
|
from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
|
|
from tinygrad.nn.state import torch_load, load_state_dict
|
|
from tinygrad.helpers import getenv, fetch
|
|
|
|
from examples.audio_helpers import mel
|
|
import numpy as np
|
|
import librosa
|
|
|
|
class MultiHeadAttention:
|
|
def __init__(self, n_state, n_head, kv_caching: Literal['cross', 'self']=None, max_self_attn_cache_len=None):
|
|
self.n_head = n_head
|
|
self.query = nn.Linear(n_state, n_state)
|
|
self.key = nn.Linear(n_state, n_state, bias=False)
|
|
self.value = nn.Linear(n_state, n_state)
|
|
self.out = nn.Linear(n_state, n_state)
|
|
|
|
self.kv_caching = kv_caching
|
|
self.max_self_attn_cache_len = max_self_attn_cache_len
|
|
|
|
def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None, len: Union[Variable,int]=None):
|
|
if self.kv_caching == 'cross':
|
|
if xa is not None:
|
|
k, v = self.key(xa), self.value(xa)
|
|
if not hasattr(self, 'cache_k'):
|
|
self.cache_k, self.cache_v = k, v
|
|
else:
|
|
self.cache_k.assign(k).realize()
|
|
self.cache_v.assign(v).realize()
|
|
else:
|
|
k, v = self.cache_k, self.cache_v
|
|
else:
|
|
k, v = self.key(x), self.value(x)
|
|
if self.kv_caching == 'self':
|
|
if not hasattr(self, 'cache_k'):
|
|
self.cache_k = Tensor.zeros(x.shape[0], self.max_self_attn_cache_len, x.shape[2])
|
|
self.cache_v = Tensor.zeros(x.shape[0], self.max_self_attn_cache_len, x.shape[2])
|
|
k = self.cache_k.shrink((None, (0, len), None)).cat(k, dim=1)
|
|
v = self.cache_v.shrink((None, (0, len), None)).cat(v, dim=1)
|
|
padding = self.max_self_attn_cache_len-len-x.shape[1]
|
|
self.cache_k.assign(k.pad((None, (0, padding), None)).contiguous()).realize()
|
|
self.cache_v.assign(v.pad((None, (0, padding), None)).contiguous()).realize()
|
|
|
|
q = self.query(x)
|
|
n_ctx = q.shape[1]
|
|
assert(q.shape[-1] == k.shape[-1] == v.shape[-1])
|
|
head_dim = q.shape[-1] // self.n_head
|
|
q = q.reshape(*q.shape[:2], self.n_head, head_dim).permute(0, 2, 1, 3)
|
|
k = k.reshape(*k.shape[:2], self.n_head, head_dim).permute(0, 2, 1, 3)
|
|
v = v.reshape(*v.shape[:2], self.n_head, head_dim).permute(0, 2, 1, 3)
|
|
attn = Tensor.scaled_dot_product_attention(q, k, v, mask[:n_ctx,:n_ctx] if mask is not None else None)
|
|
wv = attn.permute(0, 2, 1, 3).flatten(start_dim=2)
|
|
return self.out(wv)
|
|
|
|
|
|
class ResidualAttentionBlock:
|
|
def __init__(self, n_state, n_head, is_decoder_block=False, max_self_attn_cache_len=None):
|
|
self.attn = MultiHeadAttention(n_state, n_head, kv_caching='self' if is_decoder_block else None, max_self_attn_cache_len=max_self_attn_cache_len)
|
|
self.attn_ln = nn.LayerNorm(n_state)
|
|
|
|
self.cross_attn = MultiHeadAttention(n_state, n_head, kv_caching='cross') if is_decoder_block else None
|
|
self.cross_attn_ln = nn.LayerNorm(n_state) if is_decoder_block else None
|
|
|
|
self.mlp = [nn.Linear(n_state, n_state*4), Tensor.gelu, nn.Linear(n_state*4, n_state)]
|
|
self.mlp_ln = nn.LayerNorm(n_state)
|
|
|
|
def __call__(self, x, xa=None, mask=None, len: Union[Variable, int]=None):
|
|
x = x + self.attn(self.attn_ln(x), mask=mask, len=len)
|
|
if self.cross_attn: x = x + self.cross_attn(self.cross_attn_ln(x), xa)
|
|
x = x + self.mlp_ln(x).sequential(self.mlp)
|
|
return x.realize()
|
|
|
|
class AudioEncoder:
|
|
def __init__(self, n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer, **_):
|
|
self.conv1 = nn.Conv1d(n_mels, n_audio_state, kernel_size=3, padding=1)
|
|
self.conv2 = nn.Conv1d(n_audio_state, n_audio_state, kernel_size=3, stride=2, padding=1)
|
|
self.blocks = [ResidualAttentionBlock(n_audio_state, n_audio_head) for _ in range(n_audio_layer)]
|
|
self.ln_post = nn.LayerNorm(n_audio_state)
|
|
self.positional_embedding = Tensor.empty(n_audio_ctx, n_audio_state)
|
|
self.encode = TinyJit(self.__call__)
|
|
|
|
def __call__(self, x):
|
|
x = self.conv1(x).gelu()
|
|
x = self.conv2(x).gelu()
|
|
x = x.permute(0, 2, 1)
|
|
x = x + self.positional_embedding[:x.shape[1]]
|
|
x = x.sequential(self.blocks)
|
|
x = self.ln_post(x)
|
|
return x.realize()
|
|
|
|
class TextDecoder:
|
|
def __init__(self, n_vocab, n_text_ctx, n_text_state, n_text_head, n_text_layer, **_):
|
|
self.max_tokens_to_sample = n_text_ctx // 2
|
|
self.max_self_attn_cache_len = n_text_ctx
|
|
|
|
self.token_embedding = nn.Embedding(n_vocab, n_text_state)
|
|
self.positional_embedding = Tensor.empty(n_text_ctx, n_text_state)
|
|
self.blocks = [ResidualAttentionBlock(n_text_state, n_text_head, is_decoder_block=True, max_self_attn_cache_len=self.max_self_attn_cache_len) for _ in range(n_text_layer)]
|
|
self.ln = nn.LayerNorm(n_text_state)
|
|
self.mask = Tensor.full((n_text_ctx, n_text_ctx), -np.inf).triu(1).realize()
|
|
self.getjitted = collections.defaultdict(lambda: TinyJit(self.forward))
|
|
|
|
def __call__(self, x: Tensor, pos: int, encoded_audio: Tensor):
|
|
pos = Variable("self_attn_cache_len", 1, self.max_self_attn_cache_len-1).bind(pos) if pos else 0
|
|
return self.getjitted[x.shape](x, pos, encoded_audio)
|
|
|
|
def forward(self, x:Tensor, pos:Union[Variable, Literal[0]], encoded_audio:Tensor):
|
|
seqlen = x.shape[-1]
|
|
x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None))
|
|
for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos)
|
|
return self.output_tok(x)
|
|
|
|
def output_tok(self, x):
|
|
return (self.ln(x) @ self.token_embedding.weight.T).realize()
|
|
|
|
class Whisper:
|
|
def __init__(self, dims, batch_size=1):
|
|
self.encoder = AudioEncoder(**dims)
|
|
self.decoder = TextDecoder(**dims)
|
|
self.is_multilingual = dims["n_vocab"] == 51865
|
|
self.batch_size = batch_size
|
|
|
|
|
|
RATE = 16000
|
|
SEGMENT_SECONDS=30
|
|
SAMPLES_PER_SEGMENT = RATE * SEGMENT_SECONDS # 480000
|
|
N_FFT = 400
|
|
HOP_LENGTH = 160
|
|
N_MELS = 80
|
|
FRAMES_PER_SEGMENT = SAMPLES_PER_SEGMENT // HOP_LENGTH # 3000
|
|
|
|
def prep_audio(waveforms: List[np.ndarray], batch_size: int, truncate=False) -> np.ndarray:
|
|
"""
|
|
:param waveforms: A list of possibly variable length 16000Hz audio samples
|
|
:param batch_size: The batch_size associated with the Whisper model being used to transcribe the audio.
|
|
Used to prevent JIT mismatch errors since the encoder does not accept symbolic shapes
|
|
:param truncate: If true, truncates (or pads) audio to exactly 30s for a single encoder pass
|
|
:return: mel spectrogram of the given waveforms
|
|
"""
|
|
def pad_or_trim(arr, target_len):
|
|
curr_len = len(arr)
|
|
if curr_len == target_len:
|
|
return arr
|
|
elif curr_len < target_len:
|
|
return np.pad(arr, (0, target_len - curr_len), 'constant')
|
|
else:
|
|
return arr[:target_len]
|
|
|
|
max_len = SAMPLES_PER_SEGMENT if truncate else max(len(wav) for wav in waveforms)
|
|
if (r := max_len % SAMPLES_PER_SEGMENT) > 0: max_len += SAMPLES_PER_SEGMENT - r
|
|
waveforms = np.array(list(map(lambda w: pad_or_trim(w, max_len), waveforms)))
|
|
assert waveforms.shape[0] <= batch_size
|
|
if waveforms.shape[0] < batch_size:
|
|
# we could have a symbolic batch_size dim instead of manually padding here if conv/layernorm supported symbolic shapes
|
|
waveforms = np.pad(waveforms, pad_width=((0, batch_size - waveforms.shape[0]), (0, 0)))
|
|
|
|
stft = librosa.stft(waveforms, n_fft=N_FFT, hop_length=HOP_LENGTH, window='hann', dtype=np.csingle)
|
|
magnitudes = np.absolute(stft[..., :-1]) ** 2
|
|
mel_spec = mel(sr=RATE, n_fft=N_FFT, n_mels=N_MELS).numpy() @ magnitudes
|
|
|
|
log_spec = np.log10(np.clip(mel_spec, 1e-10, None))
|
|
log_spec = np.maximum(log_spec, log_spec.max((1,2), keepdims=True) - 8.0)
|
|
log_spec = (log_spec + 4.0) / 4.0
|
|
|
|
return log_spec
|
|
|
|
LANGUAGES = {
|
|
"en": "english", "zh": "chinese", "de": "german", "es": "spanish", "ru": "russian", "ko": "korean", "fr": "french", "ja": "japanese", "pt": "portuguese", "tr": "turkish",
|
|
"pl": "polish", "ca": "catalan", "nl": "dutch", "ar": "arabic", "sv": "swedish", "it": "italian", "id": "indonesian", "hi": "hindi", "fi": "finnish", "vi": "vietnamese",
|
|
"he": "hebrew", "uk": "ukrainian", "el": "greek", "ms": "malay", "cs": "czech", "ro": "romanian", "da": "danish", "hu": "hungarian", "ta": "tamil", "no": "norwegian",
|
|
"th": "thai", "ur": "urdu", "hr": "croatian", "bg": "bulgarian", "lt": "lithuanian", "la": "latin", "mi": "maori", "ml": "malayalam", "cy": "welsh", "sk": "slovak", "te": "telugu",
|
|
"fa": "persian", "lv": "latvian", "bn": "bengali", "sr": "serbian", "az": "azerbaijani", "sl": "slovenian", "kn": "kannada", "et": "estonian", "mk": "macedonian",
|
|
"br": "breton", "eu": "basque", "is": "icelandic", "hy": "armenian", "ne": "nepali", "mn": "mongolian", "bs": "bosnian", "kk": "kazakh", "sq": "albanian", "sw": "swahili",
|
|
"gl": "galician", "mr": "marathi", "pa": "punjabi", "si": "sinhala", "km": "khmer", "sn": "shona", "yo": "yoruba", "so": "somali", "af": "afrikaans", "oc": "occitan", "ka": "georgian",
|
|
"be": "belarusian", "tg": "tajik", "sd": "sindhi", "gu": "gujarati", "am": "amharic", "yi": "yiddish", "lo": "lao", "uz": "uzbek", "fo": "faroese", "ht": "haitian creole",
|
|
"ps": "pashto", "tk": "turkmen", "nn": "nynorsk", "mt": "maltese", "sa": "sanskrit", "lb": "luxembourgish", "my": "myanmar", "bo": "tibetan", "tl": "tagalog", "mg": "malagasy",
|
|
"as": "assamese", "tt": "tatar", "haw": "hawaiian", "ln": "lingala", "ha": "hausa", "ba": "bashkir", "jw": "javanese", "su": "sundanese",
|
|
}
|
|
|
|
def get_encoding(encoding_name):
|
|
with fetch(f"https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/{encoding_name}.tiktoken").open() as f:
|
|
ranks = {base64.b64decode(token): int(rank) for token, rank in (line.split() for line in f if line)}
|
|
n_vocab = len(ranks)
|
|
specials = [
|
|
"<|endoftext|>",
|
|
"<|startoftranscript|>",
|
|
*[f"<|{lang}|>" for lang in LANGUAGES.keys()],
|
|
"<|translate|>",
|
|
"<|transcribe|>",
|
|
"<|startoflm|>",
|
|
"<|startofprev|>",
|
|
"<|nospeech|>",
|
|
"<|notimestamps|>",
|
|
*[f"<|{i * 0.02:.2f}|>" for i in range(1501)],
|
|
]
|
|
special_tokens = dict(zip(specials, itertools.count(n_vocab)))
|
|
n_vocab += len(specials)
|
|
import tiktoken
|
|
return tiktoken.Encoding(
|
|
name=encoding_name,
|
|
explicit_n_vocab=n_vocab,
|
|
pat_str=r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""",
|
|
mergeable_ranks=ranks,
|
|
special_tokens=special_tokens)
|
|
|
|
MODEL_URLS = {
|
|
"tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
|
|
"tiny": "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt",
|
|
"base.en": "https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt",
|
|
"base": "https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt",
|
|
"small.en": "https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt",
|
|
"small": "https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt",
|
|
"medium.en": "https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt",
|
|
"medium": "https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt",
|
|
"large-v1": "https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt",
|
|
"large-v2": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
|
|
"large": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
|
|
}
|
|
def init_whisper(model_name="tiny.en", batch_size=1):
|
|
assert MODEL_URLS[model_name] is not None
|
|
|
|
filename = fetch(MODEL_URLS[model_name])
|
|
state = torch_load(filename)
|
|
model = Whisper(state['dims'], batch_size)
|
|
load_state_dict(model, state['model_state_dict'], strict=False)
|
|
enc = get_encoding("multilingual" if model.is_multilingual else "gpt2")
|
|
return model, enc
|
|
|
|
def load_file_waveform(filename):
|
|
waveform, _ = librosa.load(filename, sr=RATE)
|
|
return waveform
|
|
|
|
def transcribe_file(model, enc, filename):
|
|
return transcribe_waveform(model, enc, [load_file_waveform(filename)])
|
|
|
|
def transcribe_waveform(model: Whisper, enc, waveforms, truncate=False):
|
|
"""
|
|
Expects an array of shape (N,S) where N is the number waveforms to transcribe in parallel and S is number of 16000Hz samples
|
|
Returns the transcribed text if a single waveform is provided, or an array of transcriptions if multiple are provided
|
|
"""
|
|
|
|
log_spec = prep_audio(waveforms, model.batch_size, truncate)
|
|
nsample = model.decoder.max_tokens_to_sample
|
|
nctx = model.decoder.max_self_attn_cache_len
|
|
|
|
def inferloop(ctx: Union[np.ndarray, List[np.ndarray]], encoded_audio):
|
|
pos, next_tokens = 0, ctx
|
|
for i in range(nsample):
|
|
next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
|
|
next_tokens[ctx[:, -1] == eot] = eot
|
|
ctx = np.concatenate((ctx, next_tokens), axis=1)
|
|
pos = ctx.shape[-1] - 1
|
|
if (next_tokens == eot).all() or pos == nctx: break
|
|
return ctx
|
|
|
|
def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]
|
|
start_tokens = [enc._special_tokens["<|startoftranscript|>"]]
|
|
if model.is_multilingual:
|
|
# TODO detect language
|
|
language_token = enc._special_tokens["<|startoftranscript|>"] + 1 + tuple(LANGUAGES.keys()).index("en")
|
|
start_tokens.append(language_token)
|
|
start_tokens.append(enc._special_tokens["<|transcribe|>"])
|
|
start_tokens.append(enc._special_tokens["<|notimestamps|>"])
|
|
|
|
eot = enc._special_tokens["<|endoftext|>"]
|
|
|
|
ctx = np.tile(start_tokens, (model.batch_size,1))
|
|
transcriptions = [[] for _ in waveforms]
|
|
|
|
for curr_frame in range(0, log_spec.shape[-1], FRAMES_PER_SEGMENT):
|
|
encoded_audio = model.encoder.encode(Tensor(log_spec[:, :, curr_frame:curr_frame + FRAMES_PER_SEGMENT]))
|
|
|
|
if all(len(c) == len(ctx[0]) for c in ctx): ctx = inferloop(np.array(ctx), encoded_audio)
|
|
else: ctx = [inferloop((np.array([c]*model.batch_size)), encoded_audio)[i] for i,c in enumerate(ctx)]
|
|
|
|
for i, (res, arr) in enumerate(zip(transcriptions, ctx)):
|
|
if curr_frame*HOP_LENGTH <= len(waveforms[i]):res.extend(arr[np.where(arr == start_tokens[-1])[0][0]+1:eoti[0] if len (eoti:=np.where(arr == eot)[0]) else None])
|
|
ctx = [[enc._special_tokens['<|startofprev|>']]+gettexttoks(cs)+start_tokens for cs in ctx]
|
|
|
|
transcriptions = list(map(lambda tokens: enc.decode(tokens).strip(), transcriptions))
|
|
return transcriptions if len(transcriptions) > 1 else transcriptions[0]
|
|
|
|
CHUNK = 1600
|
|
RECORD_SECONDS = 10
|
|
|
|
def listener(q):
|
|
import pyaudio
|
|
p = pyaudio.PyAudio()
|
|
stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
|
|
print("listening")
|
|
for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
|
|
data = stream.read(CHUNK)
|
|
waveform = ((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3)
|
|
q.put(waveform)
|
|
print("done listening")
|
|
|
|
if __name__ == "__main__":
|
|
model, enc = init_whisper("small.en" if getenv("SMALL") else "tiny.en", batch_size=1)
|
|
|
|
if len(sys.argv) > 1:
|
|
print(transcribe_file(model, enc, sys.argv[1]))
|
|
else:
|
|
# online
|
|
q = multiprocessing.Queue()
|
|
p = multiprocessing.Process(target=listener, args=(q,))
|
|
p.daemon = True
|
|
p.start()
|
|
|
|
lst = [enc._special_tokens["<|startoftranscript|>"], enc._special_tokens["<|notimestamps|>"]]
|
|
total = None
|
|
did_read = False
|
|
for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
|
|
while not q.empty() or total is None:
|
|
waveform = q.get()
|
|
if total is None: total = waveform
|
|
else: total = np.concatenate([total, waveform])
|
|
did_read = True
|
|
if did_read:
|
|
log_spec = prep_audio(total.reshape(1, -1), model.batch_size, truncate=True)
|
|
encoded_audio = model.encoder.encode(Tensor(log_spec))
|
|
# pass the previously inferred tokens as 'prefix' - https://github.com/openai/whisper/discussions/117#discussioncomment-3727051
|
|
out = model.decoder(Tensor([lst]), 0, encoded_audio).realize()
|
|
idx = int(out[0,-1].argmax().numpy().item())
|
|
lst.append(idx)
|
|
dec = enc.decode(lst)
|
|
print(dec) # DO NOT REMOVE PRINT. IT'S VERY IMPORTANT
|
|
if dec.endswith("<|endoftext|>"):
|
|
lst.pop()
|