yum-archive/TaSTT-Whisper

High-performance GPGPU inference of OpenAI's Whisper automatic speech recognition (ASR) model

git clone https://git.yummers.dev/yum-archive/TaSTT-Whisper

KonstantinFixed the old sample project012be51

master
121.0 KiB3602 linesraw
1#define WHISPER_BUILD
2#include "whisper.h"
3
4#include "ggml.h"
5
6#include <algorithm>
7#include <cassert>
8#define _USE_MATH_DEFINES
9#include <cmath>
10#include <cstdio>
11#include <cstring>
12#include <fstream>
13#include <map>
14#include <string>
15#include <thread>
16#include <vector>
17#include <regex>
18#include "Utils/Logger.h"
19
20#define USE_FLASH_ATTN
21//#define USE_FLASH_FF
22
23// available whisper models
24enum e_model {
25    MODEL_UNKNOWN,
26    MODEL_TINY,
27    MODEL_BASE,
28    MODEL_SMALL,
29    MODEL_MEDIUM,
30    MODEL_LARGE,
31};
32
33static const std::map<std::string, std::pair<int, std::string>> g_lang = {
34    { "en",  { 0,  "english",         } },
35    { "zh",  { 1,  "chinese",         } },
36    { "de",  { 2,  "german",          } },
37    { "es",  { 3,  "spanish",         } },
38    { "ru",  { 4,  "russian",         } },
39    { "ko",  { 5,  "korean",          } },
40    { "fr",  { 6,  "french",          } },
41    { "ja",  { 7,  "japanese",        } },
42    { "pt",  { 8,  "portuguese",      } },
43    { "tr",  { 9,  "turkish",         } },
44    { "pl",  { 10, "polish",          } },
45    { "ca",  { 11,  "catalan",        } },
46    { "nl",  { 12,  "dutch",          } },
47    { "ar",  { 13,  "arabic",         } },
48    { "sv",  { 14,  "swedish",        } },
49    { "it",  { 15,  "italian",        } },
50    { "id",  { 16,  "indonesian",     } },
51    { "hi",  { 17,  "hindi",          } },
52    { "fi",  { 18,  "finnish",        } },
53    { "vi",  { 19,  "vietnamese",     } },
54    { "iw",  { 20,  "hebrew",         } },
55    { "uk",  { 21,  "ukrainian",      } },
56    { "el",  { 22,  "greek",          } },
57    { "ms",  { 23,  "malay",          } },
58    { "cs",  { 24,  "czech",          } },
59    { "ro",  { 25,  "romanian",       } },
60    { "da",  { 26,  "danish",         } },
61    { "hu",  { 27,  "hungarian",      } },
62    { "ta",  { 28,  "tamil",          } },
63    { "no",  { 29,  "norwegian",      } },
64    { "th",  { 30,  "thai",           } },
65    { "ur",  { 31,  "urdu",           } },
66    { "hr",  { 32,  "croatian",       } },
67    { "bg",  { 33,  "bulgarian",      } },
68    { "lt",  { 34,  "lithuanian",     } },
69    { "la",  { 35,  "latin",          } },
70    { "mi",  { 36,  "maori",          } },
71    { "ml",  { 37,  "malayalam",      } },
72    { "cy",  { 38,  "welsh",          } },
73    { "sk",  { 39,  "slovak",         } },
74    { "te",  { 40,  "telugu",         } },
75    { "fa",  { 41,  "persian",        } },
76    { "lv",  { 42,  "latvian",        } },
77    { "bn",  { 43,  "bengali",        } },
78    { "sr",  { 44,  "serbian",        } },
79    { "az",  { 45,  "azerbaijani",    } },
80    { "sl",  { 46,  "slovenian",      } },
81    { "kn",  { 47,  "kannada",        } },
82    { "et",  { 48,  "estonian",       } },
83    { "mk",  { 49,  "macedonian",     } },
84    { "br",  { 50,  "breton",         } },
85    { "eu",  { 51,  "basque",         } },
86    { "is",  { 52,  "icelandic",      } },
87    { "hy",  { 53,  "armenian",       } },
88    { "ne",  { 54,  "nepali",         } },
89    { "mn",  { 55,  "mongolian",      } },
90    { "bs",  { 56,  "bosnian",        } },
91    { "kk",  { 57,  "kazakh",         } },
92    { "sq",  { 58,  "albanian",       } },
93    { "sw",  { 59,  "swahili",        } },
94    { "gl",  { 60,  "galician",       } },
95    { "mr",  { 61,  "marathi",        } },
96    { "pa",  { 62,  "punjabi",        } },
97    { "si",  { 63,  "sinhala",        } },
98    { "km",  { 64,  "khmer",          } },
99    { "sn",  { 65,  "shona",          } },
100    { "yo",  { 66,  "yoruba",         } },
101    { "so",  { 67,  "somali",         } },
102    { "af",  { 68,  "afrikaans",      } },
103    { "oc",  { 69,  "occitan",        } },
104    { "ka",  { 70,  "georgian",       } },
105    { "be",  { 71,  "belarusian",     } },
106    { "tg",  { 72,  "tajik",          } },
107    { "sd",  { 73,  "sindhi",         } },
108    { "gu",  { 74,  "gujarati",       } },
109    { "am",  { 75,  "amharic",        } },
110    { "yi",  { 76,  "yiddish",        } },
111    { "lo",  { 77,  "lao",            } },
112    { "uz",  { 78,  "uzbek",          } },
113    { "fo",  { 79,  "faroese",        } },
114    { "ht",  { 80,  "haitian creole", } },
115    { "ps",  { 81,  "pashto",         } },
116    { "tk",  { 82,  "turkmen",        } },
117    { "nn",  { 83,  "nynorsk",        } },
118    { "mt",  { 84,  "maltese",        } },
119    { "sa",  { 85,  "sanskrit",       } },
120    { "lb",  { 86,  "luxembourgish",  } },
121    { "my",  { 87,  "myanmar",        } },
122    { "bo",  { 88,  "tibetan",        } },
123    { "tl",  { 89,  "tagalog",        } },
124    { "mg",  { 90,  "malagasy",       } },
125    { "as",  { 91,  "assamese",       } },
126    { "tt",  { 92,  "tatar",          } },
127    { "haw", { 93,  "hawaiian",       } },
128    { "ln",  { 94,  "lingala",        } },
129    { "ha",  { 95,  "hausa",          } },
130    { "ba",  { 96,  "bashkir",        } },
131    { "jw",  { 97,  "javanese",       } },
132    { "su",  { 98,  "sundanese",      } },
133};
134
135static const size_t MB = 1024*1024;
136
137static const std::map<e_model, size_t> MEM_REQ_MODEL = {
138    { MODEL_TINY,     74ull*MB },
139    { MODEL_BASE,    142ull*MB },
140    { MODEL_SMALL,   466ull*MB },
141    { MODEL_MEDIUM, 1464ull*MB },
142    { MODEL_LARGE,  2952ull*MB },
143};
144
145static const std::map<e_model, size_t> MEM_REQ_MEMORY = {
146    { MODEL_TINY,     12ull*MB },
147    { MODEL_BASE,     24ull*MB },
148    { MODEL_SMALL,    70ull*MB },
149    { MODEL_MEDIUM,  184ull*MB },
150    { MODEL_LARGE,   306ull*MB },
151};
152
153static const std::map<e_model, size_t> MEM_REQ_ENCODE = {
154    { MODEL_TINY,     80ull*MB },
155    { MODEL_BASE,    128ull*MB },
156    { MODEL_SMALL,   300ull*MB },
157    { MODEL_MEDIUM,  680ull*MB },
158    { MODEL_LARGE,  1100ull*MB },
159};
160
161static const std::map<e_model, size_t> MEM_REQ_ENCODE_LAYER = {
162    { MODEL_TINY,    104ull*MB },
163    { MODEL_BASE,    138ull*MB },
164    { MODEL_SMALL,   208ull*MB },
165    { MODEL_MEDIUM,  280ull*MB },
166    { MODEL_LARGE,   354ull*MB },
167};
168
169static const std::map<e_model, size_t> MEM_REQ_DECODE = {
170    { MODEL_TINY,    200ull*MB },
171    { MODEL_BASE,    202ull*MB },
172    { MODEL_SMALL,   204ull*MB },
173    { MODEL_MEDIUM,  206ull*MB },
174    { MODEL_LARGE,   208ull*MB },
175};
176
177static const std::map<e_model, size_t> MEM_REQ_DECODE_LAYER = {
178    { MODEL_TINY,     32ull*MB },
179    { MODEL_BASE,     44ull*MB },
180    { MODEL_SMALL,    64ull*MB },
181    { MODEL_MEDIUM,   84ull*MB },
182    { MODEL_LARGE,   110ull*MB },
183};
184
185struct whisper_mel {
186    int n_len;
187    int n_mel;
188
189    std::vector<float> data;
190};
191
192struct whisper_filters {
193    int32_t n_mel;
194    int32_t n_fft;
195
196    std::vector<float> data;
197};
198
199struct whisper_vocab {
200    using id    = int32_t;
201    using token = std::string;
202
203    int n_vocab = 51864;
204
205    std::map<token, id> token_to_id;
206    std::map<id, token> id_to_token;
207
208    id token_eot  = 50256;
209    id token_sot  = 50257;
210    id token_prev = 50360;
211    id token_solm = 50361; // ??
212    id token_not  = 50362; // no timestamps
213    id token_beg  = 50363;
214
215    // available tasks
216    static const id token_translate  = 50358;
217    static const id token_transcribe = 50359;
218
219    bool is_multilingual() const {
220        return n_vocab == 51865;
221    }
222};
223
224struct whisper_segment {
225    int64_t t0;
226    int64_t t1;
227
228    std::string text;
229
230    std::vector<whisper_token_data> tokens;
231};
232
233// medium
234// hparams: {
235// 'n_mels': 80,
236// 'n_vocab': 51864,
237// 'n_audio_ctx': 1500,
238// 'n_audio_state': 1024,
239// 'n_audio_head': 16,
240// 'n_audio_layer': 24,
241// 'n_text_ctx': 448,
242// 'n_text_state': 1024,
243// 'n_text_head': 16,
244// 'n_text_layer': 24
245// }
246//
247// default hparams (Whisper tiny)
248struct whisper_hparams {
249    int32_t n_vocab       = 51864;
250    int32_t n_audio_ctx   = 1500;
251    int32_t n_audio_state = 384;
252    int32_t n_audio_head  = 6;
253    int32_t n_audio_layer = 4;
254    int32_t n_text_ctx    = 448;
255    int32_t n_text_state  = 384;
256    int32_t n_text_head   = 6;
257    int32_t n_text_layer  = 4;
258    int32_t n_mels        = 80;
259    int32_t f16           = 1;
260};
261
262// audio encoding layer
263struct whisper_layer_encoder {
264    // encoder.blocks.*.attn_ln
265    struct ggml_tensor * attn_ln_0_w;
266    struct ggml_tensor * attn_ln_0_b;
267
268    // encoder.blocks.*.attn.out
269    struct ggml_tensor * attn_ln_1_w;
270    struct ggml_tensor * attn_ln_1_b;
271
272    // encoder.blocks.*.attn.query
273    struct ggml_tensor * attn_q_w;
274    struct ggml_tensor * attn_q_b;
275
276    // encoder.blocks.*.attn.key
277    struct ggml_tensor * attn_k_w;
278
279    // encoder.blocks.*.attn.value
280    struct ggml_tensor * attn_v_w;
281    struct ggml_tensor * attn_v_b;
282
283    // encoder.blocks.*.mlp_ln
284    struct ggml_tensor * mlp_ln_w;
285    struct ggml_tensor * mlp_ln_b;
286
287    // encoder.blocks.*.mlp.0
288    struct ggml_tensor * mlp_0_w;
289    struct ggml_tensor * mlp_0_b;
290
291    // encoder.blocks.*.mlp.2
292    struct ggml_tensor * mlp_1_w;
293    struct ggml_tensor * mlp_1_b;
294};
295
296// token decoding layer
297struct whisper_layer_decoder {
298    // decoder.blocks.*.attn_ln
299    struct ggml_tensor * attn_ln_0_w;
300    struct ggml_tensor * attn_ln_0_b;
301
302    // decoder.blocks.*.attn.out
303    struct ggml_tensor * attn_ln_1_w;
304    struct ggml_tensor * attn_ln_1_b;
305
306    // decoder.blocks.*.attn.query
307    struct ggml_tensor * attn_q_w;
308    struct ggml_tensor * attn_q_b;
309
310    // decoder.blocks.*.attn.key
311    struct ggml_tensor * attn_k_w;
312
313    // decoder.blocks.*.attn.value
314    struct ggml_tensor * attn_v_w;
315    struct ggml_tensor * attn_v_b;
316
317    // decoder.blocks.*.cross_attn_ln
318    struct ggml_tensor * cross_attn_ln_0_w;
319    struct ggml_tensor * cross_attn_ln_0_b;
320
321    // decoder.blocks.*.cross_attn.out
322    struct ggml_tensor * cross_attn_ln_1_w;
323    struct ggml_tensor * cross_attn_ln_1_b;
324
325    // decoder.blocks.*.cross_attn.query
326    struct ggml_tensor * cross_attn_q_w;
327    struct ggml_tensor * cross_attn_q_b;
328
329    // decoder.blocks.*.cross_attn.key
330    struct ggml_tensor * cross_attn_k_w;
331
332    // decoder.blocks.*.cross_attn.value
333    struct ggml_tensor * cross_attn_v_w;
334    struct ggml_tensor * cross_attn_v_b;
335
336    // decoder.blocks.*.mlp_ln
337    struct ggml_tensor * mlp_ln_w;
338    struct ggml_tensor * mlp_ln_b;
339
340    // decoder.blocks.*.mlp.0
341    struct ggml_tensor * mlp_0_w;
342    struct ggml_tensor * mlp_0_b;
343
344    // decoder.blocks.*.mlp.2
345    struct ggml_tensor * mlp_1_w;
346    struct ggml_tensor * mlp_1_b;
347};
348
349struct whisper_model {
350    e_model type = MODEL_UNKNOWN;
351
352    whisper_hparams hparams;
353    whisper_filters filters;
354
355    // encoder.positional_embedding
356    struct ggml_tensor * e_pe;
357
358    // encoder.conv1
359    struct ggml_tensor * e_conv_1_w;
360    struct ggml_tensor * e_conv_1_b;
361
362    // encoder.conv2
363    struct ggml_tensor * e_conv_2_w;
364    struct ggml_tensor * e_conv_2_b;
365
366    // encoder.ln_post
367    struct ggml_tensor * e_ln_w;
368    struct ggml_tensor * e_ln_b;
369
370    // decoder.positional_embedding
371    struct ggml_tensor * d_pe; // DD
372
373    // decoder.token_embedding
374    struct ggml_tensor * d_te; // DD
375
376    // decoder.ln
377    struct ggml_tensor * d_ln_w; // DD
378    struct ggml_tensor * d_ln_b; // DD
379
380    std::vector<whisper_layer_encoder> layers_encoder;
381    std::vector<whisper_layer_decoder> layers_decoder;
382
383    // key + value memory
384    struct ggml_tensor * memory_k;
385    struct ggml_tensor * memory_v;
386
387    struct ggml_tensor * memory_cross_k;
388    struct ggml_tensor * memory_cross_v;
389
390    // context
391    struct ggml_context * ctx;
392    struct ggml_context * ctx_mem;
393
394    // tensors
395    int n_loaded;
396    std::map<std::string, struct ggml_tensor *> tensors;
397};
398
399struct whisper_context {
400    int64_t t_load_us   = 0;
401    int64_t t_mel_us    = 0;
402    int64_t t_sample_us = 0;
403    int64_t t_encode_us = 0;
404    int64_t t_decode_us = 0;
405    int64_t t_start_us  = 0;
406
407    std::vector<uint8_t> * buf_model; // the model buffer is read-only and can be shared between processors
408    std::vector<uint8_t>   buf_memory;
409    std::vector<uint8_t>   buf_compute;
410    std::vector<uint8_t>   buf_compute_layer;
411
412    whisper_model model;
413    whisper_vocab vocab;
414
415    whisper_mel mel;
416
417    std::vector<float> probs;
418    std::vector<float> logits;
419
420    std::vector<whisper_segment> result_all;
421
422    std::vector<whisper_token> prompt_past;
423
424    // [EXPERIMENTAL] token-level timestamps data
425    int64_t t_beg;
426    int64_t t_last;
427    whisper_token tid_last;
428    std::vector<float> energy; // PCM signal energy
429
430    // [EXPERIMENTAL] speed-up techniques
431    int32_t exp_n_audio_ctx; // 0 - use default
432};
433
434template<typename T>
435static void read_safe(std::ifstream& fin, T& dest)
436{
437  fin.read((char*)& dest, sizeof(T));
438}
439
440// load the model from a ggml file
441//
442// file format:
443//
444//   - hparams
445//   - pre-computed mel filters
446//   - vocab
447//   - weights
448//
449// see the convert-pt-to-ggml.py script for details
450//
451static bool whisper_model_load(const std::string & fname, whisper_context & wctx) {
452	logDebug( u8"%s: loading model from '%s'", __func__, fname.c_str() );
453
454    auto & model = wctx.model;
455    auto & vocab = wctx.vocab;
456
457    auto fin = std::ifstream(fname, std::ios::binary);
458    if (!fin) {
459		logError( u8"%s: failed to open '%s'", __func__, fname.c_str() );
460        return false;
461    }
462
463    // verify magic
464    {
465        uint32_t magic;
466        read_safe(fin, magic);
467        if (magic != 0x67676d6c) {
468			logError( u8"%s: invalid model file '%s' (bad magic)", __func__, fname.c_str() );
469            return false;
470        }
471    }
472
473    //load hparams
474    {
475        auto & hparams = model.hparams;
476
477        read_safe(fin, hparams.n_vocab);
478        read_safe(fin, hparams.n_audio_ctx);
479        read_safe(fin, hparams.n_audio_state);
480        read_safe(fin, hparams.n_audio_head);
481        read_safe(fin, hparams.n_audio_layer);
482        read_safe(fin, hparams.n_text_ctx);
483        read_safe(fin, hparams.n_text_state);
484        read_safe(fin, hparams.n_text_head);
485        read_safe(fin, hparams.n_text_layer);
486        read_safe(fin, hparams.n_mels);
487        read_safe(fin, hparams.f16);
488
489        assert(hparams.n_text_state == hparams.n_audio_state);
490
491        if (hparams.n_audio_layer == 4) {
492            model.type = e_model::MODEL_TINY;
493        }
494
495        if (hparams.n_audio_layer == 6) {
496            model.type = e_model::MODEL_BASE;
497        }
498
499        if (hparams.n_audio_layer == 12) {
500            model.type = e_model::MODEL_SMALL;
501        }
502
503        if (hparams.n_audio_layer == 24) {
504            model.type = e_model::MODEL_MEDIUM;
505        }
506
507        if (hparams.n_audio_layer == 32) {
508            model.type = e_model::MODEL_LARGE;
509        }
510
511		logDebug( u8"%s: n_vocab       = %d", __func__, hparams.n_vocab);
512		logDebug( u8"%s: n_audio_ctx   = %d", __func__, hparams.n_audio_ctx);
513		logDebug( u8"%s: n_audio_state = %d", __func__, hparams.n_audio_state);
514		logDebug( u8"%s: n_audio_head  = %d", __func__, hparams.n_audio_head);
515		logDebug( u8"%s: n_audio_layer = %d", __func__, hparams.n_audio_layer);
516		logDebug( u8"%s: n_text_ctx    = %d", __func__, hparams.n_text_ctx);
517		logDebug( u8"%s: n_text_state  = %d", __func__, hparams.n_text_state);
518		logDebug( u8"%s: n_text_head   = %d", __func__, hparams.n_text_head);
519		logDebug( u8"%s: n_text_layer  = %d", __func__, hparams.n_text_layer);
520		logDebug( u8"%s: n_mels        = %d", __func__, hparams.n_mels);
521		logDebug( u8"%s: f16           = %d", __func__, hparams.f16);
522		logDebug( u8"%s: type          = %d", __func__, model.type);
523
524        wctx.buf_model = new std::vector<uint8_t>();
525        wctx.buf_model->resize(MEM_REQ_MODEL.at(model.type));
526        wctx.buf_memory.resize(MEM_REQ_MEMORY.at(model.type));
527        wctx.buf_compute.resize(std::max(MEM_REQ_ENCODE.at(model.type), MEM_REQ_DECODE.at(model.type)));
528        wctx.buf_compute_layer.resize(std::max(MEM_REQ_ENCODE_LAYER.at(model.type), MEM_REQ_DECODE_LAYER.at(model.type)));
529    }
530
531    // load mel filters
532    {
533        auto & filters = wctx.model.filters;
534
535        read_safe(fin, filters.n_mel);
536        read_safe(fin, filters.n_fft);
537
538        filters.data.resize(filters.n_mel * filters.n_fft);
539        fin.read((char *) filters.data.data(), filters.data.size() * sizeof(float));
540    }
541
542    // load vocab
543    {
544        int32_t n_vocab = 0;
545        read_safe(fin, n_vocab);
546
547        //if (n_vocab != model.hparams.n_vocab) {
548        //    fprintf(stderr, "%s: invalid model file '%s' (bad vocab size %d != %d)\n",
549        //            __func__, fname.c_str(), n_vocab, model.hparams.n_vocab);
550        //    return false;
551        //}
552
553        std::string word;
554        std::vector<char> tmp;
555        for (int i = 0; i < n_vocab; i++) {
556            uint32_t len;
557            read_safe(fin, len);
558
559            if (len > 0) {
560                tmp.resize(len);
561                fin.read(&tmp[0], tmp.size()); // read to buffer
562                word.assign(&tmp[0], tmp.size());
563            } else {
564                // seems like we have an empty-string token in multi-language models (i = 50256)
565                //fprintf(stderr, "%s: warning: empty-string token in vocab, i = %d\n", __func__, i);
566                word = "";
567            }
568
569            vocab.token_to_id[word] = i;
570            vocab.id_to_token[i] = word;
571
572            //printf("%s: vocab[%d] = '%s'\n", __func__, i, word.c_str());
573        }
574
575        vocab.n_vocab = model.hparams.n_vocab;
576        if (vocab.is_multilingual()) {
577            vocab.token_eot++;
578            vocab.token_sot++;
579            vocab.token_prev++;
580            vocab.token_solm++;
581            vocab.token_not++;
582            vocab.token_beg++;
583        }
584
585        if (n_vocab < model.hparams.n_vocab) {
586			logDebug( u8"%s: adding %d extra tokens", __func__, model.hparams.n_vocab - n_vocab );
587            for (int i = n_vocab; i < model.hparams.n_vocab; i++) {
588                if (i > vocab.token_beg) {
589                    word = "[_TT_" + std::to_string(i - vocab.token_beg) + "]";
590                } else if (i == vocab.token_eot) {
591                    word = "[_EOT_]";
592                } else if (i == vocab.token_sot) {
593                    word = "[_SOT_]";
594                } else if (i == vocab.token_prev) {
595                    word = "[_PREV_]";
596                } else if (i == vocab.token_not) {
597                    word = "[_NOT_]";
598                } else if (i == vocab.token_beg) {
599                    word = "[_BEG_]";
600                } else {
601                    word = "[_extra_token_" + std::to_string(i) + "]";
602                }
603                vocab.token_to_id[word] = i;
604                vocab.id_to_token[i] = word;
605            }
606        }
607    }
608
609    {
610        // this is the total memory required to run the inference
611        const size_t mem_required =
612                   wctx.buf_model->size() +
613                   wctx.buf_memory.size() +
614                   wctx.buf_compute.size() +
615                   wctx.buf_compute_layer.size();
616
617		logDebug( u8"%s: mem_required  = %7.2f MB", __func__, mem_required / 1024.0 / 1024.0 );
618    }
619
620    // for the big tensors, we have the option to store the data in 16-bit floats
621    // in order to save memory and also to speed up the computation
622    const ggml_type wtype = model.hparams.f16 ? GGML_TYPE_F16 : GGML_TYPE_F32;
623
624    size_t ctx_size = 0;
625
626    {
627        const auto & hparams = model.hparams;
628
629        const int n_vocab = hparams.n_vocab;
630
631        const int n_audio_ctx   = hparams.n_audio_ctx;
632        const int n_audio_state = hparams.n_audio_state;
633        const int n_audio_layer = hparams.n_audio_layer;
634
635        const int n_text_ctx   = hparams.n_text_ctx;
636        const int n_text_state = hparams.n_text_state;
637        const int n_text_layer = hparams.n_text_layer;
638
639        const int n_mels = hparams.n_mels;
640
641        // encoder
642        {
643            // TODO: F16 .. maybe not?
644            ctx_size += n_audio_ctx*n_audio_state*ggml_type_size(GGML_TYPE_F32); // e_pe;
645
646            ctx_size += 3*n_mels*n_audio_state*ggml_type_size(wtype);         // e_conv_1_w
647            ctx_size +=          n_audio_state*ggml_type_size(GGML_TYPE_F32); // e_conv_1_b
648
649            ctx_size += 3*n_audio_state*n_audio_state*ggml_type_size(wtype);         // e_conv_2_w
650            ctx_size +=                 n_audio_state*ggml_type_size(GGML_TYPE_F32); // e_conv_2_b
651
652            ctx_size += n_audio_state*ggml_type_size(GGML_TYPE_F32); // e_ln_w;
653            ctx_size += n_audio_state*ggml_type_size(GGML_TYPE_F32); // e_ln_b;
654        }
655
656        // decoder
657        {
658            // TODO: F16 .. maybe not?
659            ctx_size += n_text_ctx*n_text_state*ggml_type_size(GGML_TYPE_F32); // d_pe;
660
661            ctx_size += n_vocab*n_text_state*ggml_type_size(wtype); // d_te;
662
663            ctx_size += n_text_state*ggml_type_size(GGML_TYPE_F32); // d_ln_w;
664            ctx_size += n_text_state*ggml_type_size(GGML_TYPE_F32); // d_ln_b;
665        }
666
667        // encoder layers
668        {
669            ctx_size += n_audio_layer*(n_audio_state*ggml_type_size(GGML_TYPE_F32)); // mlp_ln_w
670            ctx_size += n_audio_layer*(n_audio_state*ggml_type_size(GGML_TYPE_F32)); // mlp_ln_b
671
672            ctx_size += n_audio_layer*(4*n_audio_state*n_audio_state*ggml_type_size(wtype));         // mlp_0_w
673            ctx_size += n_audio_layer*(              4*n_audio_state*ggml_type_size(GGML_TYPE_F32)); // mlp_0_b
674
675            ctx_size += n_audio_layer*(4*n_audio_state*n_audio_state*ggml_type_size(wtype));         // mlp_1_w
676            ctx_size += n_audio_layer*(                n_audio_state*ggml_type_size(GGML_TYPE_F32)); // mlp_1_b
677
678            ctx_size += n_audio_layer*(n_audio_state*ggml_type_size(GGML_TYPE_F32)); // attn_ln_0_w
679            ctx_size += n_audio_layer*(n_audio_state*ggml_type_size(GGML_TYPE_F32)); // attn_ln_0_b
680
681            ctx_size += n_audio_layer*(n_audio_state*n_audio_state*ggml_type_size(wtype));         // attn_q_w
682            ctx_size += n_audio_layer*(              n_audio_state*ggml_type_size(GGML_TYPE_F32)); // attn_q_b
683
684            ctx_size += n_audio_layer*(n_audio_state*n_audio_state*ggml_type_size(wtype)); // attn_k_w
685
686            ctx_size += n_audio_layer*(n_audio_state*n_audio_state*ggml_type_size(wtype));         // attn_v_w
687            ctx_size += n_audio_layer*(              n_audio_state*ggml_type_size(GGML_TYPE_F32)); // attn_v_b
688
689            ctx_size += n_audio_layer*(n_audio_state*n_audio_state*ggml_type_size(wtype));         // attn_ln_1_w
690            ctx_size += n_audio_layer*(              n_audio_state*ggml_type_size(GGML_TYPE_F32)); // attn_ln_1_b
691        }
692
693        // decoder layers
694        {
695            ctx_size += n_text_layer*(n_text_state*ggml_type_size(GGML_TYPE_F32)); // mlp_ln_w
696            ctx_size += n_text_layer*(n_text_state*ggml_type_size(GGML_TYPE_F32)); // mlp_ln_b
697
698            ctx_size += n_text_layer*(4*n_text_state*n_text_state*ggml_type_size(wtype));         // mlp_0_w
699            ctx_size += n_text_layer*(             4*n_text_state*ggml_type_size(GGML_TYPE_F32)); // mlp_0_b
700
701            ctx_size += n_text_layer*(4*n_text_state*n_text_state*ggml_type_size(wtype));         // mlp_1_w
702            ctx_size += n_text_layer*(               n_text_state*ggml_type_size(GGML_TYPE_F32)); // mlp_1_b
703
704            ctx_size += n_text_layer*(n_text_state*ggml_type_size(GGML_TYPE_F32)); // attn_ln_0_w
705            ctx_size += n_text_layer*(n_text_state*ggml_type_size(GGML_TYPE_F32)); // attn_ln_0_b
706
707            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype));         // attn_q_w
708            ctx_size += n_text_layer*(             n_text_state*ggml_type_size(GGML_TYPE_F32)); // attn_q_b
709
710            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype)); // attn_k_w
711
712            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype));         // attn_v_w
713            ctx_size += n_text_layer*(             n_text_state*ggml_type_size(GGML_TYPE_F32)); // attn_v_b
714
715            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype));         // attn_ln_1_w
716            ctx_size += n_text_layer*(             n_text_state*ggml_type_size(GGML_TYPE_F32)); // attn_ln_1_b
717                                                                                                //
718            ctx_size += n_text_layer*(n_text_state*ggml_type_size(GGML_TYPE_F32)); // cross_attn_ln_0_w
719            ctx_size += n_text_layer*(n_text_state*ggml_type_size(GGML_TYPE_F32)); // cross_attn_ln_0_b
720
721            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype));         // cross_attn_q_w
722            ctx_size += n_text_layer*(             n_text_state*ggml_type_size(GGML_TYPE_F32)); // cross_attn_q_b
723
724            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype)); // cross_attn_k_w
725
726            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype));         // cross_attn_v_w
727            ctx_size += n_text_layer*(             n_text_state*ggml_type_size(GGML_TYPE_F32)); // cross_attn_v_b
728
729            ctx_size += n_text_layer*(n_text_state*n_text_state*ggml_type_size(wtype));         // cross_attn_ln_1_w
730            ctx_size += n_text_layer*(             n_text_state*ggml_type_size(GGML_TYPE_F32)); // cross_attn_ln_1_b
731        }
732
733        ctx_size += (15 + 15*n_audio_layer + 24*n_text_layer)*256; // object overhead
734
735		logDebug( u8"%s: ggml ctx size = %7.2f MB", __func__, ctx_size / ( 1024.0 * 1024.0 ) );
736    }
737
738    // create the ggml context
739    {
740        struct ggml_init_params params;
741        params.mem_size   = wctx.buf_model->size();
742        params.mem_buffer = wctx.buf_model->data();
743
744        model.ctx = ggml_init(params);
745        if (!model.ctx) {
746			logError( u8"%s: ggml_init() failed", __func__ );
747            return false;
748        }
749    }
750
751    // prepare memory for the weights
752    {
753        auto & ctx = model.ctx;
754
755        const auto & hparams = model.hparams;
756
757        const int n_vocab = hparams.n_vocab;
758
759        const int n_audio_ctx   = hparams.n_audio_ctx;
760        const int n_audio_state = hparams.n_audio_state;
761        const int n_audio_layer = hparams.n_audio_layer;
762
763        const int n_text_ctx   = hparams.n_text_ctx;
764        const int n_text_state = hparams.n_text_state;
765        const int n_text_layer = hparams.n_text_layer;
766
767        const int n_mels = hparams.n_mels;
768
769        model.layers_encoder.resize(n_audio_layer);
770        model.layers_decoder.resize(n_text_layer);
771
772        // encoder
773        {
774            model.e_pe = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_audio_state, n_audio_ctx);
775
776            model.e_conv_1_w = ggml_new_tensor_3d(ctx, wtype,         3, n_mels, n_audio_state);
777            model.e_conv_1_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n_audio_state);
778
779            model.e_conv_2_w = ggml_new_tensor_3d(ctx, wtype,         3, n_audio_state, n_audio_state);
780            model.e_conv_2_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n_audio_state);
781
782            model.e_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
783            model.e_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
784
785            // map by name
786            model.tensors["encoder.positional_embedding"] = model.e_pe;
787
788            model.tensors["encoder.conv1.weight"] = model.e_conv_1_w;
789            model.tensors["encoder.conv1.bias"]   = model.e_conv_1_b;
790
791            model.tensors["encoder.conv2.weight"] = model.e_conv_2_w;
792            model.tensors["encoder.conv2.bias"]   = model.e_conv_2_b;
793
794            model.tensors["encoder.ln_post.weight"] = model.e_ln_w;
795            model.tensors["encoder.ln_post.bias"]   = model.e_ln_b;
796
797            for (int i = 0; i < n_audio_layer; ++i) {
798                auto & layer = model.layers_encoder[i];
799
800                layer.mlp_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
801                layer.mlp_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
802
803                layer.mlp_0_w = ggml_new_tensor_2d(ctx, wtype,           n_audio_state, 4*n_audio_state);
804                layer.mlp_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 4*n_audio_state);
805
806                layer.mlp_1_w = ggml_new_tensor_2d(ctx, wtype,         4*n_audio_state, n_audio_state);
807                layer.mlp_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32,   n_audio_state);
808
809                layer.attn_ln_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
810                layer.attn_ln_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
811
812                layer.attn_q_w = ggml_new_tensor_2d(ctx, wtype,         n_audio_state, n_audio_state);
813                layer.attn_q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
814
815                layer.attn_k_w = ggml_new_tensor_2d(ctx, wtype,         n_audio_state, n_audio_state);
816
817                layer.attn_v_w = ggml_new_tensor_2d(ctx, wtype,         n_audio_state, n_audio_state);
818                layer.attn_v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
819
820                layer.attn_ln_1_w = ggml_new_tensor_2d(ctx, wtype,         n_audio_state, n_audio_state);
821                layer.attn_ln_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
822
823                // map by name
824                model.tensors["encoder.blocks." + std::to_string(i) + ".mlp_ln.weight"] = layer.mlp_ln_w;
825                model.tensors["encoder.blocks." + std::to_string(i) + ".mlp_ln.bias"]   = layer.mlp_ln_b;
826
827                model.tensors["encoder.blocks." + std::to_string(i) + ".mlp.0.weight"] = layer.mlp_0_w;
828                model.tensors["encoder.blocks." + std::to_string(i) + ".mlp.0.bias"]   = layer.mlp_0_b;
829
830                model.tensors["encoder.blocks." + std::to_string(i) + ".mlp.2.weight"] = layer.mlp_1_w;
831                model.tensors["encoder.blocks." + std::to_string(i) + ".mlp.2.bias"]   = layer.mlp_1_b;
832
833                model.tensors["encoder.blocks." + std::to_string(i) + ".attn_ln.weight"] = layer.attn_ln_0_w;
834                model.tensors["encoder.blocks." + std::to_string(i) + ".attn_ln.bias"]   = layer.attn_ln_0_b;
835
836                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.query.weight"] = layer.attn_q_w;
837                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.query.bias"]   = layer.attn_q_b;
838
839                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.key.weight"] = layer.attn_k_w;
840
841                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.value.weight"] = layer.attn_v_w;
842                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.value.bias"]   = layer.attn_v_b;
843
844                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.out.weight"] = layer.attn_ln_1_w;
845                model.tensors["encoder.blocks." + std::to_string(i) + ".attn.out.bias"]   = layer.attn_ln_1_b;
846            }
847        }
848
849        // decoder
850        {
851            model.d_pe = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_text_state, n_text_ctx);
852
853            model.d_te = ggml_new_tensor_2d(ctx, wtype, n_text_state, n_vocab);
854
855            model.d_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
856            model.d_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
857
858            // map by name
859            model.tensors["decoder.positional_embedding"] = model.d_pe;
860
861            model.tensors["decoder.token_embedding.weight"] = model.d_te;
862
863            model.tensors["decoder.ln.weight"] = model.d_ln_w;
864            model.tensors["decoder.ln.bias"]   = model.d_ln_b;
865
866            for (int i = 0; i < n_text_layer; ++i) {
867                auto & layer = model.layers_decoder[i];
868
869                layer.mlp_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
870                layer.mlp_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
871
872                layer.mlp_0_w = ggml_new_tensor_2d(ctx, wtype,           n_text_state, 4*n_text_state);
873                layer.mlp_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 4*n_text_state);
874
875                layer.mlp_1_w = ggml_new_tensor_2d(ctx, wtype,         4*n_text_state, n_text_state);
876                layer.mlp_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32,   n_text_state);
877
878                layer.attn_ln_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
879                layer.attn_ln_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
880
881                layer.attn_q_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
882                layer.attn_q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
883
884                layer.attn_k_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
885
886                layer.attn_v_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
887                layer.attn_v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
888
889                layer.attn_ln_1_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
890                layer.attn_ln_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
891
892                layer.cross_attn_ln_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
893                layer.cross_attn_ln_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
894
895                layer.cross_attn_q_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
896                layer.cross_attn_q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
897
898                layer.cross_attn_k_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
899
900                layer.cross_attn_v_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
901                layer.cross_attn_v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
902
903                layer.cross_attn_ln_1_w = ggml_new_tensor_2d(ctx, wtype,         n_text_state, n_text_state);
904                layer.cross_attn_ln_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_text_state);
905
906                // map by name
907                model.tensors["decoder.blocks." + std::to_string(i) + ".mlp_ln.weight"] = layer.mlp_ln_w;
908                model.tensors["decoder.blocks." + std::to_string(i) + ".mlp_ln.bias"]   = layer.mlp_ln_b;
909
910                model.tensors["decoder.blocks." + std::to_string(i) + ".mlp.0.weight"] = layer.mlp_0_w;
911                model.tensors["decoder.blocks." + std::to_string(i) + ".mlp.0.bias"]   = layer.mlp_0_b;
912
913                model.tensors["decoder.blocks." + std::to_string(i) + ".mlp.2.weight"] = layer.mlp_1_w;
914                model.tensors["decoder.blocks." + std::to_string(i) + ".mlp.2.bias"]   = layer.mlp_1_b;
915
916                model.tensors["decoder.blocks." + std::to_string(i) + ".attn_ln.weight"] = layer.attn_ln_0_w;
917                model.tensors["decoder.blocks." + std::to_string(i) + ".attn_ln.bias"]   = layer.attn_ln_0_b;
918
919                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.query.weight"] = layer.attn_q_w;
920                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.query.bias"]   = layer.attn_q_b;
921
922                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.key.weight"] = layer.attn_k_w;
923
924                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.value.weight"] = layer.attn_v_w;
925                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.value.bias"]   = layer.attn_v_b;
926
927                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.out.weight"] = layer.attn_ln_1_w;
928                model.tensors["decoder.blocks." + std::to_string(i) + ".attn.out.bias"]   = layer.attn_ln_1_b;
929
930                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn_ln.weight"] = layer.cross_attn_ln_0_w;
931                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn_ln.bias"]   = layer.cross_attn_ln_0_b;
932
933                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.query.weight"] = layer.cross_attn_q_w;
934                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.query.bias"]   = layer.cross_attn_q_b;
935
936                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.key.weight"] = layer.cross_attn_k_w;
937
938                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.value.weight"] = layer.cross_attn_v_w;
939                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.value.bias"]   = layer.cross_attn_v_b;
940
941                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.out.weight"] = layer.cross_attn_ln_1_w;
942                model.tensors["decoder.blocks." + std::to_string(i) + ".cross_attn.out.bias"]   = layer.cross_attn_ln_1_b;
943            }
944        }
945    }
946
947    // create the ggml memory context
948    {
949        struct ggml_init_params params;
950        params.mem_size   = wctx.buf_memory.size();
951        params.mem_buffer = wctx.buf_memory.data();
952
953        model.ctx_mem = ggml_init(params);
954        if (!model.ctx_mem) {
955			logError( u8"%s: ggml_init() failed", __func__ );
956            return false;
957        }
958    }
959
960    // key + value memory
961    {
962        auto & ctx = model.ctx_mem;
963
964        const auto & hparams = model.hparams;
965
966        const int n_text_state = hparams.n_text_state;
967        const int n_text_layer = hparams.n_text_layer;
968        const int n_text_ctx   = hparams.n_text_ctx;
969
970        // key/value memory for the self-attention layer
971        {
972            const int n_mem      = n_text_layer*n_text_ctx;
973            const int n_elements = n_text_state*n_mem;
974
975            model.memory_k = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
976            model.memory_v = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
977        }
978
979        // key/value memory for the cross-attention layer
980        {
981            const int n_audio_ctx = hparams.n_audio_ctx;
982
983            const int n_mem      = n_text_layer*n_audio_ctx;
984            const int n_elements = n_text_state*n_mem;
985
986            model.memory_cross_k = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
987            model.memory_cross_v = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
988        }
989
990        const size_t memory_size =
991            ggml_nbytes(model.memory_k)       + ggml_nbytes(model.memory_v) +
992            ggml_nbytes(model.memory_cross_k) + ggml_nbytes(model.memory_cross_v);
993
994		logDebug( u8"%s: memory size   = %7.2f MB", __func__, memory_size/1024.0/1024.0);
995    }
996
997    // load weights
998    {
999        size_t total_size = 0;
1000
1001        model.n_loaded = 0;
1002
1003        while (true) {
1004            int32_t n_dims;
1005            int32_t length;
1006            int32_t ftype;
1007
1008            read_safe(fin, n_dims);
1009            read_safe(fin, length);
1010            read_safe(fin, ftype);
1011
1012            if (fin.eof()) {
1013                break;
1014            }
1015
1016            int32_t nelements = 1;
1017            int32_t ne[3] = { 1, 1, 1 };
1018            for (int i = 0; i < n_dims; ++i) {
1019                read_safe(fin, ne[i]);
1020                nelements *= ne[i];
1021            }
1022
1023            std::string name;
1024            std::vector<char> tmp(length); // create a buffer
1025            fin.read( &tmp[0], tmp.size() ); // read to buffer
1026            name.assign(&tmp[0], tmp.size());
1027
1028            if (model.tensors.find(name) == model.tensors.end()) {
1029				logError( u8"%s: unknown tensor '%s' in model file", __func__, name.data() );
1030                return false;
1031            }
1032
1033            auto tensor = model.tensors[name.data()];
1034            if (ggml_nelements(tensor) != nelements) {
1035				logError( u8"%s: tensor '%s' has wrong size in model file", __func__, name.data());
1036                return false;
1037            }
1038
1039            if (tensor->ne[0] != ne[0] || tensor->ne[1] != ne[1] || tensor->ne[2] != ne[2]) {
1040				logError( u8"%s: tensor '%s' has wrong shape in model file: got [%d, %d, %d], expected [%d, %d, %d]",
1041					__func__, name.data(), tensor->ne[ 0 ], tensor->ne[ 1 ], tensor->ne[ 2 ], ne[ 0 ], ne[ 1 ], ne[ 2 ] );
1042                return false;
1043            }
1044
1045            const size_t bpe = (ftype == 0) ? sizeof(float) : sizeof(ggml_fp16_t);
1046
1047            if (nelements*bpe != ggml_nbytes(tensor)) {
1048				logError( u8"%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n",
1049					__func__, name.data(), ggml_nbytes( tensor ), nelements* bpe );
1050                return false;
1051            }
1052
1053            fin.read(reinterpret_cast<char *>(tensor->data), ggml_nbytes(tensor));
1054
1055            //printf("%48s - [%5d, %5d, %5d], type = %6s, %6.2f MB\n", name.data(), ne[0], ne[1], ne[2], ftype == 0 ? "float" : "f16", ggml_nbytes(tensor)/1024.0/1024.0);
1056            total_size += ggml_nbytes(tensor);
1057            model.n_loaded++;
1058        }
1059
1060		logDebug( u8"%s: model size    = %7.2f MB", __func__, total_size / 1024.0 / 1024.0 );
1061
1062        if (model.n_loaded == 0) {
1063			logWarning( u8"%s: WARN no tensors loaded from model file - assuming empty model for testing", __func__);
1064        } else if (model.n_loaded != (int) model.tensors.size()) {
1065			logError( u8"%s: ERROR not all tensors loaded from model file - expected %zu, got %d", __func__, model.tensors.size(), model.n_loaded );
1066            return false;
1067        }
1068    }
1069
1070    fin.close();
1071
1072    return true;
1073}
1074
1075// evaluate the encoder
1076//
1077// given audio recording (more specifically, its log mel spectrogram), runs forward pass of the encoder
1078// part of the transformer model and returns the encoded features
1079//
1080//   - model:      the model
1081//   - n_threads:  number of threads to use
1082//   - mel_offset: offset in the mel spectrogram (i.e. audio offset)
1083//
1084static bool whisper_encode(
1085              whisper_context & wctx,
1086        const int n_threads,
1087        const int mel_offset) {
1088    const auto & model   = wctx.model;
1089    const auto & mel_inp = wctx.mel;
1090    const auto & hparams = model.hparams;
1091
1092    const int n_ctx   = wctx.exp_n_audio_ctx > 0 ? wctx.exp_n_audio_ctx : hparams.n_audio_ctx;
1093    const int n_state = hparams.n_audio_state;
1094    const int n_head  = hparams.n_audio_head;
1095    const int n_layer = hparams.n_audio_layer;
1096
1097    const int n_mels = hparams.n_mels;
1098    assert(mel_inp.n_mel == n_mels);
1099
1100    struct ggml_init_params params;
1101    params.mem_size   = wctx.buf_compute.size();
1102    params.mem_buffer = wctx.buf_compute.data();
1103
1104    struct ggml_context * ctx0 = ggml_init(params);
1105
1106    struct ggml_tensor * mel = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 2*n_ctx, n_mels);
1107    assert(mel->type == GGML_TYPE_F32);
1108    {
1109        float * dst = (float *) mel->data;
1110        memset(dst, 0, ggml_nbytes(mel));
1111
1112        const int i0 = std::min(mel_offset, mel_inp.n_len);
1113        const int i1 = std::min(mel_offset + 2*n_ctx, mel_inp.n_len);
1114
1115        for (int j = 0; j < mel_inp.n_mel; ++j) {
1116            for (int i = i0; i < i1; ++i) {
1117                dst[j*2*n_ctx + (i - i0)] = mel_inp.data[j*mel_inp.n_len + i];
1118            }
1119        }
1120    }
1121	Tracing::delayTensor( "enc.input", mel );
1122
1123    struct ggml_tensor * cur;
1124
1125    // convolution + gelu
1126    {
1127        cur = ggml_conv_1d_1s(ctx0, model.e_conv_1_w, mel);
1128		Tracing::delayTensor( "enc.conv1", cur );
1129        cur = ggml_add(ctx0,
1130                ggml_repeat(ctx0,
1131                    model.e_conv_1_b,
1132                    cur),
1133                cur);
1134
1135        cur = ggml_gelu(ctx0, cur);
1136		Tracing::delayTensor( "enc.temp1", cur );
1137
1138        cur = ggml_conv_1d_2s(ctx0, model.e_conv_2_w, cur);
1139        cur = ggml_add(ctx0,
1140                ggml_repeat(ctx0,
1141                    model.e_conv_2_b,
1142                    cur),
1143                cur);
1144
1145        cur = ggml_gelu(ctx0, cur);
1146    }
1147
1148    // ===================================================================
1149    // NOTE: experimenting with partial evaluation of the encoder (ignore)
1150    //static int iter = -1;
1151    //const int n_iter = 1500/n_ctx;
1152
1153    //iter = (iter + 1) % n_iter;
1154
1155    //if (iter == 0) {
1156    //    memset(model.memory_cross_k->data, 0, ggml_nbytes(model.memory_cross_k));
1157    //    memset(model.memory_cross_v->data, 0, ggml_nbytes(model.memory_cross_v));
1158    //}
1159
1160    static int iter = 0;
1161
1162    const size_t e_pe_stride = model.e_pe->ne[0]*ggml_element_size(model.e_pe);
1163    const size_t e_pe_offset = model.e_pe->ne[0]*ggml_element_size(model.e_pe)*n_ctx*iter;
1164
1165    struct ggml_tensor * e_pe = ggml_view_2d(ctx0, model.e_pe, model.e_pe->ne[0], n_ctx, e_pe_stride, e_pe_offset);
1166
1167    cur = ggml_add(ctx0, e_pe, ggml_transpose(ctx0, cur));
1168    // ===================================================================
1169
1170    // original:
1171    //cur = ggml_add(ctx0, model.e_pe, ggml_transpose(ctx0, cur));
1172
1173    struct ggml_tensor * inpL = cur;
1174
1175    for (int il = 0; il < n_layer; ++il) {
1176        const auto & layer = model.layers_encoder[il];
1177
1178        // create separate context for each layer to reduce memory usage
1179
1180        struct ggml_init_params paramsL;
1181        paramsL.mem_size   = wctx.buf_compute_layer.size();
1182        paramsL.mem_buffer = wctx.buf_compute_layer.data();
1183
1184        struct ggml_context * ctxL = ggml_init(paramsL);
1185
1186		Tracing::delayTensor( { "enc.layer[ %i ].in", il }, inpL );
1187
1188        // norm
1189        {
1190            cur = ggml_norm(ctxL, inpL);
1191			if( il == 0 )
1192				Tracing::delayTensor( "enc-norm", cur );
1193
1194            // cur = ln_0_w*cur + ln_0_b
1195            cur = ggml_add(ctxL,
1196                    ggml_mul(ctxL,
1197                        ggml_repeat(ctxL, layer.attn_ln_0_w, cur),
1198                        cur),
1199                    ggml_repeat(ctxL, layer.attn_ln_0_b, cur));
1200        }
1201
1202        // self-attention
1203        {
1204            struct ggml_tensor * Qcur = ggml_mul_mat(ctxL,
1205                    layer.attn_q_w,
1206                    cur);
1207			if( il == 0 )
1208				Tracing::delayTensor( "enc-Qcur", Qcur );
1209
1210            Qcur = ggml_add(ctxL,
1211                    ggml_repeat(ctxL,
1212                        layer.attn_q_b,
1213                        Qcur),
1214                    Qcur);
1215
1216            //Qcur = ggml_scale(ctxL, Qcur, ggml_new_f32(ctxL, pow(float(n_state)/n_head, -0.25)));
1217
1218            // note: no bias for Key
1219            struct ggml_tensor * Kcur = ggml_mul_mat(ctxL,
1220                    layer.attn_k_w,
1221                    cur);
1222			if( il == 0 )
1223				Tracing::delayTensor( "enc-Kcur", Kcur );
1224
1225            //Kcur = ggml_scale(ctxL, Kcur, ggml_new_f32(ctxL, pow(float(n_state)/n_head, -0.25)));
1226
1227            struct ggml_tensor * Vcur = ggml_mul_mat(ctxL,
1228                    layer.attn_v_w,
1229                    cur);
1230			if( il == 0 )
1231				Tracing::delayTensor( "enc-Vcur", Vcur );
1232
1233            Vcur = ggml_add(ctxL,
1234                    ggml_repeat(ctxL,
1235                        layer.attn_v_b,
1236                        Vcur),
1237                    Vcur);
1238
1239            // ------
1240
1241#ifdef USE_FLASH_ATTN
1242            struct ggml_tensor * Q =
1243                ggml_permute(ctxL,
1244                        ggml_cpy(ctxL,
1245                            Qcur,
1246                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
1247                        0, 2, 1, 3);
1248
1249            struct ggml_tensor * K =
1250                ggml_permute(ctxL,
1251                        ggml_cpy(ctxL,
1252                            Kcur,
1253                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
1254                        0, 2, 1, 3);
1255
1256            struct ggml_tensor * V =
1257                ggml_cpy(ctxL,
1258                        ggml_permute(ctxL,
1259                            ggml_reshape_3d(ctxL,
1260                                Vcur,
1261                                n_state/n_head, n_head, n_ctx),
1262                            1, 2, 0, 3),
1263                        ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_ctx, n_state/n_head, n_head)
1264                        );
1265
1266            struct ggml_tensor * KQV = ggml_flash_attn(ctxL, Q, K, V, false);
1267			if( il == 0 )
1268				Tracing::delayTensor( "enc-KQV", KQV );
1269#else
1270            struct ggml_tensor * Q =
1271                ggml_permute(ctxL,
1272                        ggml_cpy(ctxL,
1273                            Qcur,
1274                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F32, n_state/n_head, n_head, n_ctx)),
1275                        0, 2, 1, 3);
1276
1277            struct ggml_tensor * K =
1278                ggml_permute(ctxL,
1279                        ggml_cpy(ctxL,
1280                            Kcur,
1281                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
1282                        0, 2, 1, 3);
1283
1284            // K * Q
1285            struct ggml_tensor * KQ = ggml_mul_mat(ctxL, K, Q);
1286
1287            struct ggml_tensor * KQ_scaled =
1288                ggml_scale(ctxL,
1289                        KQ,
1290                        ggml_new_f32(ctxL, 1.0f/sqrt(float(n_state)/n_head))
1291                        );
1292
1293            struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctxL, KQ_scaled);
1294
1295            //struct ggml_tensor * V_trans =
1296            //    ggml_permute(ctxL,
1297            //            ggml_cpy(ctxL,
1298            //                Vcur,
1299            //                ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
1300            //            1, 2, 0, 3);
1301
1302            //struct ggml_tensor * KQV = ggml_mul_mat(ctxL, V_trans, KQ_soft_max);
1303
1304            struct ggml_tensor * V =
1305                ggml_cpy(ctxL,
1306                        ggml_permute(ctxL,
1307                            ggml_reshape_3d(ctxL,
1308                                Vcur,
1309                                n_state/n_head, n_head, n_ctx),
1310                            0, 2, 1, 3),
1311                        ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_ctx, n_head)
1312                        );
1313
1314            struct ggml_tensor * KQV = ggml_mul_mat(ctxL, ggml_transpose(ctxL, V), KQ_soft_max);
1315#endif
1316
1317            struct ggml_tensor * KQV_merged = ggml_permute(ctxL, KQV, 0, 2, 1, 3);
1318
1319            cur = ggml_cpy(ctxL,
1320                    KQV_merged,
1321                    ggml_new_tensor_2d(ctxL, GGML_TYPE_F32, n_state, n_ctx));
1322        }
1323
1324        // projection
1325        {
1326            cur = ggml_mul_mat(ctxL,
1327                    layer.attn_ln_1_w,
1328                    cur);
1329
1330            cur = ggml_add(ctxL,
1331                    ggml_repeat(ctxL, layer.attn_ln_1_b, cur),
1332                    cur);
1333        }
1334
1335        // add the input
1336        cur = ggml_add(ctxL, cur, inpL);
1337
1338        struct ggml_tensor * inpFF = cur;
1339
1340        // feed-forward network
1341        {
1342            // norm
1343            {
1344                cur = ggml_norm(ctxL, inpFF);
1345
1346                // cur = mlp_ln_w*cur + mlp_ln_b
1347                cur = ggml_add(ctxL,
1348                        ggml_mul(ctxL,
1349                            ggml_repeat(ctxL, layer.mlp_ln_w, cur),
1350                            cur),
1351                        ggml_repeat(ctxL, layer.mlp_ln_b, cur));
1352            }
1353
1354#ifdef USE_FLASH_FF
1355            cur = ggml_flash_ff(ctxL,
1356                    ggml_cpy(ctxL, cur, ggml_new_tensor_2d(ctxL, GGML_TYPE_F16, n_state, N)),
1357                    layer.mlp_0_w, layer.mlp_0_b, layer.mlp_1_w, layer.mlp_1_b);
1358#else
1359            // fully connected
1360            cur = ggml_mul_mat(ctxL,
1361                    layer.mlp_0_w,
1362                    cur);
1363
1364            cur = ggml_add(ctxL,
1365                    ggml_repeat(ctxL, layer.mlp_0_b, cur),
1366                    cur);
1367
1368            // GELU activation
1369            cur = ggml_gelu(ctxL, cur);
1370
1371            // projection
1372            cur = ggml_mul_mat(ctxL,
1373                    layer.mlp_1_w,
1374                    cur);
1375
1376            cur = ggml_add(ctxL,
1377                    ggml_repeat(ctxL, layer.mlp_1_b, cur),
1378                    cur);
1379#endif
1380        }
1381
1382        // output from this layer
1383        struct ggml_tensor * inpO = ggml_add(ctxL, cur, inpFF);
1384
1385        {
1386            struct ggml_cgraph gf = {};
1387            gf.n_threads = n_threads;
1388
1389            ggml_build_forward_expand(&gf, inpO);
1390            ggml_graph_compute       (ctxL, &gf);
1391			Tracing::writeDelayedTensors();
1392            //ggml_graph_print(&gf);
1393        }
1394
1395        // TODO: this is a hack to have per-layer computation graphs - need to come up with something better
1396        // input for next layer (inpO -> inpL)
1397        memcpy(inpL->data, inpO->data, ggml_nbytes(inpL));
1398        inpL->op = GGML_OP_NONE;
1399        inpL->src0 = nullptr;
1400        inpL->src1 = nullptr;
1401
1402        //printf("%s: - used_mem(%d) = %f MB\n", __func__, il, ggml_used_mem(ctxL)/1024.0/1024.0);
1403
1404        ggml_free(ctxL);
1405    }
1406	Tracing::tensor( "enc.layers", inpL );
1407	cur = inpL;
1408
1409    // norm
1410    {
1411        cur = ggml_norm(ctx0, cur);
1412
1413        // cur = ln_f_g*cur + ln_f_b
1414        cur = ggml_add(ctx0,
1415                ggml_mul(ctx0,
1416                    ggml_repeat(ctx0, model.e_ln_w, cur),
1417                    cur),
1418                ggml_repeat(ctx0, model.e_ln_b, cur));
1419    }
1420
1421    // run the computation
1422    {
1423        struct ggml_cgraph gf = {};
1424        gf.n_threads = n_threads;
1425
1426        ggml_build_forward_expand(&gf, cur);
1427        ggml_graph_compute       (ctx0, &gf);
1428
1429        //ggml_graph_print(&gf);
1430    }
1431
1432	Tracing::tensor( "encode-out", cur );
1433
1434    // cur
1435    //{
1436    //    printf("ne0 = %d\n", cur->ne[0]);
1437    //    printf("ne1 = %d\n", cur->ne[1]);
1438    //    for (int i = 0; i < 10; ++i) {
1439    //        printf("%8.4f ", ((float *)(cur->data))[i]);
1440    //    }
1441    //    printf("... ");
1442    //    for (int i = cur->ne[0] - 10; i < cur->ne[0]; ++i) {
1443    //        printf("%8.4f ", ((float *)(cur->data))[i]);
1444    //    }
1445    //    printf("\n");
1446    //}
1447
1448    // pre-compute cross-attention memory
1449    {
1450        struct ggml_cgraph gf = {};
1451        gf.n_threads = n_threads;
1452
1453        // TODO: hack to disconnect the encoded features from the previous graph
1454        cur->op = GGML_OP_NONE;
1455        cur->src0 = nullptr;
1456        cur->src1 = nullptr;
1457
1458        for (int il = 0; il < model.hparams.n_text_layer; ++il) {
1459            auto & layer = model.layers_decoder[il];
1460
1461            struct ggml_tensor * Kcross = ggml_mul_mat(ctx0,
1462                    layer.cross_attn_k_w,
1463                    cur);
1464
1465            Kcross = ggml_scale(ctx0, Kcross, ggml_new_f32(ctx0, pow(float(n_state)/n_head, -0.25)));
1466
1467            struct ggml_tensor * Vcross = ggml_mul_mat(ctx0,
1468                    layer.cross_attn_v_w,
1469                    cur);
1470
1471            Vcross = ggml_add(ctx0,
1472                    ggml_repeat(ctx0,
1473                        layer.cross_attn_v_b,
1474                        Vcross),
1475                    Vcross);
1476
1477            //struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_cross_k, n_state*n_ctx, (ggml_element_size(model.memory_cross_k)*n_state)*(il*hparams.n_audio_ctx + iter*n_ctx));
1478            //struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_cross_v, n_state*n_ctx, (ggml_element_size(model.memory_cross_v)*n_state)*(il*hparams.n_audio_ctx + iter*n_ctx));
1479            struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_cross_k, n_state*n_ctx, (ggml_element_size(model.memory_cross_k)*n_state)*(il*n_ctx));
1480            struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_cross_v, n_state*n_ctx, (ggml_element_size(model.memory_cross_v)*n_state)*(il*n_ctx));
1481
1482            ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Kcross, k));
1483            ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcross, v));
1484        }
1485
1486        ggml_graph_compute(ctx0, &gf);
1487    }
1488
1489    ////////////////////////////////////////////////////////////////////////////
1490
1491    //printf("%s: used_mem = %f MB\n", __func__, ggml_used_mem(ctx0)/1024.0/1024.0);
1492
1493    ggml_free(ctx0);
1494
1495    return true;
1496}
1497
1498// evaluate the decoder
1499//
1500// given text prompt + audio features -> predicts the probabilities for the next token
1501//
1502//   - model:      the model
1503//   - n_threads:  number of threads to use
1504//   - tokens:     text prompt
1505//   - n_tokens:   number of tokens in the prompt
1506//   - n_past:     number of past tokens to prefix the prompt with
1507//
1508static bool whisper_decode(
1509              whisper_context & wctx,
1510        const int n_threads,
1511        const whisper_token * tokens,
1512        const int n_tokens,
1513        const int n_past) {
1514    const auto & model   = wctx.model;
1515    const auto & hparams = model.hparams;
1516
1517    auto & logits_out = wctx.logits;
1518    auto & probs_out  = wctx.probs;
1519
1520    const int n_vocab = hparams.n_vocab;
1521
1522    const int n_ctx   = hparams.n_text_ctx;
1523    const int n_state = hparams.n_text_state;
1524    const int n_head  = hparams.n_text_head;
1525    const int n_layer = hparams.n_text_layer;
1526
1527    const int N = n_tokens;
1528    const int M = wctx.exp_n_audio_ctx > 0 ? wctx.exp_n_audio_ctx : hparams.n_audio_ctx;
1529
1530    struct ggml_init_params params;
1531    params.mem_size   = wctx.buf_compute.size();
1532    params.mem_buffer = wctx.buf_compute.data();
1533
1534    struct ggml_context * ctx0 = ggml_init(params);
1535
1536    struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N);
1537    memcpy(embd->data, tokens, N*ggml_element_size(embd));
1538
1539    struct ggml_tensor * position = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N);
1540    for (int i = 0; i < N; ++i) {
1541        ((int32_t *) position->data)[i] = n_past + i;
1542    }
1543
1544    // token encoding + position encoding
1545    struct ggml_tensor * cur =
1546        ggml_add(ctx0,
1547                ggml_get_rows(ctx0, model.d_te, embd),
1548                ggml_get_rows(ctx0, model.d_pe, position));
1549	Tracing::delayTensor( "dec-rows", cur );
1550
1551    struct ggml_tensor * inpL = cur;
1552
1553    for (int il = 0; il < n_layer; ++il) {
1554        const auto & layer = model.layers_decoder[il];
1555
1556        struct ggml_init_params paramsL;
1557        paramsL.mem_size   = wctx.buf_compute_layer.size();
1558        paramsL.mem_buffer = wctx.buf_compute_layer.data();
1559
1560        struct ggml_context * ctxL = ggml_init(paramsL);
1561        struct ggml_cgraph gf = {};
1562        gf.n_threads = n_threads;
1563
1564        // norm
1565        {
1566            cur = ggml_norm(ctxL, inpL);
1567
1568            // cur = ln_0_w*cur + ln_0_b
1569            cur = ggml_add(ctxL,
1570                    ggml_mul(ctxL,
1571                        ggml_repeat(ctxL, layer.attn_ln_0_w, cur),
1572                        cur),
1573                    ggml_repeat(ctxL, layer.attn_ln_0_b, cur));
1574        }
1575
1576        // self-attention
1577        {
1578            struct ggml_tensor * Qcur = ggml_mul_mat(ctxL,
1579                    layer.attn_q_w,
1580                    cur);
1581
1582            Qcur = ggml_add(ctxL,
1583                    ggml_repeat(ctxL,
1584                        layer.attn_q_b,
1585                        Qcur),
1586                    Qcur);
1587
1588            Qcur = ggml_scale(ctxL, Qcur, ggml_new_f32(ctxL, pow(float(n_state)/n_head, -0.25)));
1589
1590            // note: no bias for Key
1591            struct ggml_tensor * Kcur = ggml_mul_mat(ctxL,
1592                    layer.attn_k_w,
1593                    cur);
1594
1595            Kcur = ggml_scale(ctxL, Kcur, ggml_new_f32(ctxL, pow(float(n_state)/n_head, -0.25)));
1596
1597            struct ggml_tensor * Vcur = ggml_mul_mat(ctxL,
1598                    layer.attn_v_w,
1599                    cur);
1600
1601            Vcur = ggml_add(ctxL,
1602                    ggml_repeat(ctxL,
1603                        layer.attn_v_b,
1604                        Vcur),
1605                    Vcur);
1606
1607            // store key and value to memory
1608            {
1609                struct ggml_tensor * k = ggml_view_1d(ctxL, model.memory_k, N*n_state, (ggml_element_size(model.memory_k)*n_state)*(il*n_ctx + n_past));
1610                struct ggml_tensor * v = ggml_view_1d(ctxL, model.memory_v, N*n_state, (ggml_element_size(model.memory_v)*n_state)*(il*n_ctx + n_past));
1611
1612                ggml_build_forward_expand(&gf, ggml_cpy(ctxL, Kcur, k));
1613                ggml_build_forward_expand(&gf, ggml_cpy(ctxL, Vcur, v));
1614            }
1615
1616            // ------
1617
1618            struct ggml_tensor * Q =
1619                ggml_permute(ctxL,
1620                        ggml_cpy(ctxL,
1621                            Qcur,
1622                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F32, n_state/n_head, n_head, N)),
1623                        0, 2, 1, 3);
1624
1625            struct ggml_tensor * K =
1626                ggml_permute(ctxL,
1627                        ggml_reshape_3d(ctxL,
1628                            ggml_view_1d(ctxL, model.memory_k, (n_past + N)*n_state, il*n_ctx*ggml_element_size(model.memory_k)*n_state),
1629                            n_state/n_head, n_head, n_past + N),
1630                        0, 2, 1, 3);
1631
1632            // K * Q
1633            struct ggml_tensor * KQ = ggml_mul_mat(ctxL, K, Q);
1634
1635            //struct ggml_tensor * KQ_scaled =
1636            //    ggml_scale(ctxL,
1637            //            KQ,
1638            //            ggml_new_f32(ctxL, 1.0f/sqrt(float(n_state)/n_head))
1639            //            );
1640
1641            struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctxL, KQ, n_past);
1642
1643            struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctxL, KQ_masked);
1644			if( 0 == il ) Tracing::delayTensor( "dec-KQ", KQ_soft_max );
1645
1646            struct ggml_tensor * V_trans =
1647                ggml_permute(ctxL,
1648                        ggml_reshape_3d(ctxL,
1649                            ggml_view_1d(ctxL, model.memory_v, (n_past + N)*n_state, il*n_ctx*ggml_element_size(model.memory_v)*n_state),
1650                            n_state/n_head, n_head, n_past + N),
1651                        1, 2, 0, 3);
1652
1653            struct ggml_tensor * KQV = ggml_mul_mat(ctxL, V_trans, KQ_soft_max);
1654			if( 0 == il ) Tracing::delayTensor( "dec-KQV", KQV );
1655
1656            struct ggml_tensor * KQV_merged = ggml_permute(ctxL, KQV, 0, 2, 1, 3);
1657
1658            cur = ggml_cpy(ctxL,
1659                    KQV_merged,
1660                    ggml_new_tensor_2d(ctxL, GGML_TYPE_F32, n_state, N));
1661        }
1662
1663        {
1664            cur = ggml_mul_mat(ctxL,
1665                    layer.attn_ln_1_w,
1666                    cur);
1667
1668            cur = ggml_add(ctxL,
1669                    ggml_repeat(ctxL, layer.attn_ln_1_b, cur),
1670                    cur);
1671        }
1672
1673        // add the input
1674        struct ggml_tensor * inpCA = ggml_add(ctxL, cur, inpL);
1675
1676        // norm
1677        {
1678            cur = ggml_norm(ctxL, inpCA); // note: we use inpCA here
1679
1680            // cur = ln_0_w*cur + ln_0_b
1681            cur = ggml_add(ctxL,
1682                    ggml_mul(ctxL,
1683                        ggml_repeat(ctxL, layer.cross_attn_ln_0_w, cur),
1684                        cur),
1685                    ggml_repeat(ctxL, layer.cross_attn_ln_0_b, cur));
1686        }
1687
1688        // cross-attention
1689        {
1690            struct ggml_tensor * Qcur = ggml_mul_mat(ctxL,
1691                    layer.cross_attn_q_w,
1692                    cur);
1693
1694            Qcur = ggml_add(ctxL,
1695                    ggml_repeat(ctxL,
1696                        layer.cross_attn_q_b,
1697                        Qcur),
1698                    Qcur);
1699
1700            Qcur = ggml_scale(ctxL, Qcur, ggml_new_f32(ctxL, pow(float(n_state)/n_head, -0.25)));
1701
1702            // Kcross is already scaled
1703            struct ggml_tensor * Kcross =
1704                ggml_reshape_3d(ctxL,
1705                        ggml_view_1d(ctxL, model.memory_cross_k, M*n_state, il*M*ggml_element_size(model.memory_cross_k)*n_state),
1706                        n_state/n_head, n_head, M);
1707
1708            struct ggml_tensor * Vcross =
1709                ggml_reshape_3d(ctxL,
1710                        ggml_view_1d(ctxL, model.memory_cross_v, M*n_state, il*M*ggml_element_size(model.memory_cross_v)*n_state),
1711                        n_state/n_head, n_head, M);
1712
1713            // ------
1714
1715            struct ggml_tensor * Q =
1716                ggml_permute(ctxL,
1717                        ggml_cpy(ctxL,
1718                            Qcur,
1719                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F32, n_state/n_head, n_head, N)),
1720                        0, 2, 1, 3);
1721
1722            struct ggml_tensor * K = ggml_permute(ctxL, Kcross, 0, 2, 1, 3);
1723
1724            // K * Q
1725            struct ggml_tensor * KQ = ggml_mul_mat(ctxL, K, Q);
1726
1727            //struct ggml_tensor * KQ_scaled =
1728            //    ggml_scale(ctxL,
1729            //            KQ,
1730            //            ggml_new_f32(ctxL, 1.0f/sqrt(float(n_state)/n_head))
1731            //            );
1732
1733            // no masking for cross-attention
1734            //struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctxL, KQ_scaled, n_past);
1735
1736            struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctxL, KQ);
1737
1738            struct ggml_tensor * V_trans = ggml_permute(ctxL, Vcross, 1, 2, 0, 3);
1739
1740            struct ggml_tensor * KQV = ggml_mul_mat(ctxL, V_trans, KQ_soft_max);
1741			if( 0 == il ) Tracing::delayTensor( "dec-KQV", KQV );
1742
1743            struct ggml_tensor * KQV_merged = ggml_permute(ctxL, KQV, 0, 2, 1, 3);
1744
1745            // cur = KQV_merged.contiguous().view(n_state, N)
1746            cur = ggml_cpy(ctxL,
1747                    KQV_merged,
1748                    ggml_new_tensor_2d(ctxL, GGML_TYPE_F32, n_state, N));
1749        }
1750
1751        // projection
1752        {
1753            cur = ggml_mul_mat(ctxL,
1754                    layer.cross_attn_ln_1_w,
1755                    cur);
1756
1757            cur = ggml_add(ctxL,
1758                    ggml_repeat(ctxL, layer.cross_attn_ln_1_b, cur),
1759                    cur);
1760        }
1761
1762        // add the input
1763        cur = ggml_add(ctxL, cur, inpCA);
1764
1765        struct ggml_tensor * inpFF = cur;
1766
1767        // feed-forward network
1768        {
1769            // norm
1770            {
1771                cur = ggml_norm(ctxL, inpFF);
1772
1773                // cur = mlp_ln_w*cur + mlp_ln_b
1774                cur = ggml_add(ctxL,
1775                        ggml_mul(ctxL,
1776                            ggml_repeat(ctxL, layer.mlp_ln_w, cur),
1777                            cur),
1778                        ggml_repeat(ctxL, layer.mlp_ln_b, cur));
1779            }
1780
1781            // fully connected
1782            cur = ggml_mul_mat(ctxL,
1783                    layer.mlp_0_w,
1784                    cur);
1785
1786            cur = ggml_add(ctxL,
1787                    ggml_repeat(ctxL, layer.mlp_0_b, cur),
1788                    cur);
1789
1790            // GELU activation
1791            cur = ggml_gelu(ctxL, cur);
1792
1793            // projection
1794            cur = ggml_mul_mat(ctxL,
1795                    layer.mlp_1_w,
1796                    cur);
1797
1798            cur = ggml_add(ctxL,
1799                    ggml_repeat(ctxL, layer.mlp_1_b, cur),
1800                    cur);
1801        }
1802
1803        // output from this layer
1804        struct ggml_tensor * inpO = ggml_add(ctxL, cur, inpFF);
1805
1806        {
1807            ggml_build_forward_expand(&gf, inpO);
1808            ggml_graph_compute       (ctxL, &gf);
1809			Tracing::writeDelayedTensors();
1810            //ggml_graph_print(&gf);
1811        }
1812
1813        // TODO: this is a hack to have per-layer computation graphs - need to come up with something better
1814        // input for next layer (inpO -> inpL)
1815        memcpy(inpL->data, inpO->data, ggml_nbytes(inpL));
1816        inpL->op = GGML_OP_NONE;
1817        inpL->src0 = nullptr;
1818        inpL->src1 = nullptr;
1819
1820        if (N > 1) {
1821            //printf("%s: - used_mem(%d) = %f MB\n", __func__, il, ggml_used_mem(ctxL)/1024.0/1024.0);
1822        }
1823
1824        ggml_free(ctxL);
1825    }
1826
1827    cur = inpL;
1828
1829    // norm
1830    {
1831        cur = ggml_norm(ctx0, cur);
1832
1833        cur = ggml_add(ctx0,
1834                ggml_mul(ctx0,
1835                    ggml_repeat(ctx0, model.d_ln_w, cur),
1836                    cur),
1837                ggml_repeat(ctx0, model.d_ln_b, cur));
1838    }
1839
1840    struct ggml_tensor * logits = ggml_mul_mat(ctx0, model.d_te, cur);
1841
1842    // logits -> probs
1843    cur = ggml_dup(ctx0, logits);
1844    cur = ggml_soft_max(ctx0, cur); // in-place
1845
1846    // run the computation
1847    {
1848        struct ggml_cgraph gf = {};
1849        gf.n_threads = n_threads;
1850
1851        ggml_build_forward_expand(&gf, cur);
1852        ggml_graph_compute       (ctx0, &gf);
1853    }
1854
1855    logits_out.resize(N*n_vocab);
1856    memcpy(logits_out.data(), ggml_get_data(logits), sizeof(float)*N*n_vocab);
1857
1858    probs_out.resize(N*n_vocab);
1859    memcpy(probs_out.data(), ggml_get_data(cur), sizeof(float)*N*n_vocab);
1860
1861    if (N > 1) {
1862        //const float mem_per_token = ggml_used_mem(ctx0)/1024.0/1024.0/N;
1863        //printf("%s: used_mem = %f MB / %f per token\n", __func__, ggml_used_mem(ctx0)/1024.0/1024.0, mem_per_token);
1864        //printf("%s: max mem = %f MB\n", __func__, mem_per_token*model.hparams.n_text_ctx);
1865    }
1866
1867    ggml_free(ctx0);
1868	// Hash::vector( "probs", probs_out );
1869	Tracing::vector( "probs", probs_out );
1870
1871    return true;
1872}
1873
1874// the most basic sampling scheme - select the top token
1875static whisper_token_data whisper_sample_best(
1876        const whisper_vocab & vocab,
1877        const float * probs,
1878              bool force_timestamp,
1879              bool is_initial) {
1880    whisper_token_data result = {
1881        0, 0, 0.0f, 0.0f, 0.0f, -1, -1, 0.0f,
1882    };
1883
1884    int n_logits = vocab.id_to_token.size();
1885
1886    std::vector<std::pair<double, whisper_vocab::id>> probs_id;
1887    probs_id.reserve(n_logits);
1888
1889    for (int i = 0; i < n_logits; i++) {
1890        probs_id.emplace_back(probs[i], i);
1891    }
1892
1893    {
1894        double sum_ts =  0.0;
1895        double max_ts = -1.0;
1896        double max_tx = -1.0;
1897
1898        for (int i = 0; i < vocab.token_beg; i++) {
1899            max_tx = std::max(max_tx, probs_id[i].first);
1900        }
1901
1902        const auto i0 = is_initial ? vocab.token_beg + 101 : vocab.token_beg;
1903        const auto i1 = is_initial ? vocab.token_beg + 101 : n_logits;
1904
1905        // the initial timestamp cannot be larger than 100
1906        // ref: https://github.com/openai/whisper/blob/0b1ba3d46ebf7fe6f953acfd8cad62a4f851b49f/whisper/decoding.py#L426-L429
1907        if (is_initial) {
1908            for (int i = i0; i < n_logits; ++ i) {
1909                probs_id[i].first = -INFINITY;
1910            }
1911        }
1912
1913        for (int i = vocab.token_beg; i < i1; i++) {
1914            sum_ts += probs_id[i].first;
1915            if  (probs_id[i].first > max_ts) {
1916                max_ts = probs_id[i].first;
1917                result.tid = probs_id[i].second;
1918            }
1919        }
1920
1921        // if the probability sum of all timestamp tokens is higher than the max probability of the text tokens - sample a
1922        // timestamp token
1923        if (sum_ts > max_tx || force_timestamp) {
1924            // ref: https://github.com/openai/whisper/blob/0b1ba3d46ebf7fe6f953acfd8cad62a4f851b49f/whisper/decoding.py#L430-L438
1925            for (int i = 0; i < vocab.token_beg; i++) {
1926                probs_id[i].first = -INFINITY;
1927            }
1928        }
1929
1930        result.pt = max_ts/(sum_ts + 1e-10);
1931        result.ptsum = sum_ts;
1932    }
1933
1934    // find the top K tokens
1935    const int top_k = 4;
1936
1937    std::partial_sort(
1938            probs_id.begin(),
1939            probs_id.begin() + top_k, probs_id.end(),
1940            [](const std::pair<double, whisper_vocab::id> & a, const std::pair<double, whisper_vocab::id> & b) {
1941        return a.first > b.first;
1942    });
1943
1944    probs_id.resize(top_k);
1945
1946    //printf("\n");
1947    //for (int i = 0; i < (int) probs_id.size(); i++) {
1948    //    printf("%d: '%s' %f, %d\n", i, vocab.id_to_token.at(probs_id[i].second).c_str(), probs_id[i].first, probs_id[i].second);
1949    //}
1950
1951    int res = 0;
1952    while ((probs_id[res].second == vocab.token_sot ||
1953            probs_id[res].second == vocab.token_solm ||
1954            probs_id[res].second == vocab.token_not) &&
1955            res < (int) probs_id.size() - 1) {
1956        res++;
1957    }
1958
1959    result.id = probs_id[res].second;
1960    result.p  = probs_id[res].first;
1961
1962    return result;
1963}
1964
1965//  500 -> 00:05.000
1966// 6000 -> 01:00.000
1967static std::string to_timestamp(int64_t t, bool comma = false) {
1968    int64_t msec = t * 10;
1969    int64_t hr = msec / (1000 * 60 * 60);
1970    msec = msec - hr * (1000 * 60 * 60);
1971    int64_t min = msec / (1000 * 60);
1972    msec = msec - min * (1000 * 60);
1973    int64_t sec = msec / 1000;
1974    msec = msec - sec * 1000;
1975
1976    char buf[32];
1977    snprintf(buf, sizeof(buf), "%02d:%02d:%02d%s%03d", (int) hr, (int) min, (int) sec, comma ? "," : ".", (int) msec);
1978
1979    return std::string(buf);
1980}
1981
1982// naive Discrete Fourier Transform
1983// input is real-valued
1984// output is complex-valued
1985static void dft(const std::vector<float> & in, std::vector<float> & out) {
1986    int N = in.size();
1987
1988    out.resize(N*2);
1989
1990    for (int k = 0; k < N; k++) {
1991        float re = 0;
1992        float im = 0;
1993
1994        for (int n = 0; n < N; n++) {
1995            float angle = 2*M_PI*k*n/N;
1996            re += in[n]*cos(angle);
1997            im -= in[n]*sin(angle);
1998        }
1999
2000        out[k*2 + 0] = re;
2001        out[k*2 + 1] = im;
2002    }
2003}
2004
2005// Cooley-Tukey FFT
2006// poor man's implementation - use something better
2007// input is real-valued
2008// output is complex-valued
2009static void fft(const std::vector<float> & in, std::vector<float> & out) {
2010    out.resize(in.size()*2);
2011
2012    int N = in.size();
2013
2014    if (N == 1) {
2015        out[0] = in[0];
2016        out[1] = 0;
2017        return;
2018    }
2019
2020    if (N%2 == 1) {
2021        dft(in, out);
2022        return;
2023    }
2024
2025    std::vector<float> even;
2026    std::vector<float> odd;
2027
2028    for (int i = 0; i < N; i++) {
2029        if (i % 2 == 0) {
2030            even.push_back(in[i]);
2031        } else {
2032            odd.push_back(in[i]);
2033        }
2034    }
2035
2036    std::vector<float> even_fft;
2037    std::vector<float> odd_fft;
2038
2039    fft(even, even_fft);
2040    fft(odd, odd_fft);
2041
2042    for (int k = 0; k < N/2; k++) {
2043        float theta = 2*M_PI*k/N;
2044
2045        float re = cos(theta);
2046        float im = -sin(theta);
2047
2048        float re_odd = odd_fft[2*k + 0];
2049        float im_odd = odd_fft[2*k + 1];
2050
2051        out[2*k + 0] = even_fft[2*k + 0] + re*re_odd - im*im_odd;
2052        out[2*k + 1] = even_fft[2*k + 1] + re*im_odd + im*re_odd;
2053
2054        out[2*(k + N/2) + 0] = even_fft[2*k + 0] - re*re_odd + im*im_odd;
2055        out[2*(k + N/2) + 1] = even_fft[2*k + 1] - re*im_odd - im*re_odd;
2056    }
2057}
2058
2059// ref: https://github.com/openai/whisper/blob/main/whisper/audio.py#L92-L124
2060static bool log_mel_spectrogram(
2061    const float * samples,
2062    const int n_samples,
2063    const int /*sample_rate*/,
2064    const int fft_size,
2065    const int fft_step,
2066    const int n_mel,
2067    const int n_threads,
2068    const whisper_filters & filters,
2069    const bool speed_up,
2070    whisper_mel & mel) {
2071
2072    // Hanning window
2073    std::vector<float> hann;
2074    hann.resize(fft_size);
2075    for (int i = 0; i < fft_size; i++) {
2076        hann[i] = 0.5*(1.0 - cos((2.0*M_PI*i)/(fft_size)));
2077    }
2078
2079    mel.n_mel = n_mel;
2080    mel.n_len = (n_samples)/fft_step;
2081    mel.data.resize(mel.n_mel*mel.n_len);
2082
2083    const int n_fft = 1 + (speed_up ? fft_size/4 : fft_size/2);
2084
2085    //printf("%s: n_samples = %d, n_len = %d\n", __func__, n_samples, mel.n_len);
2086    //printf("%s: recording length: %f s\n", __func__, (float) n_samples/sample_rate);
2087
2088    std::vector<std::thread> workers(n_threads);
2089    for (int iw = 0; iw < n_threads; ++iw) {
2090        workers[iw] = std::thread([&](int ith) {
2091            std::vector<float> fft_in;
2092            fft_in.resize(fft_size);
2093            for (int i = 0; i < fft_size; i++) {
2094                fft_in[i] = 0.0;
2095            }
2096
2097            std::vector<float> fft_out;
2098            fft_out.resize(2*fft_size);
2099
2100            for (int i = ith; i < mel.n_len; i += n_threads) {
2101                const int offset = i*fft_step;
2102
2103                // apply Hanning window
2104                for (int j = 0; j < fft_size; j++) {
2105                    if (offset + j < n_samples) {
2106                        fft_in[j] = hann[j]*samples[offset + j];
2107                    } else {
2108                        fft_in[j] = 0.0;
2109                    }
2110                }
2111
2112                // FFT -> mag^2
2113                fft(fft_in, fft_out);
2114
2115                for (int j = 0; j < fft_size; j++) {
2116                    fft_out[j] = (fft_out[2*j + 0]*fft_out[2*j + 0] + fft_out[2*j + 1]*fft_out[2*j + 1]);
2117                }
2118                for (int j = 1; j < fft_size/2; j++) {
2119                    //if (i == 0) {
2120                    //    printf("%d: %f %f\n", j, fft_out[j], fft_out[fft_size - j]);
2121                    //}
2122                    fft_out[j] += fft_out[fft_size - j];
2123                }
2124                if (i == 0) {
2125                    //for (int j = 0; j < fft_size; j++) {
2126                    //    printf("%d: %e\n", j, fft_out[j]);
2127                    //}
2128                }
2129
2130                if (speed_up) {
2131                    // scale down in the frequency domain results in a speed up in the time domain
2132                    for (int j = 0; j < n_fft; j++) {
2133                        fft_out[j] = 0.5*(fft_out[2*j] + fft_out[2*j + 1]);
2134                    }
2135                }
2136
2137                // mel spectrogram
2138                for (int j = 0; j < mel.n_mel; j++) {
2139                    double sum = 0.0;
2140
2141                    for (int k = 0; k < n_fft; k++) {
2142                        sum += fft_out[k]*filters.data[j*n_fft + k];
2143                    }
2144                    if (sum < 1e-10) {
2145                        sum = 1e-10;
2146                    }
2147
2148                    sum = log10(sum);
2149
2150                    mel.data[j*mel.n_len + i] = sum;
2151                }
2152            }
2153        }, iw);
2154    }
2155
2156    for (int iw = 0; iw < n_threads; ++iw) {
2157        workers[iw].join();
2158    }
2159
2160    // clamping and normalization
2161    double mmax = -1e20;
2162    for (int i = 0; i < mel.n_mel*mel.n_len; i++) {
2163        if (mel.data[i] > mmax) {
2164            mmax = mel.data[i];
2165        }
2166    }
2167    //printf("%s: max = %f\n", __func__, mmax);
2168
2169    mmax -= 8.0;
2170
2171    for (int i = 0; i < mel.n_mel*mel.n_len; i++) {
2172        if (mel.data[i] < mmax) {
2173            mel.data[i] = mmax;
2174        }
2175
2176        mel.data[i] = (mel.data[i] + 4.0)/4.0;
2177    }
2178
2179    return true;
2180}
2181
2182// split text into tokens
2183//
2184// ref: https://github.com/openai/gpt-2/blob/a74da5d99abaaba920de8131d64da2862a8f213b/src/encoder.py#L53
2185//
2186// Regex (Python):
2187// r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
2188//
2189// Regex (C++):
2190// R"('s|'t|'re|'ve|'m|'ll|'d| ?[[:alpha:]]+| ?[[:digit:]]+| ?[^\s[:alpha:][:digit:]]+|\s+(?!\S)|\s+)"
2191//
2192static std::vector<whisper_vocab::id> tokenize(const whisper_vocab & vocab, const std::string & text) {
2193    std::vector<std::string> words;
2194
2195    // first split the text into words
2196    {
2197        std::string str = text;
2198        std::string pat = R"('s|'t|'re|'ve|'m|'ll|'d| ?[[:alpha:]]+| ?[[:digit:]]+| ?[^\s[:alpha:][:digit:]]+|\s+(?!\S)|\s+)";
2199
2200        std::regex re(pat);
2201        std::smatch m;
2202
2203        while (std::regex_search(str, m, re)) {
2204            for (auto x : m) {
2205                words.push_back(x);
2206            }
2207            str = m.suffix();
2208        }
2209    }
2210
2211    // find the longest tokens that form the words:
2212    std::vector<whisper_vocab::id> tokens;
2213    for (const auto & word : words) {
2214        if (word.empty()) continue;
2215
2216        int i = 0;
2217        int n = word.size();
2218        while (i < n) {
2219            int j = n;
2220            while (j > i) {
2221                auto it = vocab.token_to_id.find(word.substr(i, j-i));
2222                if (it != vocab.token_to_id.end()) {
2223                    tokens.push_back(it->second);
2224                    i = j;
2225                    break;
2226                }
2227                --j;
2228            }
2229            if (i == n) {
2230                break;
2231            }
2232            if (j == i) {
2233                auto sub = word.substr(i, 1);
2234                if (vocab.token_to_id.find(sub) != vocab.token_to_id.end()) {
2235                    tokens.push_back(vocab.token_to_id.at(sub));
2236                } else {
2237					logWarning( u8"%s: unknown token '%s'", __func__, sub.data() );
2238                }
2239                ++i;
2240            }
2241        }
2242    }
2243
2244    return tokens;
2245}
2246
2247//
2248// interface implementation
2249//
2250
2251struct whisper_context * whisper_init(const char * path_model) {
2252    ggml_time_init();
2253
2254    whisper_context * ctx = new whisper_context;
2255
2256    const int64_t t_start_us = ggml_time_us();
2257
2258    ctx->t_start_us = t_start_us;
2259
2260    if (!whisper_model_load(path_model, *ctx)) {
2261		logError( u8"%s: failed to load model from '%s'", __func__, path_model );
2262        delete ctx;
2263        return nullptr;
2264    }
2265
2266    ctx->t_load_us = ggml_time_us() - t_start_us;
2267
2268    return ctx;
2269}
2270
2271void whisper_free(struct whisper_context * ctx) {
2272    if (ctx) {
2273        if (ctx->model.ctx) {
2274            ggml_free(ctx->model.ctx);
2275        }
2276        if (ctx->model.ctx_mem) {
2277            ggml_free(ctx->model.ctx_mem);
2278        }
2279        if (ctx->buf_model) {
2280            delete ctx->buf_model;
2281        }
2282        delete ctx;
2283    }
2284}
2285
2286int whisper_pcm_to_mel(struct whisper_context * ctx, const float * samples, int n_samples, int n_threads) {
2287    const int64_t t_start_us = ggml_time_us();
2288
2289    if (!log_mel_spectrogram(samples, n_samples, WHISPER_SAMPLE_RATE, WHISPER_N_FFT, WHISPER_HOP_LENGTH, WHISPER_N_MEL, n_threads, ctx->model.filters, false, ctx->mel)) {
2290		logError( u8"%s: failed to compute mel spectrogram", __func__ );
2291        return -1;
2292    }
2293
2294    ctx->t_mel_us = ggml_time_us() - t_start_us;
2295
2296    return 0;
2297}
2298
2299// same as whisper_pcm_to_mel, but applies a Phase Vocoder to speed up the audio x2
2300int whisper_pcm_to_mel_phase_vocoder(struct whisper_context * ctx, const float * samples, int n_samples, int n_threads) {
2301    const int64_t t_start_us = ggml_time_us();
2302
2303    if (!log_mel_spectrogram(samples, n_samples, WHISPER_SAMPLE_RATE, 2*WHISPER_N_FFT, 2*WHISPER_HOP_LENGTH, WHISPER_N_MEL, n_threads, ctx->model.filters, true, ctx->mel)) {
2304		logError( u8"%s: failed to compute mel spectrogram", __func__ );
2305        return -1;
2306    }
2307
2308    ctx->t_mel_us = ggml_time_us() - t_start_us;
2309
2310    return 0;
2311}
2312
2313int whisper_set_mel(
2314        struct whisper_context * ctx,
2315        const float * data,
2316        int n_len,
2317        int n_mel) {
2318    if (n_mel != WHISPER_N_MEL) {
2319		logError( u8"%s: invalid number of mel bands: %d (expected %d)", __func__, n_mel, WHISPER_N_MEL );
2320        return -1;
2321    }
2322
2323    ctx->mel.n_len = n_len;
2324    ctx->mel.n_mel = n_mel;
2325
2326    ctx->mel.data.resize(n_len*n_mel);
2327    memcpy(ctx->mel.data.data(), data, n_len*n_mel*sizeof(float));
2328
2329    return 0;
2330}
2331
2332int whisper_encode(struct whisper_context * ctx, int offset, int n_threads) {
2333    const int64_t t_start_us = ggml_time_us();
2334
2335    if (!whisper_encode(*ctx, n_threads, offset)) {
2336		logError( u8"%s: failed to eval", __func__ );
2337        return -1;
2338    }
2339
2340    ctx->t_encode_us += ggml_time_us() - t_start_us;
2341
2342    return 0;
2343}
2344
2345int whisper_decode(struct whisper_context * ctx, const whisper_token * tokens, int n_tokens, int n_past, int n_threads) {
2346    const int64_t t_start_us = ggml_time_us();
2347
2348    if (!whisper_decode(*ctx, n_threads, tokens, n_tokens, n_past)) {
2349		logError( u8"%s: failed to eval", __func__ );
2350        return 1;
2351    }
2352
2353    ctx->t_decode_us += ggml_time_us() - t_start_us;
2354
2355    return 0;
2356}
2357
2358struct whisper_token_data whisper_sample_best(struct whisper_context * ctx) {
2359    const int64_t t_start_sample_us = ggml_time_us();
2360
2361    const auto res = whisper_sample_best(ctx->vocab, ctx->probs.data() + (ctx->probs.size() - ctx->vocab.n_vocab), false, false);
2362
2363    ctx->t_sample_us += ggml_time_us() - t_start_sample_us;
2364
2365    return res;
2366}
2367
2368struct whisper_token_data whisper_sample_timestamp(struct whisper_context * ctx, bool is_initial) {
2369    const int64_t t_start_sample_us = ggml_time_us();
2370
2371    const auto res = whisper_sample_best(ctx->vocab, ctx->probs.data() + (ctx->probs.size() - ctx->vocab.n_vocab), true, is_initial);
2372
2373    ctx->t_sample_us += ggml_time_us() - t_start_sample_us;
2374
2375    return res;
2376}
2377
2378int whisper_tokenize(struct whisper_context * ctx, const char * text, whisper_token * tokens, int n_max_tokens) {
2379    const auto res = tokenize(ctx->vocab, text);
2380
2381    if (n_max_tokens < (int) res.size()) {
2382		logError( u8"%s: too many resulting tokens: %d (max %d)", __func__, (int)res.size(), n_max_tokens );
2383        return -1;
2384    }
2385
2386    for (int i = 0; i < (int) res.size(); i++) {
2387        tokens[i] = res[i];
2388    }
2389
2390    return res.size();
2391}
2392
2393int whisper_lang_max_id() {
2394    auto max_id = 0;
2395    for (const auto & kv : g_lang) {
2396        max_id = std::max(max_id, kv.second.first);
2397    }
2398
2399    return max_id;
2400}
2401
2402int whisper_lang_id(const char * lang) {
2403    if (!g_lang.count(lang)) {
2404        for (const auto & kv : g_lang) {
2405            if (kv.second.second == lang) {
2406                return kv.second.first;
2407            }
2408        }
2409
2410		logError( u8"%s: unknown language '%s'", __func__, lang );
2411        return -1;
2412    }
2413
2414    return g_lang.at(lang).first;
2415}
2416
2417const char * whisper_lang_str(int id) {
2418    for (const auto & kv : g_lang) {
2419        if (kv.second.first == id) {
2420            return kv.first.c_str();
2421        }
2422    }
2423
2424	logError( u8"%s: unknown language id %d", __func__, id );
2425    return nullptr;
2426}
2427
2428int whisper_lang_auto_detect(
2429        struct whisper_context * ctx,
2430        int offset_ms,
2431        int n_threads,
2432        float * lang_probs) {
2433    const int seek = offset_ms/10;
2434
2435    if (seek < 0) {
2436		logError( u8"%s: offset %dms is before the start of the audio", __func__, offset_ms );
2437        return -1;
2438    }
2439
2440    if (seek >= ctx->mel.n_len) {
2441		logError( u8"%s: offset %dms is past the end of the audio (%dms)", __func__, offset_ms, ctx->mel.n_len * 10 );
2442        return -2;
2443    }
2444
2445    // run the encoder
2446    if (whisper_encode(ctx, seek, n_threads) != 0) {
2447		logError( u8"%s: failed to encode", __func__ );
2448        return -6;
2449    }
2450
2451    const std::vector<whisper_token> prompt = { whisper_token_sot(ctx) };
2452
2453    if (whisper_decode(ctx, prompt.data(), prompt.size(), 0, n_threads) != 0) {
2454		logError( u8"%s: failed to decode", __func__ );
2455        return -7;
2456    }
2457
2458    std::vector<std::pair<float, int>> probs_id;
2459    for (const auto & kv : g_lang) {
2460        const auto token_lang = whisper_token_lang(ctx, kv.second.first);
2461        probs_id.emplace_back( ctx->probs[token_lang], kv.second.first );
2462    }
2463
2464    // sort descending
2465    {
2466        using pair_type = decltype(probs_id)::value_type;
2467        std::sort(probs_id.begin(), probs_id.end(), [](const pair_type & a, const pair_type & b) {
2468            return a.first > b.first;
2469        });
2470    }
2471
2472    // softmax
2473    {
2474        float sum = 0;
2475        for (const auto & kv : probs_id) {
2476            sum += exp(kv.first);
2477        }
2478
2479        for (auto & kv : probs_id) {
2480            kv.first = exp(kv.first) / sum;
2481        }
2482    }
2483
2484    {
2485        for (int i = 0; i < (int) probs_id.size(); i++) {
2486            if (lang_probs) {
2487                lang_probs[probs_id[i].second] = probs_id[i].first;
2488            }
2489
2490            //printf("%s: lang %2d (%3s): %f\n", __func__, probs_id[i].second, whisper_lang_str(probs_id[i].second), probs_id[i].first);
2491        }
2492    }
2493
2494    return probs_id[0].second;
2495}
2496
2497int whisper_n_len(struct whisper_context * ctx) {
2498    return ctx->mel.n_len;
2499}
2500
2501int whisper_n_vocab(struct whisper_context * ctx) {
2502    return ctx->vocab.n_vocab;
2503}
2504
2505int whisper_n_text_ctx(struct whisper_context * ctx) {
2506    return ctx->model.hparams.n_text_ctx;
2507}
2508
2509int whisper_is_multilingual(struct whisper_context * ctx) {
2510    return ctx->vocab.is_multilingual() ? 1 : 0;
2511}
2512
2513float * whisper_get_probs(struct whisper_context * ctx) {
2514    return ctx->probs.data();
2515}
2516
2517const char * whisper_token_to_str(struct whisper_context * ctx, whisper_token token) {
2518    return ctx->vocab.id_to_token.at(token).c_str();
2519}
2520
2521whisper_token whisper_token_eot(struct whisper_context * ctx) {
2522    return ctx->vocab.token_eot;
2523}
2524
2525whisper_token whisper_token_sot(struct whisper_context * ctx) {
2526    return ctx->vocab.token_sot;
2527}
2528
2529whisper_token whisper_token_prev(struct whisper_context * ctx) {
2530    return ctx->vocab.token_prev;
2531}
2532
2533whisper_token whisper_token_solm(struct whisper_context * ctx) {
2534    return ctx->vocab.token_solm;
2535}
2536
2537whisper_token whisper_token_not(struct whisper_context * ctx) {
2538    return ctx->vocab.token_not;
2539}
2540
2541whisper_token whisper_token_beg(struct whisper_context * ctx) {
2542    return ctx->vocab.token_beg;
2543}
2544
2545whisper_token whisper_token_lang(struct whisper_context * ctx, int lang_id) {
2546    return whisper_token_sot(ctx) + 1 + lang_id;
2547}
2548
2549whisper_token whisper_token_translate(void) {
2550    return whisper_vocab::token_translate;
2551}
2552
2553whisper_token whisper_token_transcribe(void) {
2554    return whisper_vocab::token_transcribe;
2555}
2556
2557void whisper_print_timings(struct whisper_context * ctx) {
2558    const int64_t t_end_us = ggml_time_us();
2559
2560	logInfo( u8"%s:     load time = %8.2f ms", __func__, ctx->t_load_us / 1000.0f );
2561	logInfo( u8"%s:      mel time = %8.2f ms", __func__, ctx->t_mel_us / 1000.0f );
2562	logInfo( u8"%s:   sample time = %8.2f ms", __func__, ctx->t_sample_us / 1000.0f );
2563	logInfo( u8"%s:   encode time = %8.2f ms / %.2f ms per layer", __func__,
2564		ctx->t_encode_us / 1000.0f, ctx->t_encode_us / 1000.0f / ctx->model.hparams.n_audio_layer );
2565	logInfo( u8"%s:   decode time = %8.2f ms / %.2f ms per layer", __func__,
2566		ctx->t_decode_us / 1000.0f, ctx->t_decode_us / 1000.0f / ctx->model.hparams.n_text_layer );
2567	logInfo( u8"%s:    total time = %8.2f ms", __func__, ( t_end_us - ctx->t_start_us ) / 1000.0f );
2568}
2569
2570void whisper_reset_timings(struct whisper_context * ctx) {
2571    ctx->t_sample_us = 0;
2572    ctx->t_encode_us = 0;
2573    ctx->t_decode_us = 0;
2574}
2575
2576const char * whisper_print_system_info(void) {
2577    static std::string s;
2578
2579    s  = "";
2580    s += "AVX = "       + std::to_string(ggml_cpu_has_avx())       + " | ";
2581    s += "AVX2 = "      + std::to_string(ggml_cpu_has_avx2())      + " | ";
2582    s += "AVX512 = "    + std::to_string(ggml_cpu_has_avx512())    + " | ";
2583    s += "FMA = "       + std::to_string(ggml_cpu_has_fma())       + " | ";
2584    s += "NEON = "      + std::to_string(ggml_cpu_has_neon())      + " | ";
2585    s += "ARM_FMA = "   + std::to_string(ggml_cpu_has_arm_fma())   + " | ";
2586    s += "F16C = "      + std::to_string(ggml_cpu_has_f16c())      + " | ";
2587    s += "FP16_VA = "   + std::to_string(ggml_cpu_has_fp16_va())   + " | ";
2588    s += "WASM_SIMD = " + std::to_string(ggml_cpu_has_wasm_simd()) + " | ";
2589    s += "BLAS = "      + std::to_string(ggml_cpu_has_blas())      + " | ";
2590
2591    return s.c_str();
2592}
2593
2594////////////////////////////////////////////////////////////////////////////
2595
2596struct whisper_full_params whisper_full_default_params(enum whisper_sampling_strategy strategy) {
2597    struct whisper_full_params result;
2598
2599    switch (strategy) {
2600        case WHISPER_SAMPLING_GREEDY:
2601            {
2602                result = {
2603                    /*.strategy         =*/ WHISPER_SAMPLING_GREEDY,
2604
2605                    /*.n_threads        =*/ std::min(4, (int32_t) std::thread::hardware_concurrency()),
2606                    /*.n_max_text_ctx   =*/ 16384,
2607                    /*.offset_ms        =*/ 0,
2608                    /*.duration_ms      =*/ 0,
2609
2610                    /*.translate        =*/ false,
2611                    /*.no_context       =*/ false,
2612                    /*.single_segment   =*/ false,
2613                    /*.print_special    =*/ false,
2614                    /*.print_progress   =*/ true,
2615                    /*.print_realtime   =*/ false,
2616                    /*.print_timestamps =*/ true,
2617
2618                    /*.token_timestamps =*/ false,
2619                    /*.thold_pt         =*/ 0.01f,
2620                    /*.thold_ptsum      =*/ 0.01f,
2621                    /*.max_len          =*/ 0,
2622                    /*.max_tokens       =*/ 0,
2623
2624                    /*.speed_up         =*/ false,
2625                    /*.audio_ctx        =*/ 0,
2626
2627                    /*.prompt_tokens    =*/ nullptr,
2628                    /*.prompt_n_tokens  =*/ 0,
2629
2630                    /*.language         =*/ "en",
2631
2632                    /*.greedy           =*/ {
2633                        /*.n_past =*/ 0,
2634                    },
2635
2636                    /*.beam_search      =*/ {
2637                        /*.n_past     =*/ -1,
2638                        /*.beam_width =*/ -1,
2639                        /*.n_best     =*/ -1,
2640                    },
2641
2642                    /*.new_segment_callback           =*/ nullptr,
2643                    /*.new_segment_callback_user_data =*/ nullptr,
2644
2645                    /*.encoder_begin_callback           =*/ nullptr,
2646                    /*.encoder_begin_callback_user_data =*/ nullptr,
2647                };
2648            } break;
2649        case WHISPER_SAMPLING_BEAM_SEARCH:
2650            {
2651                result = {
2652                    /*.strategy         =*/ WHISPER_SAMPLING_BEAM_SEARCH,
2653
2654                    /*.n_threads        =*/ std::min(4, (int32_t) std::thread::hardware_concurrency()),
2655                    /*.n_max_text_ctx   =*/ 16384,
2656                    /*.offset_ms        =*/ 0,
2657                    /*.duration_ms      =*/ 0,
2658
2659                    /*.translate        =*/ false,
2660                    /*.no_context       =*/ false,
2661                    /*.single_segment   =*/ false,
2662                    /*.print_special    =*/ false,
2663                    /*.print_progress   =*/ true,
2664                    /*.print_realtime   =*/ false,
2665                    /*.print_timestamps =*/ true,
2666
2667                    /*.token_timestamps =*/ false,
2668                    /*.thold_pt         =*/ 0.01f,
2669                    /*.thold_ptsum      =*/ 0.01f,
2670                    /*.max_len          =*/ 0,
2671                    /*.max_tokens       =*/ 0,
2672
2673                    /*.speed_up         =*/ false,
2674                    /*.audio_ctx        =*/ 0,
2675
2676                    /*.prompt_tokens    =*/ nullptr,
2677                    /*.prompt_n_tokens  =*/ 0,
2678
2679                    /*.language         =*/ "en",
2680
2681                    /*.greedy           =*/ {
2682                        /*.n_past =*/ -1,
2683                    },
2684
2685                    /*.beam_search      =*/ {
2686                        /*.n_past     =*/ 0,
2687                        /*.beam_width =*/ 10,
2688                        /*.n_best     =*/ 5,
2689                    },
2690
2691                    /*.new_segment_callback           =*/ nullptr,
2692                    /*.new_segment_callback_user_data =*/ nullptr,
2693
2694                    /*.encoder_begin_callback           =*/ nullptr,
2695                    /*.encoder_begin_callback_user_data =*/ nullptr,
2696                };
2697            } break;
2698    }
2699
2700    return result;
2701}
2702
2703// forward declarations
2704static std::vector<float> get_signal_energy(const float * signal, int n_samples, int n_samples_per_half_window);
2705static void whisper_exp_compute_token_level_timestamps(
2706        struct whisper_context * ctx,
2707        int   i_segment,
2708        float thold_pt,
2709        float thold_ptsum);
2710
2711// wrap the last segment to max_len characters
2712// returns the number of new segments
2713static int whisper_wrap_segment(struct whisper_context * ctx, int max_len) {
2714    auto segment = ctx->result_all.back();
2715
2716    int res = 1;
2717    int acc = 0;
2718
2719    std::string text;
2720
2721    for (int i = 0; i < (int) segment.tokens.size(); i++) {
2722        const auto & token = segment.tokens[i];
2723        if (token.id >= whisper_token_eot(ctx)) {
2724            continue;
2725        }
2726
2727        const auto txt = whisper_token_to_str(ctx, token.id);
2728
2729        const int cur = strlen(txt);
2730
2731        if (acc + cur > max_len && i > 0) {
2732            // split here
2733            ctx->result_all.back().text = std::move(text);
2734            ctx->result_all.back().t1 = token.t0;
2735            ctx->result_all.back().tokens.resize(i);
2736
2737            ctx->result_all.push_back({});
2738            ctx->result_all.back().t0 = token.t0;
2739            ctx->result_all.back().t1 = segment.t1;
2740
2741            // add tokens [i, end] to the new segment
2742            ctx->result_all.back().tokens.insert(
2743                    ctx->result_all.back().tokens.end(),
2744                    segment.tokens.begin() + i,
2745                    segment.tokens.end());
2746
2747            acc = 0;
2748            text = "";
2749
2750            segment = ctx->result_all.back();
2751            i = -1;
2752
2753            res++;
2754        } else {
2755            acc += cur;
2756            text += txt;
2757        }
2758    }
2759
2760    ctx->result_all.back().text = std::move(text);
2761
2762    return res;
2763}
2764
2765int whisper_full(
2766        struct whisper_context * ctx,
2767        struct whisper_full_params params,
2768        const float * samples,
2769        int n_samples) {
2770    // clear old results
2771    auto & result_all = ctx->result_all;
2772
2773    result_all.clear();
2774
2775    // compute log mel spectrogram
2776    if (params.speed_up) {
2777        if (whisper_pcm_to_mel_phase_vocoder(ctx, samples, n_samples, params.n_threads) != 0) {
2778			logError( u8"%s: failed to compute log mel spectrogram", __func__ );
2779            return -1;
2780        }
2781    } else {
2782        if (whisper_pcm_to_mel(ctx, samples, n_samples, params.n_threads) != 0) {
2783			logError( u8"%s: failed to compute log mel spectrogram", __func__ );
2784            return -2;
2785        }
2786    }
2787
2788    // auto-detect language if not specified
2789    if (params.language == nullptr || strlen(params.language) == 0 || strcmp(params.language, "auto") == 0) {
2790        std::vector<float> probs(whisper_lang_max_id() + 1, 0.0f);
2791
2792        const auto lang_id = whisper_lang_auto_detect(ctx, 0, params.n_threads, probs.data());
2793        if (lang_id < 0) {
2794			logError( u8"%s: failed to auto-detect language", __func__ );
2795            return -3;
2796        }
2797
2798        params.language = whisper_lang_str(lang_id);
2799
2800		logInfo( u8"%s: auto-detected language: %s (p = %f)", __func__, params.language, probs[ whisper_lang_id( params.language ) ] );
2801    }
2802
2803    if (params.token_timestamps) {
2804        ctx->t_beg = 0;
2805        ctx->t_last = 0;
2806        ctx->tid_last = 0;
2807        ctx->energy = get_signal_energy(samples, n_samples, 32);
2808    }
2809
2810    const int seek_start = params.offset_ms/10;
2811    const int seek_end = seek_start + (params.duration_ms == 0 ? whisper_n_len(ctx) : params.duration_ms/10);
2812
2813    // if length of spectrogram is less than 1s (100 samples), then return
2814    // basically don't process anything that is less than 1s
2815    // see issue #39: https://github.com/ggerganov/whisper.cpp/issues/39
2816    if (seek_end < 100 + seek_start) {
2817        return 0;
2818    }
2819
2820    // the accumulated text context so far
2821    auto & prompt_past = ctx->prompt_past;
2822    if (params.no_context) {
2823        prompt_past.clear();
2824    }
2825
2826    // prepend the prompt tokens to the prompt_past
2827    if (params.prompt_tokens && params.prompt_n_tokens > 0) {
2828        // parse tokens from the pointer
2829        for (int i = 0; i < params.prompt_n_tokens; i++) {
2830            prompt_past.push_back(params.prompt_tokens[i]);
2831        }
2832        std::rotate(prompt_past.begin(), prompt_past.end() - params.prompt_n_tokens, prompt_past.end());
2833    }
2834
2835    // overwrite audio_ctx
2836    ctx->exp_n_audio_ctx = params.audio_ctx;
2837
2838    // these tokens determine the task that will be performed
2839    std::vector<whisper_token> prompt_init = { whisper_token_sot(ctx) };
2840    if (whisper_is_multilingual(ctx)) {
2841        const int lang_id = whisper_lang_id(params.language);
2842        prompt_init.push_back(whisper_token_lang(ctx, lang_id));
2843        if (params.translate) {
2844            prompt_init.push_back(whisper_token_translate());
2845        } else {
2846            prompt_init.push_back(whisper_token_transcribe());
2847        }
2848    }
2849
2850    int progress_prev = 0;
2851    int progress_step = 5;
2852
2853    std::vector<whisper_token_data> tokens_cur;
2854    tokens_cur.reserve(whisper_n_text_ctx(ctx));
2855
2856    std::vector<whisper_token> prompt;
2857    prompt.reserve(whisper_n_text_ctx(ctx));
2858
2859    // main loop
2860    int seek = seek_start;
2861    while (true) {
2862        const int progress_cur = (100*(seek - seek_start))/(seek_end - seek_start);
2863        while (progress_cur >= progress_prev + progress_step) {
2864            progress_prev += progress_step;
2865            if (params.print_progress) {
2866				logInfo( u8"%s: progress = %3d%%", __func__, progress_prev );
2867            }
2868        }
2869
2870        // of only 1 second left, then stop
2871        if (seek + 100 >= seek_end) {
2872            break;
2873        }
2874
2875        // if there is a very short audio segment left to process, we remove any past prompt since it tends
2876        // to confuse the decoder and often make it repeat or hallucinate stuff
2877        if (seek > seek_start && seek + 500 >= seek_end) {
2878            prompt_past.clear();
2879        }
2880
2881        if (params.encoder_begin_callback) {
2882            if (params.encoder_begin_callback(ctx, params.encoder_begin_callback_user_data) == false) {
2883				logDebug( u8"%s: encoder_begin_callback returned false - aborting", __func__ );
2884                break;
2885            }
2886        }
2887
2888        // encode audio features starting at offset seek
2889        if (whisper_encode(ctx, seek, params.n_threads) != 0) {
2890			logError( u8"%s: failed to encode", __func__ );
2891            return -4;
2892        }
2893
2894        int n_past = 0;
2895        prompt.clear();
2896
2897        // if we have already generated some text, use it as a prompt to condition the next generation
2898        if (!prompt_past.empty()) {
2899            int n_take = std::min(std::min(params.n_max_text_ctx, whisper_n_text_ctx(ctx)/2), int(prompt_past.size()));
2900
2901            prompt = { whisper_token_prev(ctx) };
2902            prompt.insert(prompt.begin() + 1, prompt_past.end() - n_take, prompt_past.end());
2903
2904            prompt_past.clear();
2905            prompt_past.insert(prompt_past.end(), prompt.begin() + 1, prompt.end());
2906        }
2907
2908        prompt.insert(prompt.end(), prompt_init.begin(), prompt_init.end());
2909
2910        int seek_delta = 100*WHISPER_CHUNK_SIZE;
2911
2912        // print the prompt
2913        //printf("\n\n");
2914        //for (int i = 0; i < prompt.size(); i++) {
2915        //    printf("%s: prompt[%d] = %s\n", __func__, i, ctx->vocab.id_to_token[prompt[i]].c_str());
2916        //}
2917        //printf("\n\n");
2918
2919        // the accumulated transcription in the current interation
2920        int result_len = 0;
2921        tokens_cur.clear();
2922
2923        bool failed = false;
2924        bool has_ts = false; // have we already sampled a non-beg timestamp token for the current segment?
2925
2926        for (int i = 0, n_max = whisper_n_text_ctx(ctx)/2 - 4; i < n_max; ++i) {
2927            if (whisper_decode(ctx, prompt.data(), prompt.size(), n_past, params.n_threads) != 0) {
2928				logError( u8"%s: failed to decode", __func__ );
2929                return -5;
2930            }
2931
2932            n_past += prompt.size();
2933            prompt.clear();
2934
2935            // very basic greedy sampling strategy:
2936            //
2937            //   - always take the most probable token
2938            //
2939            // more sophisticated sampling strategies could be implemented here, but we keep it simple
2940            // feel free to experiment!
2941            //
2942            {
2943                const auto token = (i == 0) ? whisper_sample_timestamp(ctx, true) : whisper_sample_best(ctx);
2944
2945                // timestamp token - update sliding window
2946                if (token.id > whisper_token_beg(ctx)) {
2947                    const int seek_delta_new = 2*(token.id - whisper_token_beg(ctx));
2948
2949                    // do not allow to go back in time
2950                    if (has_ts && seek_delta > seek_delta_new && result_len < i) {
2951                        break;
2952                    }
2953
2954                    seek_delta = seek_delta_new;
2955                    result_len = i + 1;
2956                    has_ts = true;
2957                }
2958
2959                // add it to the context
2960                prompt.push_back(token.id);
2961                tokens_cur.push_back(token);
2962
2963                //{
2964                //    const auto tt = token.pt > 0.10 ? ctx->vocab.id_to_token[token.tid] : "[?]";
2965                //    printf("%s: %3d %10s %6d %6.3f '%s'\n", __func__, i, tt.c_str(), token.id, token.pt, ctx->vocab.id_to_token[token.id].c_str());
2966                //}
2967
2968                // end of segment
2969                if (token.id == whisper_token_eot(ctx) ||                // end of text token
2970                    (params.max_tokens > 0 && i >= params.max_tokens) || // max tokens per segment reached
2971                    (has_ts && seek + seek_delta + 100 >= seek_end)      // end of audio reached
2972                    ) {
2973                    if (result_len == 0) {
2974                        if (seek + seek_delta + 100 >= seek_end) {
2975                            result_len = i + 1;
2976                        } else {
2977                            failed = true;
2978                            break;
2979                        }
2980                    }
2981
2982                    if (params.single_segment) {
2983                        result_len = i + 1;
2984                        seek_delta = 100*WHISPER_CHUNK_SIZE;
2985                    }
2986
2987                    break;
2988                }
2989
2990                // TESTS: if no tensors are loaded, it means we are running tests
2991                if (ctx->model.n_loaded == 0) {
2992                    seek_delta = 100*WHISPER_CHUNK_SIZE;
2993                    break;
2994                }
2995            }
2996
2997            // sometimes, the decoding can get stuck in a repetition loop
2998            // this is a simple strategy to avoid such cases - we simply flag the decoding as failed and advance
2999            // the sliding window by 1 second
3000            if (i == n_max - 1 && (result_len == 0 || seek_delta < 100*WHISPER_CHUNK_SIZE/2)) {
3001                failed = true;
3002                break;
3003            }
3004        }
3005
3006        if (failed) {
3007            // when we fail to sample timestamp token, retry by clearing the past prompt
3008            // if it fails again, then we advance the window by 1 second
3009            if (!prompt_past.empty()) {
3010                prompt_past.clear();
3011            } else {
3012				logWarning( u8"%s: failed to generate timestamp token - skipping one second", __func__ );
3013                seek += 100;
3014            }
3015            continue;
3016        }
3017
3018        // shrink down to result_len
3019        tokens_cur.resize(result_len);
3020
3021        for (const auto & r : tokens_cur) {
3022            prompt_past.push_back(r.id);
3023        }
3024
3025        // store the text from this iteration
3026        if (!tokens_cur.empty()) {
3027            int  i0 = 0;
3028            auto t0 = seek + 2*(tokens_cur.front().tid - whisper_token_beg(ctx));
3029
3030            std::string text;
3031
3032            for (int i = 0; i < (int) tokens_cur.size(); i++) {
3033                //printf("%s: %18s %6.3f %18s %6.3f\n", __func__,
3034                //        ctx->vocab.id_to_token[tokens_cur[i].id].c_str(), tokens_cur[i].p,
3035                //        ctx->vocab.id_to_token[tokens_cur[i].tid].c_str(), tokens_cur[i].pt);
3036
3037                if (params.print_special == false && tokens_cur[i].id >= whisper_token_eot(ctx)) {
3038                } else {
3039                    text += whisper_token_to_str(ctx, tokens_cur[i].id);
3040                }
3041                if (tokens_cur[i].id > whisper_token_beg(ctx) && !params.single_segment) {
3042                    const auto t1 = seek + 2*(tokens_cur[i].tid - whisper_token_beg(ctx));
3043                    if (!text.empty()) {
3044                        const auto tt0 = params.speed_up ? 2*t0 : t0;
3045                        const auto tt1 = params.speed_up ? 2*t1 : t1;
3046
3047                        if (params.print_realtime) {
3048                            if (params.print_timestamps) {
3049                                printf("[%s --> %s]  %s\n", to_timestamp(tt0).c_str(), to_timestamp(tt1).c_str(), text.c_str());
3050                            } else {
3051                                printf("%s", text.c_str());
3052                                fflush(stdout);
3053                            }
3054                        }
3055
3056                        result_all.push_back({ tt0, tt1, text, {} });
3057                        for (int j = i0; j <= i; j++) {
3058                            result_all.back().tokens.push_back(tokens_cur[j]);
3059                        }
3060
3061                        int n_new = 1;
3062
3063                        if (params.token_timestamps) {
3064                            whisper_exp_compute_token_level_timestamps(
3065                                    ctx, result_all.size() - 1, params.thold_pt, params.thold_ptsum);
3066
3067                            if (params.max_len > 0) {
3068                                n_new = whisper_wrap_segment(ctx, params.max_len);
3069                            }
3070                        }
3071                        if (params.new_segment_callback) {
3072                            params.new_segment_callback(ctx, n_new, params.new_segment_callback_user_data);
3073                        }
3074                    }
3075                    text = "";
3076                    while (i < (int) tokens_cur.size() && tokens_cur[i].id > whisper_token_beg(ctx)) {
3077                        i++;
3078                    }
3079                    i--;
3080                    t0 = t1;
3081                    i0 = i + 1;
3082                }
3083            }
3084
3085            if (!text.empty()) {
3086                const auto t1 = seek + seek_delta;
3087
3088                const auto tt0 = params.speed_up ? 2*t0 : t0;
3089                const auto tt1 = params.speed_up ? 2*t1 : t1;
3090
3091                if (params.print_realtime) {
3092                    if (params.print_timestamps) {
3093                        printf("[%s --> %s]  %s\n", to_timestamp(tt0).c_str(), to_timestamp(tt1).c_str(), text.c_str());
3094                    } else {
3095                        printf("%s", text.c_str());
3096                        fflush(stdout);
3097                    }
3098                }
3099
3100                result_all.push_back({ tt0, tt1, text, {} });
3101                for (int j = i0; j < (int) tokens_cur.size(); j++) {
3102                    result_all.back().tokens.push_back(tokens_cur[j]);
3103                }
3104
3105                int n_new = 1;
3106
3107                if (params.token_timestamps) {
3108                    whisper_exp_compute_token_level_timestamps(
3109                            ctx, result_all.size() - 1, params.thold_pt, params.thold_ptsum);
3110
3111                    if (params.max_len > 0) {
3112                        n_new = whisper_wrap_segment(ctx, params.max_len);
3113                    }
3114                }
3115                if (params.new_segment_callback) {
3116                    params.new_segment_callback(ctx, n_new, params.new_segment_callback_user_data);
3117                }
3118            }
3119        }
3120
3121        seek += seek_delta;
3122    }
3123
3124    return 0;
3125}
3126
3127int whisper_full_parallel(
3128        struct whisper_context * ctx,
3129        struct whisper_full_params params,
3130        const float * samples,
3131        int n_samples,
3132        int n_processors) {
3133    if (n_processors == 1) {
3134        return whisper_full(ctx, params, samples, n_samples);
3135    }
3136
3137    int ret = 0;
3138
3139    // prepare separate contexts for each thread
3140    std::vector<struct whisper_context> ctxs(n_processors - 1);
3141
3142    for (int i = 0; i < n_processors - 1; ++i) {
3143        ctxs[i] = *ctx;
3144
3145        auto & model = ctxs[i].model;
3146
3147        // create the ggml memory context
3148        {
3149            struct ggml_init_params params;
3150            params.mem_size   = ctxs[i].buf_memory.size();
3151            params.mem_buffer = ctxs[i].buf_memory.data();
3152
3153            model.ctx_mem = ggml_init(params);
3154            if (!model.ctx_mem) {
3155				logError( u8"%s: ggml_init() failed", __func__ );
3156                return false;
3157            }
3158        }
3159
3160        // separate key + value memory for each processor
3161        {
3162            auto & ctx = model.ctx_mem;
3163
3164            const auto & hparams = model.hparams;
3165
3166            const int n_text_state = hparams.n_text_state;
3167            const int n_text_layer = hparams.n_text_layer;
3168            const int n_text_ctx   = hparams.n_text_ctx;
3169
3170            // key/value memory for the self-attention layer
3171            {
3172                const int n_mem      = n_text_layer*n_text_ctx;
3173                const int n_elements = n_text_state*n_mem;
3174
3175                model.memory_k = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
3176                model.memory_v = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
3177            }
3178
3179            // key/value memory for the cross-attention layer
3180            {
3181                const int n_audio_ctx = hparams.n_audio_ctx;
3182
3183                const int n_mem      = n_text_layer*n_audio_ctx;
3184                const int n_elements = n_text_state*n_mem;
3185
3186                model.memory_cross_k = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
3187                model.memory_cross_v = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, n_elements);
3188            }
3189        }
3190    }
3191
3192    const int offset_samples = (WHISPER_SAMPLE_RATE*params.offset_ms)/1000;
3193    const int n_samples_per_processor = (n_samples - offset_samples)/n_processors;
3194
3195    // the calling thread will process the first chunk
3196    // while the other threads will process the remaining chunks
3197
3198    std::vector<std::thread> workers(n_processors - 1);
3199    for (int i = 0; i < n_processors - 1; ++i) {
3200        const int start_samples = offset_samples + (i + 1)*n_samples_per_processor;
3201        const int n_samples_cur = (i == n_processors - 2) ? n_samples - start_samples : n_samples_per_processor;
3202
3203        auto params_cur = params;
3204
3205        params_cur.offset_ms = 0;
3206        params_cur.print_progress = false;
3207        params_cur.print_realtime = false;
3208
3209        params_cur.new_segment_callback = nullptr;
3210        params_cur.new_segment_callback_user_data = nullptr;
3211
3212        workers[i] = std::thread(whisper_full, &ctxs[i], std::move(params_cur), samples + start_samples, n_samples_cur);
3213    }
3214
3215    {
3216        auto params_cur = params;
3217
3218        ret = whisper_full(ctx, std::move(params_cur), samples, offset_samples + n_samples_per_processor);
3219    }
3220
3221    for (int i = 0; i < n_processors - 1; ++i) {
3222        workers[i].join();
3223    }
3224
3225    const int64_t offset_t = (int64_t) params.offset_ms/10.0;
3226
3227    // combine results into ctx->result_all
3228    for (int i = 0; i < n_processors - 1; ++i) {
3229        auto & results_i = ctxs[i].result_all;
3230
3231        for (int j = 0; j < (int) results_i.size(); ++j) {
3232            // correct the segment timestamp taking into account the offset
3233            results_i[j].t0 += 100*((i + 1)*n_samples_per_processor)/WHISPER_SAMPLE_RATE + offset_t;
3234            results_i[j].t1 += 100*((i + 1)*n_samples_per_processor)/WHISPER_SAMPLE_RATE + offset_t;
3235
3236            // make sure that segments are not overlapping
3237            if (!ctx->result_all.empty()) {
3238                results_i[j].t0 = std::max(results_i[j].t0, ctx->result_all.back().t1);
3239            }
3240
3241            ctx->result_all.push_back(std::move(results_i[j]));
3242
3243            // call the new_segment_callback for each segment
3244            if (params.new_segment_callback) {
3245                params.new_segment_callback(ctx, 1, params.new_segment_callback_user_data);
3246            }
3247        }
3248
3249        ctx->t_mel_us    += ctxs[i].t_mel_us;
3250        ctx->t_sample_us += ctxs[i].t_sample_us;
3251        ctx->t_encode_us += ctxs[i].t_encode_us;
3252        ctx->t_decode_us += ctxs[i].t_decode_us;
3253    }
3254
3255    // average the timings
3256    ctx->t_mel_us    /= n_processors;
3257    ctx->t_sample_us /= n_processors;
3258    ctx->t_encode_us /= n_processors;
3259    ctx->t_decode_us /= n_processors;
3260
3261    // print information about the audio boundaries
3262	logDebug( u8"%s: the audio has been split into %d chunks at the following times:", __func__, n_processors );
3263	for( int i = 0; i < n_processors - 1; ++i )
3264		logDebug( u8"%s: split %d - %s", __func__, ( i + 1 ), to_timestamp( 100 * ( ( i + 1 ) * n_samples_per_processor ) / WHISPER_SAMPLE_RATE + offset_t ).c_str() );
3265	logDebug( u8"%s: the transcription quality may be degraded near these boundaries", __func__ );
3266
3267    return ret;
3268}
3269
3270int whisper_full_n_segments(struct whisper_context * ctx) {
3271    return ctx->result_all.size();
3272}
3273
3274int64_t whisper_full_get_segment_t0(struct whisper_context * ctx, int i_segment) {
3275    return ctx->result_all[i_segment].t0;
3276}
3277
3278int64_t whisper_full_get_segment_t1(struct whisper_context * ctx, int i_segment) {
3279    return ctx->result_all[i_segment].t1;
3280}
3281
3282const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment) {
3283    return ctx->result_all[i_segment].text.c_str();
3284}
3285
3286int whisper_full_n_tokens(struct whisper_context * ctx, int i_segment) {
3287    return ctx->result_all[i_segment].tokens.size();
3288}
3289
3290const char * whisper_full_get_token_text(struct whisper_context * ctx, int i_segment, int i_token) {
3291    return ctx->vocab.id_to_token[ctx->result_all[i_segment].tokens[i_token].id].c_str();
3292}
3293
3294whisper_token whisper_full_get_token_id(struct whisper_context * ctx, int i_segment, int i_token) {
3295    return ctx->result_all[i_segment].tokens[i_token].id;
3296}
3297
3298struct whisper_token_data whisper_full_get_token_data(struct whisper_context * ctx, int i_segment, int i_token) {
3299    return ctx->result_all[i_segment].tokens[i_token];
3300}
3301
3302float whisper_full_get_token_p(struct whisper_context * ctx, int i_segment, int i_token) {
3303    return ctx->result_all[i_segment].tokens[i_token].p;
3304}
3305
3306// =================================================================================================
3307
3308//
3309// Experimental stuff below
3310//
3311// Not sure if these should be part of the library at all, because the quality of the results is not
3312// guaranteed. Might get removed at some point unless a robust algorithm implementation is found
3313//
3314
3315// =================================================================================================
3316
3317//
3318// token-level timestamps
3319//
3320
3321static int timestamp_to_sample(int64_t t, int n_samples) {
3322    return std::max(0, std::min((int) n_samples - 1, (int) ((t*WHISPER_SAMPLE_RATE)/100)));
3323}
3324
3325static int64_t sample_to_timestamp(int i_sample) {
3326    return (100*i_sample)/WHISPER_SAMPLE_RATE;
3327}
3328
3329// a cost-function / heuristic that is high for text that takes longer to pronounce
3330// obviously, can be improved
3331static float voice_length(const std::string & text) {
3332    float res = 0.0f;
3333
3334    for (size_t i = 0; i < text.size(); ++i) {
3335        if (text[i] == ' ') {
3336            res += 0.01f;
3337        } else if (text[i] == ',') {
3338            res += 2.00f;
3339        } else if (text[i] == '.') {
3340            res += 3.00f;
3341        } else if (text[i] == '!') {
3342            res += 3.00f;
3343        } else if (text[i] == '?') {
3344            res += 3.00f;
3345        } else if (text[i] >= '0' && text[i] <= '9') {
3346            res += 3.00f;
3347        } else {
3348            res += 1.00f;
3349        }
3350    }
3351
3352    return res;
3353}
3354
3355// average the fabs of the signal
3356static std::vector<float> get_signal_energy(const float * signal, int n_samples, int n_samples_per_half_window) {
3357    const int hw = n_samples_per_half_window;
3358
3359    std::vector<float> result(n_samples);
3360
3361    for (int i = 0; i < n_samples; i++) {
3362        float sum = 0;
3363        for (int j = -hw; j <= hw; j++) {
3364            if (i + j >= 0 && i + j < n_samples) {
3365                sum += fabs(signal[i + j]);
3366            }
3367        }
3368        result[i] = sum/(2*hw + 1);
3369    }
3370
3371    return result;
3372}
3373
3374static void whisper_exp_compute_token_level_timestamps(
3375        struct whisper_context * ctx,
3376        int   i_segment,
3377        float thold_pt,
3378        float thold_ptsum) {
3379    auto & segment = ctx->result_all[i_segment];
3380    auto & tokens  = segment.tokens;
3381
3382    const int n_samples = ctx->energy.size();
3383
3384    if (n_samples == 0) {
3385		logWarning( u8"%s: no signal data available", __func__ );
3386        return;
3387    }
3388
3389    const int64_t t0 = segment.t0;
3390    const int64_t t1 = segment.t1;
3391
3392    const int n = tokens.size();
3393
3394    if (n == 0) {
3395        return;
3396    }
3397
3398    if (n == 1) {
3399        tokens[0].t0 = t0;
3400        tokens[0].t1 = t1;
3401
3402        return;
3403    }
3404
3405    auto & t_beg    = ctx->t_beg;
3406    auto & t_last   = ctx->t_last;
3407    auto & tid_last = ctx->tid_last;
3408
3409    for (int j = 0; j < n; ++j) {
3410        auto & token = tokens[j];
3411
3412        if (j == 0) {
3413            if (token.id == whisper_token_beg(ctx)) {
3414                tokens[j    ].t0 = t0;
3415                tokens[j    ].t1 = t0;
3416                tokens[j + 1].t0 = t0;
3417
3418                t_beg    = t0;
3419                t_last   = t0;
3420                tid_last = whisper_token_beg(ctx);
3421            } else {
3422                tokens[j    ].t0 = t_last;
3423            }
3424        }
3425
3426        const int64_t tt = t_beg + 2*(token.tid - whisper_token_beg(ctx));
3427
3428        tokens[j].id    = token.id;
3429        tokens[j].tid   = token.tid;
3430        tokens[j].p     = token.p;
3431        tokens[j].pt    = token.pt;
3432        tokens[j].ptsum = token.ptsum;
3433
3434        tokens[j].vlen = voice_length(whisper_token_to_str(ctx, token.id));
3435
3436        if (token.pt > thold_pt && token.ptsum > thold_ptsum && token.tid > tid_last && tt <= t1) {
3437            if (j > 0) {
3438                tokens[j - 1].t1 = tt;
3439            }
3440            tokens[j].t0 = tt;
3441            tid_last = token.tid;
3442        }
3443    }
3444
3445    tokens[n - 2].t1 = t1;
3446    tokens[n - 1].t0 = t1;
3447    tokens[n - 1].t1 = t1;
3448
3449    t_last = t1;
3450
3451    // find intervals of tokens with unknown timestamps
3452    // fill the timestamps by proportionally splitting the interval based on the token voice lengths
3453    {
3454        int p0 = 0;
3455        int p1 = 0;
3456
3457        while (true) {
3458            while (p1 < n && tokens[p1].t1 < 0) {
3459                p1++;
3460            }
3461
3462            if (p1 >= n) {
3463                p1--;
3464            }
3465
3466            if (p1 > p0) {
3467                double psum = 0.0;
3468                for (int j = p0; j <= p1; j++) {
3469                    psum += tokens[j].vlen;
3470                }
3471
3472                //printf("analyzing %d - %d, psum = %f\n", p0, p1, psum);
3473
3474                const double dt = tokens[p1].t1 - tokens[p0].t0;
3475
3476                // split the time proportionally to the voice length
3477                for (int j = p0 + 1; j <= p1; j++) {
3478                    const double ct = tokens[j - 1].t0 + dt*tokens[j - 1].vlen/psum;
3479
3480                    tokens[j - 1].t1 = ct;
3481                    tokens[j    ].t0 = ct;
3482                }
3483            }
3484
3485            p1++;
3486            p0 = p1;
3487            if (p1 >= n) {
3488                break;
3489            }
3490        }
3491    }
3492
3493    // fix up (just in case)
3494    for (int j = 0; j < n - 1; j++) {
3495        if (tokens[j].t1 < 0) {
3496            tokens[j + 1].t0 = tokens[j].t1;
3497        }
3498
3499        if (j > 0) {
3500            if (tokens[j - 1].t1 > tokens[j].t0) {
3501                tokens[j].t0 = tokens[j - 1].t1;
3502                tokens[j].t1 = std::max(tokens[j].t0, tokens[j].t1);
3503            }
3504        }
3505    }
3506
3507    // VAD
3508    // expand or contract tokens based on voice activity
3509    {
3510        const int hw = WHISPER_SAMPLE_RATE/8;
3511
3512        for (int j = 0; j < n; j++) {
3513            if (tokens[j].id >= whisper_token_eot(ctx)) {
3514                continue;
3515            }
3516
3517            int s0 = timestamp_to_sample(tokens[j].t0, n_samples);
3518            int s1 = timestamp_to_sample(tokens[j].t1, n_samples);
3519
3520            const int ss0 = std::max(s0 - hw, 0);
3521            const int ss1 = std::min(s1 + hw, n_samples);
3522
3523            const int ns = ss1 - ss0;
3524
3525            float sum = 0.0f;
3526
3527            for (int k = ss0; k < ss1; k++) {
3528                sum += ctx->energy[k];
3529            }
3530
3531            const float thold = 0.5*sum/ns;
3532
3533            {
3534                int k = s0;
3535                if (ctx->energy[k] > thold && j > 0) {
3536                    while (k > 0 && ctx->energy[k] > thold) {
3537                        k--;
3538                    }
3539                    tokens[j].t0 = sample_to_timestamp(k);
3540                    if (tokens[j].t0 < tokens[j - 1].t1) {
3541                        tokens[j].t0 = tokens[j - 1].t1;
3542                    } else {
3543                        s0 = k;
3544                    }
3545                } else {
3546                    while (ctx->energy[k] < thold && k < s1) {
3547                        k++;
3548                    }
3549                    s0 = k;
3550                    tokens[j].t0 = sample_to_timestamp(k);
3551                }
3552            }
3553
3554            {
3555                int k = s1;
3556                if (ctx->energy[k] > thold) {
3557                    while (k < n_samples - 1 && ctx->energy[k] > thold) {
3558                        k++;
3559                    }
3560                    tokens[j].t1 = sample_to_timestamp(k);
3561                    if (j < ns - 1 && tokens[j].t1 > tokens[j + 1].t0) {
3562                        tokens[j].t1 = tokens[j + 1].t0;
3563                    } else {
3564                        s1 = k;
3565                    }
3566                } else {
3567                    while (ctx->energy[k] < thold && k > s0) {
3568                        k--;
3569                    }
3570                    s1 = k;
3571                    tokens[j].t1 = sample_to_timestamp(k);
3572                }
3573            }
3574        }
3575    }
3576
3577    // fixed token expand (optional)
3578    //{
3579    //    const int t_expand = 0;
3580
3581    //    for (int j = 0; j < n; j++) {
3582    //        if (j > 0) {
3583    //            tokens[j].t0 = std::max(0, (int) (tokens[j].t0 - t_expand));
3584    //        }
3585    //        if (j < n - 1) {
3586    //            tokens[j].t1 = tokens[j].t1 + t_expand;
3587    //        }
3588    //    }
3589    //}
3590
3591    // debug info
3592    //for (int j = 0; j < n; ++j) {
3593    //    const auto & token = tokens[j];
3594    //    const auto tt = token.pt > thold_pt && token.ptsum > 0.01 ? whisper_token_to_str(ctx, token.tid) : "[?]";
3595    //    printf("%s: %10s %6.3f %6.3f %6.3f %6.3f %5d %5d '%s'\n", __func__,
3596    //            tt, token.p, token.pt, token.ptsum, token.vlen, (int) token.t0, (int) token.t1, whisper_token_to_str(ctx, token.id));
3597
3598    //    if (tokens[j].id >= whisper_token_eot(ctx)) {
3599    //        continue;
3600    //    }
3601    //}
3602}