FrankZxShen commited on
Commit
8913fdd
1 Parent(s): ce252ec
.gitattributes CHANGED
@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
32
  *.zip filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
32
  *.zip filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
35
+ pretrain/nsf_hifigan/model filter=lfs diff=lfs merge=lfs -text
pretrain/checkpoint_best_legacy_500.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:60d936ec5a566776fc392e69ad8b630d14eb588111233fe313436e200a7b187b
3
+ size 1330114945
pretrain/meta.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ def download_dict():
2
+ return {
3
+ "vec768l12": {
4
+ "url": "https://ibm.ent.box.com/shared/static/z1wgl1stco8ffooyatzdwsqn2psd9lrr",
5
+ "output": "./pretrain/checkpoint_best_legacy_500.pt"
6
+ },
7
+ "vec256l9": {
8
+ "url": "https://ibm.ent.box.com/shared/static/z1wgl1stco8ffooyatzdwsqn2psd9lrr",
9
+ "output": "./pretrain/checkpoint_best_legacy_500.pt"
10
+ },
11
+ "hubertsoft": {
12
+ "url": "https://github.com/bshall/hubert/releases/download/v0.1/hubert-soft-0d54a1f4.pt",
13
+ "output": "./pretrain/hubert-soft-0d54a1f4.pt"
14
+ },
15
+ "whisper-ppg": {
16
+ "url": "https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt",
17
+ "output": "./pretrain/medium.pt"
18
+ }
19
+ }
20
+
21
+
22
+ def get_speech_encoder(config_path="configs/config.json"):
23
+ import json
24
+
25
+ with open(config_path, "r") as f:
26
+ data = f.read()
27
+ config = json.loads(data)
28
+ speech_encoder = config["model"]["speech_encoder"]
29
+ dict = download_dict()
30
+
31
+ return dict[speech_encoder]["url"], dict[speech_encoder]["output"]
pretrain/nsf_hifigan/NOTICE.txt ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ --- DiffSinger Community Vocoder ---
2
+
3
+ ARCHITECTURE: NSF-HiFiGAN
4
+ RELEASE DATE: 2022-12-11
5
+
6
+ HYPER PARAMETERS:
7
+ - 44100 sample rate
8
+ - 128 mel bins
9
+ - 512 hop size
10
+ - 2048 window size
11
+ - fmin at 40Hz
12
+ - fmax at 16000Hz
13
+
14
+
15
+ NOTICE:
16
+
17
+ All model weights in the [DiffSinger Community Vocoder Project](https://openvpi.github.io/vocoders/), including
18
+ model weights in this directory, are provided by the [OpenVPI Team](https://github.com/openvpi/), under the
19
+ [Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/) license.
20
+
21
+
22
+ ACKNOWLEDGEMENTS:
23
+
24
+ Training data of this vocoder is provided and permitted by the following organizations, societies and individuals:
25
+
26
+ 孙飒 https://www.qfssr.cn
27
+ 赤松_Akamatsu https://www.zhibin.club
28
+ 乐威 https://www.zhibin.club
29
+ 伯添 https://space.bilibili.com/24087011
30
+ 雲宇光 https://space.bilibili.com/660675050
31
+ 橙子言 https://space.bilibili.com/318486464
32
+ 人衣大人 https://space.bilibili.com/2270344
33
+ 玖蝶 https://space.bilibili.com/676771003
34
+ Yuuko
35
+ 白夜零BYL https://space.bilibili.com/1605040503
36
+ 嗷天 https://space.bilibili.com/5675252
37
+ 洛泠羽 https://space.bilibili.com/347373318
38
+ 灰条纹的灰猫君 https://space.bilibili.com/2083633
39
+ 幽寂 https://space.bilibili.com/478860
40
+ 恶魔王女 https://space.bilibili.com/2475098
41
+ AlexYHX 芮晴
42
+ 绮萱 https://y.qq.com/n/ryqq/singer/003HjD6H4aZn1K
43
+ 诗芸 https://y.qq.com/n/ryqq/singer/0005NInj142zm0
44
+ 汐蕾 https://y.qq.com/n/ryqq/singer/0023cWMH1Bq1PJ
45
+ 1262917464
46
+ 炜阳
47
+ 叶卡yolka
48
+ 幸の夏 https://space.bilibili.com/1017297686
49
+ 暮色未量 https://space.bilibili.com/272904686
50
+ 晓寞sama https://space.bilibili.com/3463394
51
+ 没头绪的节操君
52
+ 串串BunC https://space.bilibili.com/95817834
53
+ 落雨 https://space.bilibili.com/1292427
54
+ 长尾巴的翎艾 https://space.bilibili.com/1638666
55
+ 声闻计划 https://space.bilibili.com/392812269
56
+ 唐家大小姐 http://5sing.kugou.com/palmusic/default.html
57
+ 不伊子
58
+
59
+ Training machines are provided by:
60
+
61
+ 花儿不哭 https://space.bilibili.com/5760446
62
+
63
+
64
+ TERMS OF REDISTRIBUTIONS:
65
+
66
+ 1. Do not sell this vocoder, or charge any fees from redistributing it, as prohibited by
67
+ the license.
68
+ 2. Include a copy of the CC BY-NC-SA 4.0 license, or a link referring to it.
69
+ 3. Include a copy of this notice, or any other notices informing that this vocoder is
70
+ provided by the OpenVPI Team, that this vocoder is licensed under CC BY-NC-SA 4.0, and
71
+ with a complete acknowledgement list as shown above.
72
+ 4. If you fine-tuned or modified the weights, leave a notice about what has been changed.
73
+ 5. (Optional) Leave a link to the official release page of the vocoder, and tell users
74
+ that other versions and future updates of this vocoder can be obtained from the website.
pretrain/nsf_hifigan/NOTICE.zh-CN.txt ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ --- DiffSinger 社区声码器 ---
2
+
3
+ 架构:NSF-HiFiGAN
4
+ 发布日期:2022-12-11
5
+
6
+ 超参数:
7
+ - 44100 sample rate
8
+ - 128 mel bins
9
+ - 512 hop size
10
+ - 2048 window size
11
+ - fmin at 40Hz
12
+ - fmax at 16000Hz
13
+
14
+
15
+ 注意事项:
16
+
17
+ [DiffSinger 社区声码器企划](https://openvpi.github.io/vocoders/) 中的所有模型权重,
18
+ 包括此目录下的模型权重,均由 [OpenVPI Team](https://github.com/openvpi/) 提供,并基于
19
+ [Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/)
20
+ 进行许可。
21
+
22
+
23
+ 致谢:
24
+
25
+ 此声码器的训练数据由以下组织、社团和个人提供并许可:
26
+
27
+ 孙飒 https://www.qfssr.cn
28
+ 赤松_Akamatsu https://www.zhibin.club
29
+ 乐威 https://www.zhibin.club
30
+ 伯添 https://space.bilibili.com/24087011
31
+ 雲宇光 https://space.bilibili.com/660675050
32
+ 橙子言 https://space.bilibili.com/318486464
33
+ 人衣大人 https://space.bilibili.com/2270344
34
+ 玖蝶 https://space.bilibili.com/676771003
35
+ Yuuko
36
+ 白夜零BYL https://space.bilibili.com/1605040503
37
+ 嗷天 https://space.bilibili.com/5675252
38
+ 洛泠羽 https://space.bilibili.com/347373318
39
+ 灰条纹的灰猫君 https://space.bilibili.com/2083633
40
+ 幽寂 https://space.bilibili.com/478860
41
+ 恶魔王女 https://space.bilibili.com/2475098
42
+ AlexYHX 芮晴
43
+ 绮萱 https://y.qq.com/n/ryqq/singer/003HjD6H4aZn1K
44
+ 诗芸 https://y.qq.com/n/ryqq/singer/0005NInj142zm0
45
+ 汐蕾 https://y.qq.com/n/ryqq/singer/0023cWMH1Bq1PJ
46
+ 1262917464
47
+ 炜阳
48
+ 叶卡yolka
49
+ 幸の夏 https://space.bilibili.com/1017297686
50
+ 暮色未量 https://space.bilibili.com/272904686
51
+ 晓寞sama https://space.bilibili.com/3463394
52
+ 没头绪的节操君
53
+ 串串BunC https://space.bilibili.com/95817834
54
+ 落雨 https://space.bilibili.com/1292427
55
+ 长尾巴的翎艾 https://space.bilibili.com/1638666
56
+ 声闻计划 https://space.bilibili.com/392812269
57
+ 唐家大小姐 http://5sing.kugou.com/palmusic/default.html
58
+ 不伊子
59
+
60
+ 训练算力的提供者如下:
61
+
62
+ 花儿不哭 https://space.bilibili.com/5760446
63
+
64
+
65
+ 二次分发条款:
66
+
67
+ 1. 请勿售卖此声码器或从其二次分发过程中收取任何费用,因为此类行为受到许可证的禁止。
68
+ 2. 请在二次分发文件中包含一份 CC BY-NC-SA 4.0 许可证的副本或指向该许可证的链接。
69
+ 3. 请在二次分发文件中包含这份声明,或以其他形式声明此声码器由 OpenVPI Team 提供并基于 CC BY-NC-SA 4.0 许可,
70
+ 并附带上述完整的致谢名单。
71
+ 4. 如果您微调或修改了权重,请留下一份关于其受到了何种修改的说明。
72
+ 5.(可选)留下一份指向此声码器的官方发布页面的链接,并告知使用者可从该网站获取此声码器的其他版本和未来的更新。
pretrain/nsf_hifigan/config.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "resblock": "1",
3
+ "num_gpus": 4,
4
+ "batch_size": 10,
5
+ "learning_rate": 0.0002,
6
+ "adam_b1": 0.8,
7
+ "adam_b2": 0.99,
8
+ "lr_decay": 0.999,
9
+ "seed": 1234,
10
+
11
+ "upsample_rates": [ 8, 8, 2, 2, 2],
12
+ "upsample_kernel_sizes": [16,16, 4, 4, 4],
13
+ "upsample_initial_channel": 512,
14
+ "resblock_kernel_sizes": [3,7,11],
15
+ "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
16
+ "discriminator_periods": [3, 5, 7, 11, 17, 23, 37],
17
+
18
+ "segment_size": 16384,
19
+ "num_mels": 128,
20
+ "num_freq": 1025,
21
+ "n_fft" : 2048,
22
+ "hop_size": 512,
23
+ "win_size": 2048,
24
+
25
+ "sampling_rate": 44100,
26
+
27
+ "fmin": 40,
28
+ "fmax": 16000,
29
+ "fmax_for_loss": null,
30
+
31
+ "num_workers": 16,
32
+
33
+ "dist_config": {
34
+ "dist_backend": "nccl",
35
+ "dist_url": "tcp://localhost:54321",
36
+ "world_size": 1
37
+ }
38
+ }
pretrain/nsf_hifigan/model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2c576b63b7ed952161b70fad34e0562ace502ce689195520d8a2a6c051de29d6
3
+ size 56825430
pretrain/put_hubert_ckpt_here ADDED
File without changes
vencoder/ContentVec256L12_Onnx.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import onnxruntime
3
+ import torch
4
+
5
+ class ContentVec256L12_Onnx(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/vec-256-layer-12.onnx",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ self.hidden_dim = 256
9
+ if device is None:
10
+ self.dev = torch.device("cpu")
11
+ else:
12
+ self.dev = torch.device(device)
13
+ if device == 'cpu' or device == torch.device("cpu") or device is None:
14
+ providers = ['CPUExecutionProvider']
15
+ elif device == 'cuda' or device == torch.device("cuda"):
16
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
17
+ self.model = onnxruntime.InferenceSession(vec_path, providers=providers)
18
+
19
+ def encoder(self, wav):
20
+ feats = wav
21
+ if feats.dim() == 2: # double channels
22
+ feats = feats.mean(-1)
23
+ assert feats.dim() == 1, feats.dim()
24
+ feats = feats.view(1, -1)
25
+ feats = feats.unsqueeze(0).cpu().detach().numpy()
26
+ onnx_input = {self.model.get_inputs()[0].name: feats}
27
+ logits = self.model.run(None, onnx_input)
28
+ return torch.tensor(logits[0]).transpose(1, 2).to(self.dev)
vencoder/ContentVec256L9.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import torch
3
+ from fairseq import checkpoint_utils
4
+
5
+ class ContentVec256L9(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/checkpoint_best_legacy_500.pt",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
9
+ [vec_path],
10
+ suffix="",
11
+ )
12
+ self.hidden_dim = 256
13
+ if device is None:
14
+ self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
+ else:
16
+ self.dev = torch.device(device)
17
+ self.model = models[0].to(self.dev)
18
+ self.model.eval()
19
+
20
+ def encoder(self, wav):
21
+ feats = wav
22
+ if feats.dim() == 2: # double channels
23
+ feats = feats.mean(-1)
24
+ assert feats.dim() == 1, feats.dim()
25
+ feats = feats.view(1, -1)
26
+ padding_mask = torch.BoolTensor(feats.shape).fill_(False)
27
+ inputs = {
28
+ "source": feats.to(wav.device),
29
+ "padding_mask": padding_mask.to(wav.device),
30
+ "output_layer": 9, # layer 9
31
+ }
32
+ with torch.no_grad():
33
+ logits = self.model.extract_features(**inputs)
34
+ feats = self.model.final_proj(logits[0])
35
+ return feats.transpose(1, 2)
vencoder/ContentVec256L9_Onnx.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import onnxruntime
3
+ import torch
4
+
5
+ class ContentVec256L9_Onnx(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/vec-256-layer-9.onnx",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ self.hidden_dim = 256
9
+ if device is None:
10
+ self.dev = torch.device("cpu")
11
+ else:
12
+ self.dev = torch.device(device)
13
+ if device == 'cpu' or device == torch.device("cpu") or device is None:
14
+ providers = ['CPUExecutionProvider']
15
+ elif device == 'cuda' or device == torch.device("cuda"):
16
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
17
+ self.model = onnxruntime.InferenceSession(vec_path, providers=providers)
18
+
19
+ def encoder(self, wav):
20
+ feats = wav
21
+ if feats.dim() == 2: # double channels
22
+ feats = feats.mean(-1)
23
+ assert feats.dim() == 1, feats.dim()
24
+ feats = feats.view(1, -1)
25
+ feats = feats.unsqueeze(0).cpu().detach().numpy()
26
+ onnx_input = {self.model.get_inputs()[0].name: feats}
27
+ logits = self.model.run(None, onnx_input)
28
+ return torch.tensor(logits[0]).transpose(1, 2).to(self.dev)
vencoder/ContentVec768L12.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import torch
3
+ from fairseq import checkpoint_utils
4
+
5
+ class ContentVec768L12(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/checkpoint_best_legacy_500.pt",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ self.hidden_dim = 768
9
+ models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
10
+ [vec_path],
11
+ suffix="",
12
+ )
13
+ if device is None:
14
+ self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
+ else:
16
+ self.dev = torch.device(device)
17
+ self.model = models[0].to(self.dev)
18
+ self.model.eval()
19
+
20
+ def encoder(self, wav):
21
+ feats = wav
22
+ if feats.dim() == 2: # double channels
23
+ feats = feats.mean(-1)
24
+ assert feats.dim() == 1, feats.dim()
25
+ feats = feats.view(1, -1)
26
+ padding_mask = torch.BoolTensor(feats.shape).fill_(False)
27
+ inputs = {
28
+ "source": feats.to(wav.device),
29
+ "padding_mask": padding_mask.to(wav.device),
30
+ "output_layer": 12, # layer 12
31
+ }
32
+ with torch.no_grad():
33
+ logits = self.model.extract_features(**inputs)
34
+ return logits[0].transpose(1, 2)
vencoder/ContentVec768L12_Onnx.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import onnxruntime
3
+ import torch
4
+
5
+ class ContentVec768L12_Onnx(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/vec-768-layer-12.onnx",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ self.hidden_dim = 768
9
+ if device is None:
10
+ self.dev = torch.device("cpu")
11
+ else:
12
+ self.dev = torch.device(device)
13
+ if device == 'cpu' or device == torch.device("cpu") or device is None:
14
+ providers = ['CPUExecutionProvider']
15
+ elif device == 'cuda' or device == torch.device("cuda"):
16
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
17
+ self.model = onnxruntime.InferenceSession(vec_path, providers=providers)
18
+
19
+ def encoder(self, wav):
20
+ feats = wav
21
+ if feats.dim() == 2: # double channels
22
+ feats = feats.mean(-1)
23
+ assert feats.dim() == 1, feats.dim()
24
+ feats = feats.view(1, -1)
25
+ feats = feats.unsqueeze(0).cpu().detach().numpy()
26
+ onnx_input = {self.model.get_inputs()[0].name: feats}
27
+ logits = self.model.run(None, onnx_input)
28
+ return torch.tensor(logits[0]).transpose(1, 2).to(self.dev)
vencoder/ContentVec768L9_Onnx.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import onnxruntime
3
+ import torch
4
+
5
+ class ContentVec768L9_Onnx(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/vec-768-layer-9.onnx",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ self.hidden_dim = 768
9
+ if device is None:
10
+ self.dev = torch.device("cpu")
11
+ else:
12
+ self.dev = torch.device(device)
13
+ if device == 'cpu' or device == torch.device("cpu") or device is None:
14
+ providers = ['CPUExecutionProvider']
15
+ elif device == 'cuda' or device == torch.device("cuda"):
16
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
17
+ self.model = onnxruntime.InferenceSession(vec_path, providers=providers)
18
+
19
+ def encoder(self, wav):
20
+ feats = wav
21
+ if feats.dim() == 2: # double channels
22
+ feats = feats.mean(-1)
23
+ assert feats.dim() == 1, feats.dim()
24
+ feats = feats.view(1, -1)
25
+ feats = feats.unsqueeze(0).cpu().detach().numpy()
26
+ onnx_input = {self.model.get_inputs()[0].name: feats}
27
+ logits = self.model.run(None, onnx_input)
28
+ return torch.tensor(logits[0]).transpose(1, 2).to(self.dev)
vencoder/HubertSoft.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import torch
3
+ from vencoder.hubert import hubert_model
4
+ class HubertSoft(SpeechEncoder):
5
+ def __init__(self,vec_path = "pretrain/hubert-soft-0d54a1f4.pt",device=None):
6
+ print("load model(s) from {}".format(vec_path))
7
+ hubert_soft = hubert_model.hubert_soft(vec_path)
8
+ if device is None:
9
+ self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
+ else:
11
+ self.dev = torch.device(device)
12
+ self.hidden_dim = 256
13
+ self.model = hubert_soft.to(self.dev)
14
+
15
+ def encoder(self, wav):
16
+ feats = wav
17
+ if feats.dim() == 2: # double channels
18
+ feats = feats.mean(-1)
19
+ assert feats.dim() == 1, feats.dim()
20
+ feats = feats[None,None,:]
21
+ with torch.inference_mode():
22
+ units = self.model.units(feats)
23
+ return units.transpose(1,2)
vencoder/HubertSoft_Onnx.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import onnxruntime
3
+ import torch
4
+
5
+ class HubertSoft_Onnx(SpeechEncoder):
6
+ def __init__(self,vec_path = "pretrain/hubert-soft.onnx",device=None):
7
+ print("load model(s) from {}".format(vec_path))
8
+ self.hidden_dim = 256
9
+ if device is None:
10
+ self.dev = torch.device("cpu")
11
+ else:
12
+ self.dev = torch.device(device)
13
+ if device == 'cpu' or device == torch.device("cpu") or device is None:
14
+ providers = ['CPUExecutionProvider']
15
+ elif device == 'cuda' or device == torch.device("cuda"):
16
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
17
+ self.model = onnxruntime.InferenceSession(vec_path, providers=providers)
18
+
19
+ def encoder(self, wav):
20
+ feats = wav
21
+ if feats.dim() == 2: # double channels
22
+ feats = feats.mean(-1)
23
+ assert feats.dim() == 1, feats.dim()
24
+ feats = feats.view(1, -1)
25
+ feats = feats.unsqueeze(0).cpu().detach().numpy()
26
+ onnx_input = {self.model.get_inputs()[0].name: feats}
27
+ logits = self.model.run(None, onnx_input)
28
+ return torch.tensor(logits[0]).transpose(1, 2).to(self.dev)
vencoder/WhisperPPG.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from vencoder.encoder import SpeechEncoder
2
+ import torch
3
+
4
+ from vencoder.whisper.model import Whisper, ModelDimensions
5
+ from vencoder.whisper.audio import pad_or_trim, log_mel_spectrogram
6
+
7
+
8
+ class WhisperPPG(SpeechEncoder):
9
+ def __init__(self,vec_path = "pretrain/medium.pt",device=None):
10
+ if device is None:
11
+ self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
+ else:
13
+ self.dev = torch.device(device)
14
+ checkpoint = torch.load(vec_path, map_location=device)
15
+ dims = ModelDimensions(**checkpoint["dims"])
16
+ model = Whisper(dims)
17
+ model.load_state_dict(checkpoint["model_state_dict"])
18
+ self.hidden_dim = dims
19
+ self.model = model.to(self.dev)
20
+
21
+ def encoder(self, wav):
22
+ audio = wav
23
+ audln = audio.shape[0]
24
+ ppgln = audln // 320
25
+ audio = pad_or_trim(audio)
26
+ mel = log_mel_spectrogram(audio).to(self.dev)
27
+ with torch.no_grad():
28
+ ppg = self.model.encoder(mel.unsqueeze(0)).squeeze().data.cpu().float().numpy()
29
+ ppg = torch.FloatTensor(ppg[:ppgln,]).to(self.dev)
30
+ return ppg[None,:,:].transpose(1, 2)
vencoder/__init__.py ADDED
File without changes
vencoder/encoder.py ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ class SpeechEncoder(object):
2
+ def __init__(self,vec_path = "pretrain/checkpoint_best_legacy_500.pt",device=None):
3
+ self.model = None #This is Model
4
+ self.hidden_dim = 768
5
+ pass
6
+
7
+ def encoder(self,wav):
8
+ '''
9
+ input: wav:[batchsize,signal_length]
10
+ output: embedding:[batchsize,wav_frame,hidden_dim]
11
+ '''
12
+ pass
vencoder/hubert/__init__.py ADDED
File without changes
vencoder/hubert/hubert_model.py ADDED
@@ -0,0 +1,222 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import random
3
+ from typing import Optional, Tuple
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.nn.functional as t_func
8
+ from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
9
+
10
+
11
+ class Hubert(nn.Module):
12
+ def __init__(self, num_label_embeddings: int = 100, mask: bool = True):
13
+ super().__init__()
14
+ self._mask = mask
15
+ self.feature_extractor = FeatureExtractor()
16
+ self.feature_projection = FeatureProjection()
17
+ self.positional_embedding = PositionalConvEmbedding()
18
+ self.norm = nn.LayerNorm(768)
19
+ self.dropout = nn.Dropout(0.1)
20
+ self.encoder = TransformerEncoder(
21
+ nn.TransformerEncoderLayer(
22
+ 768, 12, 3072, activation="gelu", batch_first=True
23
+ ),
24
+ 12,
25
+ )
26
+ self.proj = nn.Linear(768, 256)
27
+
28
+ self.masked_spec_embed = nn.Parameter(torch.FloatTensor(768).uniform_())
29
+ self.label_embedding = nn.Embedding(num_label_embeddings, 256)
30
+
31
+ def mask(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
32
+ mask = None
33
+ if self.training and self._mask:
34
+ mask = _compute_mask((x.size(0), x.size(1)), 0.8, 10, x.device, 2)
35
+ x[mask] = self.masked_spec_embed.to(x.dtype)
36
+ return x, mask
37
+
38
+ def encode(
39
+ self, x: torch.Tensor, layer: Optional[int] = None
40
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
41
+ x = self.feature_extractor(x)
42
+ x = self.feature_projection(x.transpose(1, 2))
43
+ x, mask = self.mask(x)
44
+ x = x + self.positional_embedding(x)
45
+ x = self.dropout(self.norm(x))
46
+ x = self.encoder(x, output_layer=layer)
47
+ return x, mask
48
+
49
+ def logits(self, x: torch.Tensor) -> torch.Tensor:
50
+ logits = torch.cosine_similarity(
51
+ x.unsqueeze(2),
52
+ self.label_embedding.weight.unsqueeze(0).unsqueeze(0),
53
+ dim=-1,
54
+ )
55
+ return logits / 0.1
56
+
57
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
58
+ x, mask = self.encode(x)
59
+ x = self.proj(x)
60
+ logits = self.logits(x)
61
+ return logits, mask
62
+
63
+
64
+ class HubertSoft(Hubert):
65
+ def __init__(self):
66
+ super().__init__()
67
+
68
+ @torch.inference_mode()
69
+ def units(self, wav: torch.Tensor) -> torch.Tensor:
70
+ wav = t_func.pad(wav, ((400 - 320) // 2, (400 - 320) // 2))
71
+ x, _ = self.encode(wav)
72
+ return self.proj(x)
73
+
74
+
75
+ class FeatureExtractor(nn.Module):
76
+ def __init__(self):
77
+ super().__init__()
78
+ self.conv0 = nn.Conv1d(1, 512, 10, 5, bias=False)
79
+ self.norm0 = nn.GroupNorm(512, 512)
80
+ self.conv1 = nn.Conv1d(512, 512, 3, 2, bias=False)
81
+ self.conv2 = nn.Conv1d(512, 512, 3, 2, bias=False)
82
+ self.conv3 = nn.Conv1d(512, 512, 3, 2, bias=False)
83
+ self.conv4 = nn.Conv1d(512, 512, 3, 2, bias=False)
84
+ self.conv5 = nn.Conv1d(512, 512, 2, 2, bias=False)
85
+ self.conv6 = nn.Conv1d(512, 512, 2, 2, bias=False)
86
+
87
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
88
+ x = t_func.gelu(self.norm0(self.conv0(x)))
89
+ x = t_func.gelu(self.conv1(x))
90
+ x = t_func.gelu(self.conv2(x))
91
+ x = t_func.gelu(self.conv3(x))
92
+ x = t_func.gelu(self.conv4(x))
93
+ x = t_func.gelu(self.conv5(x))
94
+ x = t_func.gelu(self.conv6(x))
95
+ return x
96
+
97
+
98
+ class FeatureProjection(nn.Module):
99
+ def __init__(self):
100
+ super().__init__()
101
+ self.norm = nn.LayerNorm(512)
102
+ self.projection = nn.Linear(512, 768)
103
+ self.dropout = nn.Dropout(0.1)
104
+
105
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
106
+ x = self.norm(x)
107
+ x = self.projection(x)
108
+ x = self.dropout(x)
109
+ return x
110
+
111
+
112
+ class PositionalConvEmbedding(nn.Module):
113
+ def __init__(self):
114
+ super().__init__()
115
+ self.conv = nn.Conv1d(
116
+ 768,
117
+ 768,
118
+ kernel_size=128,
119
+ padding=128 // 2,
120
+ groups=16,
121
+ )
122
+ self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
123
+
124
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
125
+ x = self.conv(x.transpose(1, 2))
126
+ x = t_func.gelu(x[:, :, :-1])
127
+ return x.transpose(1, 2)
128
+
129
+
130
+ class TransformerEncoder(nn.Module):
131
+ def __init__(
132
+ self, encoder_layer: nn.TransformerEncoderLayer, num_layers: int
133
+ ) -> None:
134
+ super(TransformerEncoder, self).__init__()
135
+ self.layers = nn.ModuleList(
136
+ [copy.deepcopy(encoder_layer) for _ in range(num_layers)]
137
+ )
138
+ self.num_layers = num_layers
139
+
140
+ def forward(
141
+ self,
142
+ src: torch.Tensor,
143
+ mask: torch.Tensor = None,
144
+ src_key_padding_mask: torch.Tensor = None,
145
+ output_layer: Optional[int] = None,
146
+ ) -> torch.Tensor:
147
+ output = src
148
+ for layer in self.layers[:output_layer]:
149
+ output = layer(
150
+ output, src_mask=mask, src_key_padding_mask=src_key_padding_mask
151
+ )
152
+ return output
153
+
154
+
155
+ def _compute_mask(
156
+ shape: Tuple[int, int],
157
+ mask_prob: float,
158
+ mask_length: int,
159
+ device: torch.device,
160
+ min_masks: int = 0,
161
+ ) -> torch.Tensor:
162
+ batch_size, sequence_length = shape
163
+
164
+ if mask_length < 1:
165
+ raise ValueError("`mask_length` has to be bigger than 0.")
166
+
167
+ if mask_length > sequence_length:
168
+ raise ValueError(
169
+ f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length} and `sequence_length`: {sequence_length}`"
170
+ )
171
+
172
+ # compute number of masked spans in batch
173
+ num_masked_spans = int(mask_prob * sequence_length / mask_length + random.random())
174
+ num_masked_spans = max(num_masked_spans, min_masks)
175
+
176
+ # make sure num masked indices <= sequence_length
177
+ if num_masked_spans * mask_length > sequence_length:
178
+ num_masked_spans = sequence_length // mask_length
179
+
180
+ # SpecAugment mask to fill
181
+ mask = torch.zeros((batch_size, sequence_length), device=device, dtype=torch.bool)
182
+
183
+ # uniform distribution to sample from, make sure that offset samples are < sequence_length
184
+ uniform_dist = torch.ones(
185
+ (batch_size, sequence_length - (mask_length - 1)), device=device
186
+ )
187
+
188
+ # get random indices to mask
189
+ mask_indices = torch.multinomial(uniform_dist, num_masked_spans)
190
+
191
+ # expand masked indices to masked spans
192
+ mask_indices = (
193
+ mask_indices.unsqueeze(dim=-1)
194
+ .expand((batch_size, num_masked_spans, mask_length))
195
+ .reshape(batch_size, num_masked_spans * mask_length)
196
+ )
197
+ offsets = (
198
+ torch.arange(mask_length, device=device)[None, None, :]
199
+ .expand((batch_size, num_masked_spans, mask_length))
200
+ .reshape(batch_size, num_masked_spans * mask_length)
201
+ )
202
+ mask_idxs = mask_indices + offsets
203
+
204
+ # scatter indices to mask
205
+ mask = mask.scatter(1, mask_idxs, True)
206
+
207
+ return mask
208
+
209
+
210
+ def hubert_soft(
211
+ path: str,
212
+ ) -> HubertSoft:
213
+ r"""HuBERT-Soft from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`.
214
+ Args:
215
+ path (str): path of a pretrained model
216
+ """
217
+ hubert = HubertSoft()
218
+ checkpoint = torch.load(path)
219
+ consume_prefix_in_state_dict_if_present(checkpoint, "module.")
220
+ hubert.load_state_dict(checkpoint)
221
+ hubert.eval()
222
+ return hubert
vencoder/hubert/hubert_model_onnx.py ADDED
@@ -0,0 +1,217 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import random
3
+ from typing import Optional, Tuple
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.nn.functional as t_func
8
+ from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
9
+
10
+
11
+ class Hubert(nn.Module):
12
+ def __init__(self, num_label_embeddings: int = 100, mask: bool = True):
13
+ super().__init__()
14
+ self._mask = mask
15
+ self.feature_extractor = FeatureExtractor()
16
+ self.feature_projection = FeatureProjection()
17
+ self.positional_embedding = PositionalConvEmbedding()
18
+ self.norm = nn.LayerNorm(768)
19
+ self.dropout = nn.Dropout(0.1)
20
+ self.encoder = TransformerEncoder(
21
+ nn.TransformerEncoderLayer(
22
+ 768, 12, 3072, activation="gelu", batch_first=True
23
+ ),
24
+ 12,
25
+ )
26
+ self.proj = nn.Linear(768, 256)
27
+
28
+ self.masked_spec_embed = nn.Parameter(torch.FloatTensor(768).uniform_())
29
+ self.label_embedding = nn.Embedding(num_label_embeddings, 256)
30
+
31
+ def mask(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
32
+ mask = None
33
+ if self.training and self._mask:
34
+ mask = _compute_mask((x.size(0), x.size(1)), 0.8, 10, x.device, 2)
35
+ x[mask] = self.masked_spec_embed.to(x.dtype)
36
+ return x, mask
37
+
38
+ def encode(
39
+ self, x: torch.Tensor, layer: Optional[int] = None
40
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
41
+ x = self.feature_extractor(x)
42
+ x = self.feature_projection(x.transpose(1, 2))
43
+ x, mask = self.mask(x)
44
+ x = x + self.positional_embedding(x)
45
+ x = self.dropout(self.norm(x))
46
+ x = self.encoder(x, output_layer=layer)
47
+ return x, mask
48
+
49
+ def logits(self, x: torch.Tensor) -> torch.Tensor:
50
+ logits = torch.cosine_similarity(
51
+ x.unsqueeze(2),
52
+ self.label_embedding.weight.unsqueeze(0).unsqueeze(0),
53
+ dim=-1,
54
+ )
55
+ return logits / 0.1
56
+
57
+
58
+ class HubertSoft(Hubert):
59
+ def __init__(self):
60
+ super().__init__()
61
+
62
+ def units(self, wav: torch.Tensor) -> torch.Tensor:
63
+ wav = t_func.pad(wav, ((400 - 320) // 2, (400 - 320) // 2))
64
+ x, _ = self.encode(wav)
65
+ return self.proj(x)
66
+
67
+ def forward(self, x):
68
+ return self.units(x)
69
+
70
+ class FeatureExtractor(nn.Module):
71
+ def __init__(self):
72
+ super().__init__()
73
+ self.conv0 = nn.Conv1d(1, 512, 10, 5, bias=False)
74
+ self.norm0 = nn.GroupNorm(512, 512)
75
+ self.conv1 = nn.Conv1d(512, 512, 3, 2, bias=False)
76
+ self.conv2 = nn.Conv1d(512, 512, 3, 2, bias=False)
77
+ self.conv3 = nn.Conv1d(512, 512, 3, 2, bias=False)
78
+ self.conv4 = nn.Conv1d(512, 512, 3, 2, bias=False)
79
+ self.conv5 = nn.Conv1d(512, 512, 2, 2, bias=False)
80
+ self.conv6 = nn.Conv1d(512, 512, 2, 2, bias=False)
81
+
82
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
83
+ x = t_func.gelu(self.norm0(self.conv0(x)))
84
+ x = t_func.gelu(self.conv1(x))
85
+ x = t_func.gelu(self.conv2(x))
86
+ x = t_func.gelu(self.conv3(x))
87
+ x = t_func.gelu(self.conv4(x))
88
+ x = t_func.gelu(self.conv5(x))
89
+ x = t_func.gelu(self.conv6(x))
90
+ return x
91
+
92
+
93
+ class FeatureProjection(nn.Module):
94
+ def __init__(self):
95
+ super().__init__()
96
+ self.norm = nn.LayerNorm(512)
97
+ self.projection = nn.Linear(512, 768)
98
+ self.dropout = nn.Dropout(0.1)
99
+
100
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
101
+ x = self.norm(x)
102
+ x = self.projection(x)
103
+ x = self.dropout(x)
104
+ return x
105
+
106
+
107
+ class PositionalConvEmbedding(nn.Module):
108
+ def __init__(self):
109
+ super().__init__()
110
+ self.conv = nn.Conv1d(
111
+ 768,
112
+ 768,
113
+ kernel_size=128,
114
+ padding=128 // 2,
115
+ groups=16,
116
+ )
117
+ self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
118
+
119
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
120
+ x = self.conv(x.transpose(1, 2))
121
+ x = t_func.gelu(x[:, :, :-1])
122
+ return x.transpose(1, 2)
123
+
124
+
125
+ class TransformerEncoder(nn.Module):
126
+ def __init__(
127
+ self, encoder_layer: nn.TransformerEncoderLayer, num_layers: int
128
+ ) -> None:
129
+ super(TransformerEncoder, self).__init__()
130
+ self.layers = nn.ModuleList(
131
+ [copy.deepcopy(encoder_layer) for _ in range(num_layers)]
132
+ )
133
+ self.num_layers = num_layers
134
+
135
+ def forward(
136
+ self,
137
+ src: torch.Tensor,
138
+ mask: torch.Tensor = None,
139
+ src_key_padding_mask: torch.Tensor = None,
140
+ output_layer: Optional[int] = None,
141
+ ) -> torch.Tensor:
142
+ output = src
143
+ for layer in self.layers[:output_layer]:
144
+ output = layer(
145
+ output, src_mask=mask, src_key_padding_mask=src_key_padding_mask
146
+ )
147
+ return output
148
+
149
+
150
+ def _compute_mask(
151
+ shape: Tuple[int, int],
152
+ mask_prob: float,
153
+ mask_length: int,
154
+ device: torch.device,
155
+ min_masks: int = 0,
156
+ ) -> torch.Tensor:
157
+ batch_size, sequence_length = shape
158
+
159
+ if mask_length < 1:
160
+ raise ValueError("`mask_length` has to be bigger than 0.")
161
+
162
+ if mask_length > sequence_length:
163
+ raise ValueError(
164
+ f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length} and `sequence_length`: {sequence_length}`"
165
+ )
166
+
167
+ # compute number of masked spans in batch
168
+ num_masked_spans = int(mask_prob * sequence_length / mask_length + random.random())
169
+ num_masked_spans = max(num_masked_spans, min_masks)
170
+
171
+ # make sure num masked indices <= sequence_length
172
+ if num_masked_spans * mask_length > sequence_length:
173
+ num_masked_spans = sequence_length // mask_length
174
+
175
+ # SpecAugment mask to fill
176
+ mask = torch.zeros((batch_size, sequence_length), device=device, dtype=torch.bool)
177
+
178
+ # uniform distribution to sample from, make sure that offset samples are < sequence_length
179
+ uniform_dist = torch.ones(
180
+ (batch_size, sequence_length - (mask_length - 1)), device=device
181
+ )
182
+
183
+ # get random indices to mask
184
+ mask_indices = torch.multinomial(uniform_dist, num_masked_spans)
185
+
186
+ # expand masked indices to masked spans
187
+ mask_indices = (
188
+ mask_indices.unsqueeze(dim=-1)
189
+ .expand((batch_size, num_masked_spans, mask_length))
190
+ .reshape(batch_size, num_masked_spans * mask_length)
191
+ )
192
+ offsets = (
193
+ torch.arange(mask_length, device=device)[None, None, :]
194
+ .expand((batch_size, num_masked_spans, mask_length))
195
+ .reshape(batch_size, num_masked_spans * mask_length)
196
+ )
197
+ mask_idxs = mask_indices + offsets
198
+
199
+ # scatter indices to mask
200
+ mask = mask.scatter(1, mask_idxs, True)
201
+
202
+ return mask
203
+
204
+
205
+ def hubert_soft(
206
+ path: str,
207
+ ) -> HubertSoft:
208
+ r"""HuBERT-Soft from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`.
209
+ Args:
210
+ path (str): path of a pretrained model
211
+ """
212
+ hubert = HubertSoft()
213
+ checkpoint = torch.load(path)
214
+ consume_prefix_in_state_dict_if_present(checkpoint, "module.")
215
+ hubert.load_state_dict(checkpoint)
216
+ hubert.eval()
217
+ return hubert
vencoder/whisper/__init__.py ADDED
File without changes
vencoder/whisper/audio.py ADDED
@@ -0,0 +1,125 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from functools import lru_cache
3
+ from typing import Union
4
+
5
+ import ffmpeg
6
+ import numpy as np
7
+ import torch
8
+ import torch.nn.functional as F
9
+
10
+ from .utils import exact_div
11
+
12
+ from librosa.filters import mel as librosa_mel_fn
13
+
14
+ # hard-coded audio hyperparameters
15
+ SAMPLE_RATE = 16000
16
+ N_FFT = 400
17
+ N_MELS = 80
18
+ HOP_LENGTH = 160
19
+ CHUNK_LENGTH = 30
20
+ N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000: number of samples in a chunk
21
+ N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH) # 3000: number of frames in a mel spectrogram input
22
+
23
+
24
+ def load_audio(file: str, sr: int = SAMPLE_RATE):
25
+ """
26
+ Open an audio file and read as mono waveform, resampling as necessary
27
+
28
+ Parameters
29
+ ----------
30
+ file: str
31
+ The audio file to open
32
+
33
+ sr: int
34
+ The sample rate to resample the audio if necessary
35
+
36
+ Returns
37
+ -------
38
+ A NumPy array containing the audio waveform, in float32 dtype.
39
+ """
40
+ try:
41
+ # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
42
+ # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
43
+ out, _ = (
44
+ ffmpeg.input(file, threads=0)
45
+ .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
46
+ .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
47
+ )
48
+ except ffmpeg.Error as e:
49
+ raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
50
+
51
+ return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
52
+
53
+
54
+ def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
55
+ """
56
+ Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
57
+ """
58
+ if torch.is_tensor(array):
59
+ if array.shape[axis] > length:
60
+ array = array.index_select(dim=axis, index=torch.arange(length, device=array.device))
61
+
62
+ if array.shape[axis] < length:
63
+ pad_widths = [(0, 0)] * array.ndim
64
+ pad_widths[axis] = (0, length - array.shape[axis])
65
+ array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
66
+ else:
67
+ if array.shape[axis] > length:
68
+ array = array.take(indices=range(length), axis=axis)
69
+
70
+ if array.shape[axis] < length:
71
+ pad_widths = [(0, 0)] * array.ndim
72
+ pad_widths[axis] = (0, length - array.shape[axis])
73
+ array = np.pad(array, pad_widths)
74
+
75
+ return array
76
+
77
+
78
+ @lru_cache(maxsize=None)
79
+ def mel_filters(device, n_mels: int = N_MELS) -> torch.Tensor:
80
+ """
81
+ load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
82
+ Allows decoupling librosa dependency; saved using:
83
+
84
+ np.savez_compressed(
85
+ "mel_filters.npz",
86
+ mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
87
+ )
88
+ """
89
+ assert n_mels == 80, f"Unsupported n_mels: {n_mels}"
90
+ return torch.from_numpy(librosa_mel_fn(sr=SAMPLE_RATE,n_fft=N_FFT,n_mels=n_mels)).to(device)
91
+
92
+
93
+ def log_mel_spectrogram(audio: Union[str, np.ndarray, torch.Tensor], n_mels: int = N_MELS):
94
+ """
95
+ Compute the log-Mel spectrogram of
96
+
97
+ Parameters
98
+ ----------
99
+ audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
100
+ The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
101
+
102
+ n_mels: int
103
+ The number of Mel-frequency filters, only 80 is supported
104
+
105
+ Returns
106
+ -------
107
+ torch.Tensor, shape = (80, n_frames)
108
+ A Tensor that contains the Mel spectrogram
109
+ """
110
+ if not torch.is_tensor(audio):
111
+ if isinstance(audio, str):
112
+ audio = load_audio(audio)
113
+ audio = torch.from_numpy(audio)
114
+
115
+ window = torch.hann_window(N_FFT).to(audio.device)
116
+ stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
117
+ magnitudes = stft[..., :-1].abs() ** 2
118
+
119
+ filters = mel_filters(audio.device, n_mels)
120
+ mel_spec = filters @ magnitudes
121
+
122
+ log_spec = torch.clamp(mel_spec, min=1e-10).log10()
123
+ log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
124
+ log_spec = (log_spec + 4.0) / 4.0
125
+ return log_spec
vencoder/whisper/decoding.py ADDED
@@ -0,0 +1,712 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass, field
2
+ from typing import Dict, List, Tuple, Iterable, Optional, Sequence, Union, TYPE_CHECKING
3
+
4
+ import numpy as np
5
+ import torch
6
+ import torch.nn.functional as F
7
+ from torch import Tensor
8
+ from torch.distributions import Categorical
9
+
10
+ from .audio import CHUNK_LENGTH
11
+ from .tokenizer import Tokenizer, get_tokenizer
12
+ from .utils import compression_ratio
13
+
14
+ if TYPE_CHECKING:
15
+ from .model import Whisper
16
+
17
+
18
+ @torch.no_grad()
19
+ def detect_language(model: "Whisper", mel: Tensor, tokenizer: Tokenizer = None) -> Tuple[Tensor, List[dict]]:
20
+ """
21
+ Detect the spoken language in the audio, and return them as list of strings, along with the ids
22
+ of the most probable language tokens and the probability distribution over all language tokens.
23
+ This is performed outside the main decode loop in order to not interfere with kv-caching.
24
+
25
+ Returns
26
+ -------
27
+ language_tokens : Tensor, shape = (n_audio,)
28
+ ids of the most probable language tokens, which appears after the startoftranscript token.
29
+ language_probs : List[Dict[str, float]], length = n_audio
30
+ list of dictionaries containing the probability distribution over all languages.
31
+ """
32
+ if tokenizer is None:
33
+ tokenizer = get_tokenizer(model.is_multilingual)
34
+ if tokenizer.language is None or tokenizer.language_token not in tokenizer.sot_sequence:
35
+ raise ValueError(f"This model doesn't have language tokens so it can't perform lang id")
36
+
37
+ single = mel.ndim == 2
38
+ if single:
39
+ mel = mel.unsqueeze(0)
40
+
41
+ # skip encoder forward pass if already-encoded audio features were given
42
+ if mel.shape[-2:] != (model.dims.n_audio_ctx, model.dims.n_audio_state):
43
+ mel = model.encoder(mel)
44
+
45
+ # forward pass using a single token, startoftranscript
46
+ n_audio = mel.shape[0]
47
+ x = torch.tensor([[tokenizer.sot]] * n_audio).to(mel.device) # [n_audio, 1]
48
+ logits = model.logits(x, mel)[:, 0]
49
+
50
+ # collect detected languages; suppress all non-language tokens
51
+ mask = torch.ones(logits.shape[-1], dtype=torch.bool)
52
+ mask[list(tokenizer.all_language_tokens)] = False
53
+ logits[:, mask] = -np.inf
54
+ language_tokens = logits.argmax(dim=-1)
55
+ language_token_probs = logits.softmax(dim=-1).cpu()
56
+ language_probs = [
57
+ {
58
+ c: language_token_probs[i, j].item()
59
+ for j, c in zip(tokenizer.all_language_tokens, tokenizer.all_language_codes)
60
+ }
61
+ for i in range(n_audio)
62
+ ]
63
+
64
+ if single:
65
+ language_tokens = language_tokens[0]
66
+ language_probs = language_probs[0]
67
+
68
+ return language_tokens, language_probs
69
+
70
+
71
+ @dataclass(frozen=True)
72
+ class DecodingOptions:
73
+ task: str = "transcribe" # whether to perform X->X "transcribe" or X->English "translate"
74
+ language: Optional[str] = None # language that the audio is in; uses detected language if None
75
+
76
+ # sampling-related options
77
+ temperature: float = 0.0
78
+ sample_len: Optional[int] = None # maximum number of tokens to sample
79
+ best_of: Optional[int] = None # number of independent samples to collect, when t > 0
80
+ beam_size: Optional[int] = None # number of beams in beam search, when t == 0
81
+ patience: Optional[float] = None # patience in beam search (https://arxiv.org/abs/2204.05424)
82
+
83
+ # options for ranking generations (either beams or best-of-N samples)
84
+ length_penalty: Optional[float] = None # "alpha" in Google NMT, None defaults to length norm
85
+
86
+ # prompt, prefix, and token suppression
87
+ prompt: Optional[Union[str, List[int]]] = None # text or tokens for the previous context
88
+ prefix: Optional[Union[str, List[int]]] = None # text or tokens to prefix the current context
89
+ suppress_blank: bool = True # this will suppress blank outputs
90
+
91
+ # list of tokens ids (or comma-separated token ids) to suppress
92
+ # "-1" will suppress a set of symbols as defined in `tokenizer.non_speech_tokens()`
93
+ suppress_tokens: Optional[Union[str, Iterable[int]]] = "-1"
94
+
95
+ # timestamp sampling options
96
+ without_timestamps: bool = False # use <|notimestamps|> to sample text tokens only
97
+ max_initial_timestamp: Optional[float] = 1.0 # the initial timestamp cannot be later than this
98
+
99
+ # implementation details
100
+ fp16: bool = True # use fp16 for most of the calculation
101
+
102
+
103
+ @dataclass(frozen=True)
104
+ class DecodingResult:
105
+ audio_features: Tensor
106
+ language: str
107
+ language_probs: Optional[Dict[str, float]] = None
108
+ tokens: List[int] = field(default_factory=list)
109
+ text: str = ""
110
+ avg_logprob: float = np.nan
111
+ no_speech_prob: float = np.nan
112
+ temperature: float = np.nan
113
+ compression_ratio: float = np.nan
114
+
115
+
116
+ class Inference:
117
+ def logits(self, tokens: Tensor, audio_features: Tensor) -> Tensor:
118
+ """Perform a forward pass on the decoder and return per-token logits"""
119
+ raise NotImplementedError
120
+
121
+ def rearrange_kv_cache(self, source_indices) -> None:
122
+ """Update the key-value cache according to the updated beams"""
123
+ raise NotImplementedError
124
+
125
+ def cleanup_caching(self) -> None:
126
+ """Clean up any resources or hooks after decoding is finished"""
127
+ pass
128
+
129
+
130
+ class PyTorchInference(Inference):
131
+ def __init__(self, model: "Whisper", initial_token_length: int):
132
+ self.model: "Whisper" = model
133
+ self.initial_token_length = initial_token_length
134
+ self.kv_cache = {}
135
+ self.hooks = []
136
+
137
+ def logits(self, tokens: Tensor, audio_features: Tensor) -> Tensor:
138
+ if not self.kv_cache:
139
+ self.kv_cache, self.hooks = self.model.install_kv_cache_hooks()
140
+
141
+ if tokens.shape[-1] > self.initial_token_length:
142
+ # only need to use the last token except in the first forward pass
143
+ tokens = tokens[:, -1:]
144
+
145
+ return self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache)
146
+
147
+ def cleanup_caching(self):
148
+ for hook in self.hooks:
149
+ hook.remove()
150
+
151
+ self.kv_cache = {}
152
+ self.hooks = []
153
+
154
+ def rearrange_kv_cache(self, source_indices):
155
+ for module, tensor in self.kv_cache.items():
156
+ # update the key/value cache to contain the selected sequences
157
+ self.kv_cache[module] = tensor[source_indices].detach()
158
+
159
+
160
+ class SequenceRanker:
161
+ def rank(self, tokens: List[List[Tensor]], sum_logprobs: List[List[float]]) -> List[int]:
162
+ """
163
+ Given a list of groups of samples and their cumulative log probabilities,
164
+ return the indices of the samples in each group to select as the final result
165
+ """
166
+ raise NotImplementedError
167
+
168
+
169
+ class MaximumLikelihoodRanker(SequenceRanker):
170
+ """
171
+ Select the sample with the highest log probabilities, penalized using either
172
+ a simple length normalization or Google NMT paper's length penalty
173
+ """
174
+
175
+ def __init__(self, length_penalty: Optional[float]):
176
+ self.length_penalty = length_penalty
177
+
178
+ def rank(self, tokens: List[List[Tensor]], sum_logprobs: List[List[float]]):
179
+ def scores(logprobs, lengths):
180
+ result = []
181
+ for logprob, length in zip(logprobs, lengths):
182
+ if self.length_penalty is None:
183
+ penalty = length
184
+ else:
185
+ # from the Google NMT paper
186
+ penalty = ((5 + length) / 6) ** self.length_penalty
187
+ result.append(logprob / penalty)
188
+ return result
189
+
190
+ # get the sequence with the highest score
191
+ lengths = [[len(t) for t in s] for s in tokens]
192
+ return [np.argmax(scores(p, l)) for p, l in zip(sum_logprobs, lengths)]
193
+
194
+
195
+ class TokenDecoder:
196
+ def reset(self):
197
+ """Initialize any stateful variables for decoding a new sequence"""
198
+
199
+ def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]:
200
+ """Specify how to select the next token, based on the current trace and logits
201
+
202
+ Parameters
203
+ ----------
204
+ tokens : Tensor, shape = (n_batch, current_sequence_length)
205
+ all tokens in the context so far, including the prefix and sot_sequence tokens
206
+
207
+ logits : Tensor, shape = (n_batch, vocab_size)
208
+ per-token logits of the probability distribution at the current step
209
+
210
+ sum_logprobs : Tensor, shape = (n_batch)
211
+ cumulative log probabilities for each sequence
212
+
213
+ Returns
214
+ -------
215
+ tokens : Tensor, shape = (n_batch, current_sequence_length + 1)
216
+ the tokens, appended with the selected next token
217
+
218
+ completed : bool
219
+ True if all sequences has reached the end of text
220
+
221
+ """
222
+ raise NotImplementedError
223
+
224
+ def finalize(
225
+ self, tokens: Tensor, sum_logprobs: Tensor
226
+ ) -> Tuple[Sequence[Sequence[Tensor]], List[List[float]]]:
227
+ """Finalize search and return the final candidate sequences
228
+
229
+ Parameters
230
+ ----------
231
+ tokens : Tensor, shape = (n_audio, n_group, current_sequence_length)
232
+ all tokens in the context so far, including the prefix and sot_sequence
233
+
234
+ sum_logprobs : Tensor, shape = (n_audio, n_group)
235
+ cumulative log probabilities for each sequence
236
+
237
+ Returns
238
+ -------
239
+ tokens : Sequence[Sequence[Tensor]], length = n_audio
240
+ sequence of Tensors containing candidate token sequences, for each audio input
241
+
242
+ sum_logprobs : List[List[float]], length = n_audio
243
+ sequence of cumulative log probabilities corresponding to the above
244
+
245
+ """
246
+ raise NotImplementedError
247
+
248
+
249
+ class GreedyDecoder(TokenDecoder):
250
+ def __init__(self, temperature: float, eot: int):
251
+ self.temperature = temperature
252
+ self.eot = eot
253
+
254
+ def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]:
255
+ temperature = self.temperature
256
+ if temperature == 0:
257
+ next_tokens = logits.argmax(dim=-1)
258
+ else:
259
+ next_tokens = Categorical(logits=logits / temperature).sample()
260
+
261
+ logprobs = F.log_softmax(logits.float(), dim=-1)
262
+ current_logprobs = logprobs[torch.arange(logprobs.shape[0]), next_tokens]
263
+ sum_logprobs += current_logprobs * (tokens[:, -1] != self.eot)
264
+
265
+ next_tokens[tokens[:, -1] == self.eot] = self.eot
266
+ tokens = torch.cat([tokens, next_tokens[:, None]], dim=-1)
267
+
268
+ completed = (tokens[:, -1] == self.eot).all()
269
+ return tokens, completed
270
+
271
+ def finalize(self, tokens: Tensor, sum_logprobs: Tensor):
272
+ # make sure each sequence has at least one EOT token at the end
273
+ tokens = F.pad(tokens, (0, 1), value=self.eot)
274
+ return tokens, sum_logprobs.tolist()
275
+
276
+
277
+ class BeamSearchDecoder(TokenDecoder):
278
+ def __init__(self, beam_size: int, eot: int, inference: Inference, patience: Optional[float] = None):
279
+ self.beam_size = beam_size
280
+ self.eot = eot
281
+ self.inference = inference
282
+ self.patience = patience or 1.0
283
+ self.max_candidates: int = round(beam_size * self.patience)
284
+ self.finished_sequences = None
285
+
286
+ assert self.max_candidates > 0, f"Invalid beam size ({beam_size}) or patience ({patience})"
287
+
288
+ def reset(self):
289
+ self.finished_sequences = None
290
+
291
+ def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]:
292
+ if tokens.shape[0] % self.beam_size != 0:
293
+ raise ValueError(f"{tokens.shape}[0] % {self.beam_size} != 0")
294
+
295
+ n_audio = tokens.shape[0] // self.beam_size
296
+ if self.finished_sequences is None: # for the first update
297
+ self.finished_sequences = [{} for _ in range(n_audio)]
298
+
299
+ logprobs = F.log_softmax(logits.float(), dim=-1)
300
+ next_tokens, source_indices, finished_sequences = [], [], []
301
+ for i in range(n_audio):
302
+ scores, sources, finished = {}, {}, {}
303
+
304
+ # STEP 1: calculate the cumulative log probabilities for possible candidates
305
+ for j in range(self.beam_size):
306
+ idx = i * self.beam_size + j
307
+ prefix = tokens[idx].tolist()
308
+ for logprob, token in zip(*logprobs[idx].topk(self.beam_size + 1)):
309
+ new_logprob = (sum_logprobs[idx] + logprob).item()
310
+ sequence = tuple(prefix + [token.item()])
311
+ scores[sequence] = new_logprob
312
+ sources[sequence] = idx
313
+
314
+ # STEP 2: rank the candidates and keep the top beam_size sequences for each audio
315
+ saved = 0
316
+ for sequence in sorted(scores, key=scores.get, reverse=True):
317
+ if sequence[-1] == self.eot:
318
+ finished[sequence] = scores[sequence]
319
+ else:
320
+ sum_logprobs[len(next_tokens)] = scores[sequence]
321
+ next_tokens.append(sequence)
322
+ source_indices.append(sources[sequence])
323
+
324
+ saved += 1
325
+ if saved == self.beam_size:
326
+ break
327
+
328
+ finished_sequences.append(finished)
329
+
330
+ tokens = torch.tensor(next_tokens, device=tokens.device)
331
+ self.inference.rearrange_kv_cache(source_indices)
332
+
333
+ # add newly finished sequences to self.finished_sequences
334
+ assert len(self.finished_sequences) == len(finished_sequences)
335
+ for previously_finished, newly_finished in zip(self.finished_sequences, finished_sequences):
336
+ for seq in sorted(newly_finished, key=newly_finished.get, reverse=True):
337
+ if len(previously_finished) >= self.max_candidates:
338
+ break # the candidate list is full
339
+ previously_finished[seq] = newly_finished[seq]
340
+
341
+ # mark as completed if all audio has enough number of samples
342
+ completed = all(
343
+ len(sequences) >= self.max_candidates for sequences in self.finished_sequences
344
+ )
345
+ return tokens, completed
346
+
347
+ def finalize(self, preceding_tokens: Tensor, sum_logprobs: Tensor):
348
+ # collect all finished sequences, including patience, and add unfinished ones if not enough
349
+ sum_logprobs = sum_logprobs.cpu()
350
+ for i, sequences in enumerate(self.finished_sequences):
351
+ if len(sequences) < self.beam_size: # when not enough sequences are finished
352
+ for j in list(np.argsort(sum_logprobs[i]))[::-1]:
353
+ sequence = preceding_tokens[i, j].tolist() + [self.eot]
354
+ sequences[tuple(sequence)] = sum_logprobs[i][j].item()
355
+ if len(sequences) >= self.beam_size:
356
+ break
357
+
358
+ tokens: List[List[Tensor]] = [
359
+ [torch.tensor(seq) for seq in sequences.keys()] for sequences in self.finished_sequences
360
+ ]
361
+ sum_logprobs: List[List[float]] = [
362
+ list(sequences.values()) for sequences in self.finished_sequences
363
+ ]
364
+ return tokens, sum_logprobs
365
+
366
+
367
+ class LogitFilter:
368
+ def apply(self, logits: Tensor, tokens: Tensor) -> None:
369
+ """Apply any filtering or masking to logits in-place
370
+
371
+ Parameters
372
+ ----------
373
+ logits : Tensor, shape = (n_batch, vocab_size)
374
+ per-token logits of the probability distribution at the current step
375
+
376
+ tokens : Tensor, shape = (n_batch, current_sequence_length)
377
+ all tokens in the context so far, including the prefix and sot_sequence tokens
378
+
379
+ """
380
+ raise NotImplementedError
381
+
382
+
383
+ class SuppressBlank(LogitFilter):
384
+ def __init__(self, tokenizer: Tokenizer, sample_begin: int):
385
+ self.tokenizer = tokenizer
386
+ self.sample_begin = sample_begin
387
+
388
+ def apply(self, logits: Tensor, tokens: Tensor):
389
+ if tokens.shape[1] == self.sample_begin:
390
+ logits[:, self.tokenizer.encode(" ") + [self.tokenizer.eot]] = -np.inf
391
+
392
+
393
+ class SuppressTokens(LogitFilter):
394
+ def __init__(self, suppress_tokens: Sequence[int]):
395
+ self.suppress_tokens = list(suppress_tokens)
396
+
397
+ def apply(self, logits: Tensor, tokens: Tensor):
398
+ logits[:, self.suppress_tokens] = -np.inf
399
+
400
+
401
+ class ApplyTimestampRules(LogitFilter):
402
+ def __init__(
403
+ self, tokenizer: Tokenizer, sample_begin: int, max_initial_timestamp_index: Optional[int]
404
+ ):
405
+ self.tokenizer = tokenizer
406
+ self.sample_begin = sample_begin
407
+ self.max_initial_timestamp_index = max_initial_timestamp_index
408
+
409
+ def apply(self, logits: Tensor, tokens: Tensor):
410
+ # suppress <|notimestamps|> which is handled by without_timestamps
411
+ if self.tokenizer.no_timestamps is not None:
412
+ logits[:, self.tokenizer.no_timestamps] = -np.inf
413
+
414
+ # timestamps have to appear in pairs, except directly before EOT; mask logits accordingly
415
+ for k in range(tokens.shape[0]):
416
+ seq = [t for t in tokens[k, self.sample_begin :].tolist()]
417
+ last_was_timestamp = len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin
418
+ penultimate_was_timestamp = len(seq) < 2 or seq[-2] >= self.tokenizer.timestamp_begin
419
+
420
+ if last_was_timestamp:
421
+ if penultimate_was_timestamp: # has to be non-timestamp
422
+ logits[k, self.tokenizer.timestamp_begin :] = -np.inf
423
+ else: # cannot be normal text tokens
424
+ logits[k, : self.tokenizer.eot] = -np.inf
425
+
426
+ if tokens.shape[1] == self.sample_begin:
427
+ # suppress generating non-timestamp tokens at the beginning
428
+ logits[:, : self.tokenizer.timestamp_begin] = -np.inf
429
+
430
+ # apply the `max_initial_timestamp` option
431
+ if self.max_initial_timestamp_index is not None:
432
+ last_allowed = self.tokenizer.timestamp_begin + self.max_initial_timestamp_index
433
+ logits[:, last_allowed + 1 :] = -np.inf
434
+
435
+ # if sum of probability over timestamps is above any other token, sample timestamp
436
+ logprobs = F.log_softmax(logits.float(), dim=-1)
437
+ for k in range(tokens.shape[0]):
438
+ timestamp_logprob = logprobs[k, self.tokenizer.timestamp_begin :].logsumexp(dim=-1)
439
+ max_text_token_logprob = logprobs[k, : self.tokenizer.timestamp_begin].max()
440
+ if timestamp_logprob > max_text_token_logprob:
441
+ logits[k, : self.tokenizer.timestamp_begin] = -np.inf
442
+
443
+
444
+ class DecodingTask:
445
+ inference: Inference
446
+ sequence_ranker: SequenceRanker
447
+ decoder: TokenDecoder
448
+ logit_filters: List[LogitFilter]
449
+
450
+ def __init__(self, model: "Whisper", options: DecodingOptions):
451
+ self.model = model
452
+
453
+ language = options.language or "en"
454
+ tokenizer = get_tokenizer(model.is_multilingual, language=language, task=options.task)
455
+ self.tokenizer: Tokenizer = tokenizer
456
+ self.options: DecodingOptions = self._verify_options(options)
457
+
458
+ self.n_group: int = options.beam_size or options.best_of or 1
459
+ self.n_ctx: int = model.dims.n_text_ctx
460
+ self.sample_len: int = options.sample_len or model.dims.n_text_ctx // 2
461
+
462
+ self.sot_sequence: Tuple[int] = tokenizer.sot_sequence
463
+ if self.options.without_timestamps:
464
+ self.sot_sequence = tokenizer.sot_sequence_including_notimestamps
465
+
466
+ self.initial_tokens: Tuple[int] = self._get_initial_tokens()
467
+ self.sample_begin: int = len(self.initial_tokens)
468
+ self.sot_index: int = self.initial_tokens.index(tokenizer.sot)
469
+
470
+ # inference: implements the forward pass through the decoder, including kv caching
471
+ self.inference = PyTorchInference(model, len(self.initial_tokens))
472
+
473
+ # sequence ranker: implements how to rank a group of sampled sequences
474
+ self.sequence_ranker = MaximumLikelihoodRanker(options.length_penalty)
475
+
476
+ # decoder: implements how to select the next tokens, given the autoregressive distribution
477
+ if options.beam_size is not None:
478
+ self.decoder = BeamSearchDecoder(
479
+ options.beam_size, tokenizer.eot, self.inference, options.patience
480
+ )
481
+ else:
482
+ self.decoder = GreedyDecoder(options.temperature, tokenizer.eot)
483
+
484
+ # logit filters: applies various rules to suppress or penalize certain tokens
485
+ self.logit_filters = []
486
+ if self.options.suppress_blank:
487
+ self.logit_filters.append(SuppressBlank(self.tokenizer, self.sample_begin))
488
+ if self.options.suppress_tokens:
489
+ self.logit_filters.append(SuppressTokens(self._get_suppress_tokens()))
490
+ if not options.without_timestamps:
491
+ precision = CHUNK_LENGTH / model.dims.n_audio_ctx # usually 0.02 seconds
492
+ max_initial_timestamp_index = None
493
+ if options.max_initial_timestamp:
494
+ max_initial_timestamp_index = round(self.options.max_initial_timestamp / precision)
495
+ self.logit_filters.append(
496
+ ApplyTimestampRules(tokenizer, self.sample_begin, max_initial_timestamp_index)
497
+ )
498
+
499
+ def _verify_options(self, options: DecodingOptions) -> DecodingOptions:
500
+ if options.beam_size is not None and options.best_of is not None:
501
+ raise ValueError("beam_size and best_of can't be given together")
502
+ if options.temperature == 0:
503
+ if options.best_of is not None:
504
+ raise ValueError("best_of with greedy sampling (T=0) is not compatible")
505
+ if options.patience is not None and options.beam_size is None:
506
+ raise ValueError("patience requires beam_size to be given")
507
+ if options.length_penalty is not None and not (0 <= options.length_penalty <= 1):
508
+ raise ValueError("length_penalty (alpha) should be a value between 0 and 1")
509
+
510
+ return options
511
+
512
+ def _get_initial_tokens(self) -> Tuple[int]:
513
+ tokens = list(self.sot_sequence)
514
+ prefix = self.options.prefix
515
+ prompt = self.options.prompt
516
+
517
+ if prefix:
518
+ prefix_tokens = (
519
+ self.tokenizer.encode(" " + prefix.strip()) if isinstance(prefix, str) else prefix
520
+ )
521
+ if self.sample_len is not None:
522
+ max_prefix_len = self.n_ctx // 2 - self.sample_len
523
+ prefix_tokens = prefix_tokens[-max_prefix_len:]
524
+ tokens = tokens + prefix_tokens
525
+
526
+ if prompt:
527
+ prompt_tokens = (
528
+ self.tokenizer.encode(" " + prompt.strip()) if isinstance(prompt, str) else prompt
529
+ )
530
+ tokens = [self.tokenizer.sot_prev] + prompt_tokens[-(self.n_ctx // 2 - 1) :] + tokens
531
+
532
+ return tuple(tokens)
533
+
534
+ def _get_suppress_tokens(self) -> Tuple[int]:
535
+ suppress_tokens = self.options.suppress_tokens
536
+
537
+ if isinstance(suppress_tokens, str):
538
+ suppress_tokens = [int(t) for t in suppress_tokens.split(",")]
539
+
540
+ if -1 in suppress_tokens:
541
+ suppress_tokens = [t for t in suppress_tokens if t >= 0]
542
+ suppress_tokens.extend(self.tokenizer.non_speech_tokens)
543
+ elif suppress_tokens is None or len(suppress_tokens) == 0:
544
+ suppress_tokens = [] # interpret empty string as an empty list
545
+ else:
546
+ assert isinstance(suppress_tokens, list), "suppress_tokens must be a list"
547
+
548
+ suppress_tokens.extend(
549
+ [self.tokenizer.sot, self.tokenizer.sot_prev, self.tokenizer.sot_lm]
550
+ )
551
+ if self.tokenizer.no_speech is not None:
552
+ # no-speech probability is collected separately
553
+ suppress_tokens.append(self.tokenizer.no_speech)
554
+
555
+ return tuple(sorted(set(suppress_tokens)))
556
+
557
+ def _get_audio_features(self, mel: Tensor):
558
+ if self.options.fp16:
559
+ mel = mel.half()
560
+
561
+ if mel.shape[-2:] == (self.model.dims.n_audio_ctx, self.model.dims.n_audio_state):
562
+ # encoded audio features are given; skip audio encoding
563
+ print("encoded audio features are given; skip audio encoding")
564
+ audio_features = mel
565
+ else:
566
+ print(mel.shape)
567
+ print("===============================")
568
+ audio_features = self.model.encoder(mel)
569
+
570
+ if audio_features.dtype != (torch.float16 if self.options.fp16 else torch.float32):
571
+ return TypeError(f"audio_features has an incorrect dtype: {audio_features.dtype}")
572
+
573
+ return audio_features
574
+
575
+ def _detect_language(self, audio_features: Tensor, tokens: Tensor):
576
+ languages = [self.options.language] * audio_features.shape[0]
577
+ lang_probs = None
578
+
579
+ if self.options.language is None or self.options.task == "lang_id":
580
+ lang_tokens, lang_probs = self.model.detect_language(audio_features, self.tokenizer)
581
+ languages = [max(probs, key=probs.get) for probs in lang_probs]
582
+ if self.options.language is None:
583
+ tokens[:, self.sot_index + 1] = lang_tokens # write language tokens
584
+
585
+ return languages, lang_probs
586
+
587
+ def _main_loop(self, audio_features: Tensor, tokens: Tensor):
588
+ assert audio_features.shape[0] == tokens.shape[0]
589
+ n_batch = tokens.shape[0]
590
+ sum_logprobs: Tensor = torch.zeros(n_batch, device=audio_features.device)
591
+ no_speech_probs = [np.nan] * n_batch
592
+
593
+ try:
594
+ for i in range(self.sample_len):
595
+ logits = self.inference.logits(tokens, audio_features)
596
+
597
+ if i == 0 and self.tokenizer.no_speech is not None: # save no_speech_probs
598
+ probs_at_sot = logits[:, self.sot_index].float().softmax(dim=-1)
599
+ no_speech_probs = probs_at_sot[:, self.tokenizer.no_speech].tolist()
600
+
601
+ # now we need to consider the logits at the last token only
602
+ logits = logits[:, -1]
603
+
604
+ # apply the logit filters, e.g. for suppressing or applying penalty to
605
+ for logit_filter in self.logit_filters:
606
+ logit_filter.apply(logits, tokens)
607
+
608
+ # expand the tokens tensor with the selected next tokens
609
+ tokens, completed = self.decoder.update(tokens, logits, sum_logprobs)
610
+
611
+ if completed or tokens.shape[-1] > self.n_ctx:
612
+ break
613
+ finally:
614
+ self.inference.cleanup_caching()
615
+
616
+ return tokens, sum_logprobs, no_speech_probs
617
+
618
+ @torch.no_grad()
619
+ def run(self, mel: Tensor) -> List[DecodingResult]:
620
+ self.decoder.reset()
621
+ tokenizer: Tokenizer = self.tokenizer
622
+ n_audio: int = mel.shape[0]
623
+
624
+ audio_features: Tensor = self._get_audio_features(mel) # encoder forward pass
625
+ tokens: Tensor = torch.tensor([self.initial_tokens]).repeat(n_audio, 1)
626
+
627
+ # detect language if requested, overwriting the language token
628
+ languages, language_probs = self._detect_language(audio_features, tokens)
629
+ if self.options.task == "lang_id":
630
+ return [
631
+ DecodingResult(audio_features=features, language=language, language_probs=probs)
632
+ for features, language, probs in zip(audio_features, languages, language_probs)
633
+ ]
634
+
635
+ # repeat the audio & text tensors by the group size, for beam search or best-of-n sampling
636
+ audio_features = audio_features.repeat_interleave(self.n_group, dim=0)
637
+ tokens = tokens.repeat_interleave(self.n_group, dim=0).to(audio_features.device)
638
+
639
+ # call the main sampling loop
640
+ tokens, sum_logprobs, no_speech_probs = self._main_loop(audio_features, tokens)
641
+
642
+ # reshape the tensors to have (n_audio, n_group) as the first two dimensions
643
+ audio_features = audio_features[:: self.n_group]
644
+ no_speech_probs = no_speech_probs[:: self.n_group]
645
+ assert audio_features.shape[0] == len(no_speech_probs) == n_audio
646
+
647
+ tokens = tokens.reshape(n_audio, self.n_group, -1)
648
+ sum_logprobs = sum_logprobs.reshape(n_audio, self.n_group)
649
+
650
+ # get the final candidates for each group, and slice between the first sampled token and EOT
651
+ tokens, sum_logprobs = self.decoder.finalize(tokens, sum_logprobs)
652
+ tokens: List[List[Tensor]] = [
653
+ [t[self.sample_begin : (t == tokenizer.eot).nonzero()[0, 0]] for t in s] for s in tokens
654
+ ]
655
+
656
+ # select the top-ranked sample in each group
657
+ selected = self.sequence_ranker.rank(tokens, sum_logprobs)
658
+ tokens: List[List[int]] = [t[i].tolist() for i, t in zip(selected, tokens)]
659
+ texts: List[str] = [tokenizer.decode(t).strip() for t in tokens]
660
+
661
+ sum_logprobs: List[float] = [lp[i] for i, lp in zip(selected, sum_logprobs)]
662
+ avg_logprobs: List[float] = [lp / (len(t) + 1) for t, lp in zip(tokens, sum_logprobs)]
663
+
664
+ fields = (texts, languages, tokens, audio_features, avg_logprobs, no_speech_probs)
665
+ if len(set(map(len, fields))) != 1:
666
+ raise RuntimeError(f"inconsistent result lengths: {list(map(len, fields))}")
667
+
668
+ return [
669
+ DecodingResult(
670
+ audio_features=features,
671
+ language=language,
672
+ tokens=tokens,
673
+ text=text,
674
+ avg_logprob=avg_logprob,
675
+ no_speech_prob=no_speech_prob,
676
+ temperature=self.options.temperature,
677
+ compression_ratio=compression_ratio(text),
678
+ )
679
+ for text, language, tokens, features, avg_logprob, no_speech_prob in zip(*fields)
680
+ ]
681
+
682
+
683
+ @torch.no_grad()
684
+ def decode(model: "Whisper", mel: Tensor, options: DecodingOptions = DecodingOptions()) -> Union[DecodingResult, List[DecodingResult]]:
685
+ """
686
+ Performs decoding of 30-second audio segment(s), provided as Mel spectrogram(s).
687
+
688
+ Parameters
689
+ ----------
690
+ model: Whisper
691
+ the Whisper model instance
692
+
693
+ mel: torch.Tensor, shape = (80, 3000) or (*, 80, 3000)
694
+ A tensor containing the Mel spectrogram(s)
695
+
696
+ options: DecodingOptions
697
+ A dataclass that contains all necessary options for decoding 30-second segments
698
+
699
+ Returns
700
+ -------
701
+ result: Union[DecodingResult, List[DecodingResult]]
702
+ The result(s) of decoding contained in `DecodingResult` dataclass instance(s)
703
+ """
704
+ single = mel.ndim == 2
705
+ if single:
706
+ mel = mel.unsqueeze(0)
707
+ result = DecodingTask(model, options).run(mel)
708
+
709
+ if single:
710
+ result = result[0]
711
+
712
+ return result
vencoder/whisper/model.py ADDED
@@ -0,0 +1,269 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from typing import Dict
3
+ from typing import Iterable, Optional
4
+
5
+ import numpy as np
6
+ import torch
7
+ import torch.nn.functional as F
8
+ from torch import Tensor
9
+ from torch import nn
10
+
11
+ from .decoding import detect_language as detect_language_function, decode as decode_function
12
+
13
+
14
+ @dataclass
15
+ class ModelDimensions:
16
+ n_mels: int
17
+ n_audio_ctx: int
18
+ n_audio_state: int
19
+ n_audio_head: int
20
+ n_audio_layer: int
21
+ n_vocab: int
22
+ n_text_ctx: int
23
+ n_text_state: int
24
+ n_text_head: int
25
+ n_text_layer: int
26
+
27
+
28
+ class LayerNorm(nn.LayerNorm):
29
+ def forward(self, x: Tensor) -> Tensor:
30
+ return super().forward(x.float()).type(x.dtype)
31
+
32
+
33
+ class Linear(nn.Linear):
34
+ def forward(self, x: Tensor) -> Tensor:
35
+ return F.linear(
36
+ x, self.weight.to(x.dtype), None if self.bias is None else self.bias.to(x.dtype)
37
+ )
38
+
39
+
40
+ class Conv1d(nn.Conv1d):
41
+ def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor:
42
+ return super()._conv_forward(
43
+ x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype)
44
+ )
45
+
46
+
47
+ def sinusoids(length, channels, max_timescale=10000):
48
+ """Returns sinusoids for positional embedding"""
49
+ assert channels % 2 == 0
50
+ log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
51
+ inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
52
+ scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
53
+ return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
54
+
55
+
56
+ class MultiHeadAttention(nn.Module):
57
+ def __init__(self, n_state: int, n_head: int):
58
+ super().__init__()
59
+ self.n_head = n_head
60
+ self.query = Linear(n_state, n_state)
61
+ self.key = Linear(n_state, n_state, bias=False)
62
+ self.value = Linear(n_state, n_state)
63
+ self.out = Linear(n_state, n_state)
64
+
65
+ def forward(
66
+ self,
67
+ x: Tensor,
68
+ xa: Optional[Tensor] = None,
69
+ mask: Optional[Tensor] = None,
70
+ kv_cache: Optional[dict] = None,
71
+ ):
72
+ q = self.query(x)
73
+
74
+ if kv_cache is None or xa is None or self.key not in kv_cache:
75
+ # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
76
+ # otherwise, perform key/value projections for self- or cross-attention as usual.
77
+ k = self.key(x if xa is None else xa)
78
+ v = self.value(x if xa is None else xa)
79
+ else:
80
+ # for cross-attention, calculate keys and values once and reuse in subsequent calls.
81
+ k = kv_cache[self.key]
82
+ v = kv_cache[self.value]
83
+
84
+ wv, qk = self.qkv_attention(q, k, v, mask)
85
+ return self.out(wv), qk
86
+
87
+ def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
88
+ n_batch, n_ctx, n_state = q.shape
89
+ scale = (n_state // self.n_head) ** -0.25
90
+ q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
91
+ k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
92
+ v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
93
+
94
+ qk = q @ k
95
+ if mask is not None:
96
+ qk = qk + mask[:n_ctx, :n_ctx]
97
+ qk = qk.float()
98
+
99
+ w = F.softmax(qk, dim=-1).to(q.dtype)
100
+ return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
101
+
102
+
103
+ class ResidualAttentionBlock(nn.Module):
104
+ def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):
105
+ super().__init__()
106
+
107
+ self.attn = MultiHeadAttention(n_state, n_head)
108
+ self.attn_ln = LayerNorm(n_state)
109
+
110
+ self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
111
+ self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
112
+
113
+ n_mlp = n_state * 4
114
+ self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state))
115
+ self.mlp_ln = LayerNorm(n_state)
116
+
117
+ def forward(
118
+ self,
119
+ x: Tensor,
120
+ xa: Optional[Tensor] = None,
121
+ mask: Optional[Tensor] = None,
122
+ kv_cache: Optional[dict] = None,
123
+ ):
124
+ x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)[0]
125
+ if self.cross_attn:
126
+ x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
127
+ x = x + self.mlp(self.mlp_ln(x))
128
+ return x
129
+
130
+
131
+ class AudioEncoder(nn.Module):
132
+ def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
133
+ super().__init__()
134
+ self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
135
+ self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
136
+ self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state))
137
+
138
+ self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
139
+ [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
140
+ )
141
+ self.ln_post = LayerNorm(n_state)
142
+
143
+ def forward(self, x: Tensor):
144
+ """
145
+ x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)
146
+ the mel spectrogram of the audio
147
+ """
148
+ x = F.gelu(self.conv1(x))
149
+ x = F.gelu(self.conv2(x))
150
+ x = x.permute(0, 2, 1)
151
+
152
+ len_x = x.shape[1]
153
+ len_e = self.positional_embedding.shape[0]
154
+ assert len_x <= len_e, "incorrect audio shape"
155
+ pos_e = self.positional_embedding[:len_x, :]
156
+ x = (x + pos_e).to(x.dtype)
157
+
158
+ for block in self.blocks:
159
+ x = block(x)
160
+
161
+ x = self.ln_post(x)
162
+ return x
163
+
164
+
165
+ class TextDecoder(nn.Module):
166
+ def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
167
+ super().__init__()
168
+
169
+ self.token_embedding = nn.Embedding(n_vocab, n_state)
170
+ self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state))
171
+
172
+ self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
173
+ [ResidualAttentionBlock(n_state, n_head, cross_attention=True) for _ in range(n_layer)]
174
+ )
175
+ self.ln = LayerNorm(n_state)
176
+
177
+ mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1)
178
+ self.register_buffer("mask", mask, persistent=False)
179
+
180
+ def forward(self, x: Tensor, xa: Tensor, kv_cache: Optional[dict] = None):
181
+ """
182
+ x : torch.LongTensor, shape = (batch_size, <= n_ctx)
183
+ the text tokens
184
+ xa : torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)
185
+ the encoded audio features to be attended on
186
+ """
187
+ offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0
188
+ x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]]
189
+ x = x.to(xa.dtype)
190
+
191
+ for block in self.blocks:
192
+ x = block(x, xa, mask=self.mask, kv_cache=kv_cache)
193
+
194
+ x = self.ln(x)
195
+ logits = (x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)).float()
196
+
197
+ return logits
198
+
199
+
200
+ class Whisper(nn.Module):
201
+ def __init__(self, dims: ModelDimensions):
202
+ super().__init__()
203
+ self.dims = dims
204
+ self.encoder = AudioEncoder(
205
+ self.dims.n_mels,
206
+ self.dims.n_audio_ctx,
207
+ self.dims.n_audio_state,
208
+ self.dims.n_audio_head,
209
+ self.dims.n_audio_layer,
210
+ )
211
+ self.decoder = TextDecoder(
212
+ self.dims.n_vocab,
213
+ self.dims.n_text_ctx,
214
+ self.dims.n_text_state,
215
+ self.dims.n_text_head,
216
+ self.dims.n_text_layer,
217
+ )
218
+
219
+ def embed_audio(self, mel: torch.Tensor):
220
+ return self.encoder(mel)
221
+
222
+ def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor):
223
+ return self.decoder(tokens, audio_features)
224
+
225
+ def forward(self, mel: torch.Tensor, tokens: torch.Tensor) -> Dict[str, torch.Tensor]:
226
+ return self.decoder(tokens, self.encoder(mel))
227
+
228
+ @property
229
+ def device(self):
230
+ return next(self.parameters()).device
231
+
232
+ @property
233
+ def is_multilingual(self):
234
+ return self.dims.n_vocab == 51865
235
+
236
+ def install_kv_cache_hooks(self, cache: Optional[dict] = None):
237
+ """
238
+ The `MultiHeadAttention` module optionally accepts `kv_cache` which stores the key and value
239
+ tensors calculated for the previous positions. This method returns a dictionary that stores
240
+ all caches, and the necessary hooks for the key and value projection modules that save the
241
+ intermediate tensors to be reused during later calculations.
242
+
243
+ Returns
244
+ -------
245
+ cache : Dict[nn.Module, torch.Tensor]
246
+ A dictionary object mapping the key/value projection modules to its cache
247
+ hooks : List[RemovableHandle]
248
+ List of PyTorch RemovableHandle objects to stop the hooks to be called
249
+ """
250
+ cache = {**cache} if cache is not None else {}
251
+ hooks = []
252
+
253
+ def save_to_cache(module, _, output):
254
+ if module not in cache or output.shape[1] > self.decoder.positional_embedding.shape[0]:
255
+ cache[module] = output # save as-is, for the first token or cross attention
256
+ else:
257
+ cache[module] = torch.cat([cache[module], output], dim=1).detach()
258
+ return cache[module]
259
+
260
+ def install_hooks(layer: nn.Module):
261
+ if isinstance(layer, MultiHeadAttention):
262
+ hooks.append(layer.key.register_forward_hook(save_to_cache))
263
+ hooks.append(layer.value.register_forward_hook(save_to_cache))
264
+
265
+ self.decoder.apply(install_hooks)
266
+ return cache, hooks
267
+
268
+ detect_language = detect_language_function
269
+ decode = decode_function
vencoder/whisper/tokenizer.py ADDED
@@ -0,0 +1,331 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from dataclasses import dataclass
3
+ from functools import lru_cache
4
+ from typing import List, Optional, Tuple, Union
5
+
6
+ import numpy as np
7
+ import torch
8
+ from transformers import GPT2TokenizerFast
9
+
10
+ LANGUAGES = {
11
+ "en": "english",
12
+ "zh": "chinese",
13
+ "de": "german",
14
+ "es": "spanish",
15
+ "ru": "russian",
16
+ "ko": "korean",
17
+ "fr": "french",
18
+ "ja": "japanese",
19
+ "pt": "portuguese",
20
+ "tr": "turkish",
21
+ "pl": "polish",
22
+ "ca": "catalan",
23
+ "nl": "dutch",
24
+ "ar": "arabic",
25
+ "sv": "swedish",
26
+ "it": "italian",
27
+ "id": "indonesian",
28
+ "hi": "hindi",
29
+ "fi": "finnish",
30
+ "vi": "vietnamese",
31
+ "he": "hebrew",
32
+ "uk": "ukrainian",
33
+ "el": "greek",
34
+ "ms": "malay",
35
+ "cs": "czech",
36
+ "ro": "romanian",
37
+ "da": "danish",
38
+ "hu": "hungarian",
39
+ "ta": "tamil",
40
+ "no": "norwegian",
41
+ "th": "thai",
42
+ "ur": "urdu",
43
+ "hr": "croatian",
44
+ "bg": "bulgarian",
45
+ "lt": "lithuanian",
46
+ "la": "latin",
47
+ "mi": "maori",
48
+ "ml": "malayalam",
49
+ "cy": "welsh",
50
+ "sk": "slovak",
51
+ "te": "telugu",
52
+ "fa": "persian",
53
+ "lv": "latvian",
54
+ "bn": "bengali",
55
+ "sr": "serbian",
56
+ "az": "azerbaijani",
57
+ "sl": "slovenian",
58
+ "kn": "kannada",
59
+ "et": "estonian",
60
+ "mk": "macedonian",
61
+ "br": "breton",
62
+ "eu": "basque",
63
+ "is": "icelandic",
64
+ "hy": "armenian",
65
+ "ne": "nepali",
66
+ "mn": "mongolian",
67
+ "bs": "bosnian",
68
+ "kk": "kazakh",
69
+ "sq": "albanian",
70
+ "sw": "swahili",
71
+ "gl": "galician",
72
+ "mr": "marathi",
73
+ "pa": "punjabi",
74
+ "si": "sinhala",
75
+ "km": "khmer",
76
+ "sn": "shona",
77
+ "yo": "yoruba",
78
+ "so": "somali",
79
+ "af": "afrikaans",
80
+ "oc": "occitan",
81
+ "ka": "georgian",
82
+ "be": "belarusian",
83
+ "tg": "tajik",
84
+ "sd": "sindhi",
85
+ "gu": "gujarati",
86
+ "am": "amharic",
87
+ "yi": "yiddish",
88
+ "lo": "lao",
89
+ "uz": "uzbek",
90
+ "fo": "faroese",
91
+ "ht": "haitian creole",
92
+ "ps": "pashto",
93
+ "tk": "turkmen",
94
+ "nn": "nynorsk",
95
+ "mt": "maltese",
96
+ "sa": "sanskrit",
97
+ "lb": "luxembourgish",
98
+ "my": "myanmar",
99
+ "bo": "tibetan",
100
+ "tl": "tagalog",
101
+ "mg": "malagasy",
102
+ "as": "assamese",
103
+ "tt": "tatar",
104
+ "haw": "hawaiian",
105
+ "ln": "lingala",
106
+ "ha": "hausa",
107
+ "ba": "bashkir",
108
+ "jw": "javanese",
109
+ "su": "sundanese",
110
+ }
111
+
112
+ # language code lookup by name, with a few language aliases
113
+ TO_LANGUAGE_CODE = {
114
+ **{language: code for code, language in LANGUAGES.items()},
115
+ "burmese": "my",
116
+ "valencian": "ca",
117
+ "flemish": "nl",
118
+ "haitian": "ht",
119
+ "letzeburgesch": "lb",
120
+ "pushto": "ps",
121
+ "panjabi": "pa",
122
+ "moldavian": "ro",
123
+ "moldovan": "ro",
124
+ "sinhalese": "si",
125
+ "castilian": "es",
126
+ }
127
+
128
+
129
+ @dataclass(frozen=True)
130
+ class Tokenizer:
131
+ """A thin wrapper around `GPT2TokenizerFast` providing quick access to special tokens"""
132
+
133
+ tokenizer: "GPT2TokenizerFast"
134
+ language: Optional[str]
135
+ sot_sequence: Tuple[int]
136
+
137
+ def encode(self, text, **kwargs):
138
+ return self.tokenizer.encode(text, **kwargs)
139
+
140
+ def decode(self, token_ids: Union[int, List[int], np.ndarray, torch.Tensor], **kwargs):
141
+ return self.tokenizer.decode(token_ids, **kwargs)
142
+
143
+ def decode_with_timestamps(self, tokens) -> str:
144
+ """
145
+ Timestamp tokens are above the special tokens' id range and are ignored by `decode()`.
146
+ This method decodes given tokens with timestamps tokens annotated, e.g. "<|1.08|>".
147
+ """
148
+ outputs = [[]]
149
+ for token in tokens:
150
+ if token >= self.timestamp_begin:
151
+ timestamp = f"<|{(token - self.timestamp_begin) * 0.02:.2f}|>"
152
+ outputs.append(timestamp)
153
+ outputs.append([])
154
+ else:
155
+ outputs[-1].append(token)
156
+ outputs = [s if isinstance(s, str) else self.tokenizer.decode(s) for s in outputs]
157
+ return "".join(outputs)
158
+
159
+ @property
160
+ @lru_cache()
161
+ def eot(self) -> int:
162
+ return self.tokenizer.eos_token_id
163
+
164
+ @property
165
+ @lru_cache()
166
+ def sot(self) -> int:
167
+ return self._get_single_token_id("<|startoftranscript|>")
168
+
169
+ @property
170
+ @lru_cache()
171
+ def sot_lm(self) -> int:
172
+ return self._get_single_token_id("<|startoflm|>")
173
+
174
+ @property
175
+ @lru_cache()
176
+ def sot_prev(self) -> int:
177
+ return self._get_single_token_id("<|startofprev|>")
178
+
179
+ @property
180
+ @lru_cache()
181
+ def no_speech(self) -> int:
182
+ return self._get_single_token_id("<|nospeech|>")
183
+
184
+ @property
185
+ @lru_cache()
186
+ def no_timestamps(self) -> int:
187
+ return self._get_single_token_id("<|notimestamps|>")
188
+
189
+ @property
190
+ @lru_cache()
191
+ def timestamp_begin(self) -> int:
192
+ return self.tokenizer.all_special_ids[-1] + 1
193
+
194
+ @property
195
+ @lru_cache()
196
+ def language_token(self) -> int:
197
+ """Returns the token id corresponding to the value of the `language` field"""
198
+ if self.language is None:
199
+ raise ValueError(f"This tokenizer does not have language token configured")
200
+
201
+ additional_tokens = dict(
202
+ zip(
203
+ self.tokenizer.additional_special_tokens,
204
+ self.tokenizer.additional_special_tokens_ids,
205
+ )
206
+ )
207
+ candidate = f"<|{self.language}|>"
208
+ if candidate in additional_tokens:
209
+ return additional_tokens[candidate]
210
+
211
+ raise KeyError(f"Language {self.language} not found in tokenizer.")
212
+
213
+ @property
214
+ @lru_cache()
215
+ def all_language_tokens(self) -> Tuple[int]:
216
+ result = []
217
+ for token, token_id in zip(
218
+ self.tokenizer.additional_special_tokens,
219
+ self.tokenizer.additional_special_tokens_ids,
220
+ ):
221
+ if token.strip("<|>") in LANGUAGES:
222
+ result.append(token_id)
223
+ return tuple(result)
224
+
225
+ @property
226
+ @lru_cache()
227
+ def all_language_codes(self) -> Tuple[str]:
228
+ return tuple(self.decode([l]).strip("<|>") for l in self.all_language_tokens)
229
+
230
+ @property
231
+ @lru_cache()
232
+ def sot_sequence_including_notimestamps(self) -> Tuple[int]:
233
+ return tuple(list(self.sot_sequence) + [self.no_timestamps])
234
+
235
+ @property
236
+ @lru_cache()
237
+ def non_speech_tokens(self) -> Tuple[int]:
238
+ """
239
+ Returns the list of tokens to suppress in order to avoid any speaker tags or non-speech
240
+ annotations, to prevent sampling texts that are not actually spoken in the audio, e.g.
241
+
242
+ - ♪♪♪
243
+ - ( SPEAKING FOREIGN LANGUAGE )
244
+ - [DAVID] Hey there,
245
+
246
+ keeping basic punctuations like commas, periods, question marks, exclamation points, etc.
247
+ """
248
+ symbols = list("\"#()*+/:;<=>@[\\]^_`{|}~「」『』")
249
+ symbols += "<< >> <<< >>> -- --- -( -[ (' (\" (( )) ((( ))) [[ ]] {{ }} ♪♪ ♪♪♪".split()
250
+
251
+ # symbols that may be a single token or multiple tokens depending on the tokenizer.
252
+ # In case they're multiple tokens, suppress the first token, which is safe because:
253
+ # These are between U+2640 and U+267F miscellaneous symbols that are okay to suppress
254
+ # in generations, and in the 3-byte UTF-8 representation they share the first two bytes.
255
+ miscellaneous = set("♩♪♫♬♭♮♯")
256
+ assert all(0x2640 <= ord(c) <= 0x267F for c in miscellaneous)
257
+
258
+ # allow hyphens "-" and single quotes "'" between words, but not at the beginning of a word
259
+ result = {self.tokenizer.encode(" -")[0], self.tokenizer.encode(" '")[0]}
260
+ for symbol in symbols + list(miscellaneous):
261
+ for tokens in [self.tokenizer.encode(symbol), self.tokenizer.encode(" " + symbol)]:
262
+ if len(tokens) == 1 or symbol in miscellaneous:
263
+ result.add(tokens[0])
264
+
265
+ return tuple(sorted(result))
266
+
267
+ def _get_single_token_id(self, text) -> int:
268
+ tokens = self.tokenizer.encode(text)
269
+ assert len(tokens) == 1, f"{text} is not encoded as a single token"
270
+ return tokens[0]
271
+
272
+
273
+ @lru_cache(maxsize=None)
274
+ def build_tokenizer(name: str = "gpt2"):
275
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
276
+ path = os.path.join(os.path.dirname(__file__), "assets", name)
277
+ tokenizer = GPT2TokenizerFast.from_pretrained(path)
278
+
279
+ specials = [
280
+ "<|startoftranscript|>",
281
+ *[f"<|{lang}|>" for lang in LANGUAGES.keys()],
282
+ "<|translate|>",
283
+ "<|transcribe|>",
284
+ "<|startoflm|>",
285
+ "<|startofprev|>",
286
+ "<|nospeech|>",
287
+ "<|notimestamps|>",
288
+ ]
289
+
290
+ tokenizer.add_special_tokens(dict(additional_special_tokens=specials))
291
+ return tokenizer
292
+
293
+
294
+ @lru_cache(maxsize=None)
295
+ def get_tokenizer(
296
+ multilingual: bool,
297
+ *,
298
+ task: Optional[str] = None, # Literal["transcribe", "translate", None]
299
+ language: Optional[str] = None,
300
+ ) -> Tokenizer:
301
+ if language is not None:
302
+ language = language.lower()
303
+ if language not in LANGUAGES:
304
+ if language in TO_LANGUAGE_CODE:
305
+ language = TO_LANGUAGE_CODE[language]
306
+ else:
307
+ raise ValueError(f"Unsupported language: {language}")
308
+
309
+ if multilingual:
310
+ tokenizer_name = "multilingual"
311
+ task = task or "transcribe"
312
+ language = language or "en"
313
+ else:
314
+ tokenizer_name = "gpt2"
315
+ task = None
316
+ language = None
317
+
318
+ tokenizer = build_tokenizer(name=tokenizer_name)
319
+ all_special_ids: List[int] = tokenizer.all_special_ids
320
+ sot: int = all_special_ids[1]
321
+ translate: int = all_special_ids[-6]
322
+ transcribe: int = all_special_ids[-5]
323
+
324
+ langs = tuple(LANGUAGES.keys())
325
+ sot_sequence = [sot]
326
+ if language is not None:
327
+ sot_sequence.append(sot + 1 + langs.index(language))
328
+ if task is not None:
329
+ sot_sequence.append(transcribe if task == "transcribe" else translate)
330
+
331
+ return Tokenizer(tokenizer=tokenizer, language=language, sot_sequence=tuple(sot_sequence))
vencoder/whisper/utils.py ADDED
@@ -0,0 +1,163 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import os
3
+ import sys
4
+ import zlib
5
+ from typing import Callable, TextIO
6
+
7
+ system_encoding = sys.getdefaultencoding()
8
+
9
+ if system_encoding != "utf-8":
10
+ def make_safe(string):
11
+ # replaces any character not representable using the system default encoding with an '?',
12
+ # avoiding UnicodeEncodeError (https://github.com/openai/whisper/discussions/729).
13
+ return string.encode(system_encoding, errors="replace").decode(system_encoding)
14
+ else:
15
+ def make_safe(string):
16
+ # utf-8 can encode any Unicode code point, so no need to do the round-trip encoding
17
+ return string
18
+
19
+
20
+ def exact_div(x, y):
21
+ assert x % y == 0
22
+ return x // y
23
+
24
+
25
+ def str2bool(string):
26
+ str2val = {"True": True, "False": False}
27
+ if string in str2val:
28
+ return str2val[string]
29
+ else:
30
+ raise ValueError(f"Expected one of {set(str2val.keys())}, got {string}")
31
+
32
+
33
+ def optional_int(string):
34
+ return None if string == "None" else int(string)
35
+
36
+
37
+ def optional_float(string):
38
+ return None if string == "None" else float(string)
39
+
40
+
41
+ def compression_ratio(text) -> float:
42
+ text_bytes = text.encode("utf-8")
43
+ return len(text_bytes) / len(zlib.compress(text_bytes))
44
+
45
+
46
+ def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = '.'):
47
+ assert seconds >= 0, "non-negative timestamp expected"
48
+ milliseconds = round(seconds * 1000.0)
49
+
50
+ hours = milliseconds // 3_600_000
51
+ milliseconds -= hours * 3_600_000
52
+
53
+ minutes = milliseconds // 60_000
54
+ milliseconds -= minutes * 60_000
55
+
56
+ seconds = milliseconds // 1_000
57
+ milliseconds -= seconds * 1_000
58
+
59
+ hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
60
+ return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
61
+
62
+
63
+ class ResultWriter:
64
+ extension: str
65
+
66
+ def __init__(self, output_dir: str):
67
+ self.output_dir = output_dir
68
+
69
+ def __call__(self, result: dict, audio_path: str):
70
+ audio_basename = os.path.basename(audio_path)
71
+ output_path = os.path.join(self.output_dir, audio_basename + "." + self.extension)
72
+
73
+ with open(output_path, "w", encoding="utf-8") as f:
74
+ self.write_result(result, file=f)
75
+
76
+ def write_result(self, result: dict, file: TextIO):
77
+ raise NotImplementedError
78
+
79
+
80
+ class WriteTXT(ResultWriter):
81
+ extension: str = "txt"
82
+
83
+ def write_result(self, result: dict, file: TextIO):
84
+ for segment in result["segments"]:
85
+ print(segment['text'].strip(), file=file, flush=True)
86
+
87
+
88
+ class WriteVTT(ResultWriter):
89
+ extension: str = "vtt"
90
+
91
+ def write_result(self, result: dict, file: TextIO):
92
+ print("WEBVTT\n", file=file)
93
+ for segment in result["segments"]:
94
+ print(
95
+ f"{format_timestamp(segment['start'])} --> {format_timestamp(segment['end'])}\n"
96
+ f"{segment['text'].strip().replace('-->', '->')}\n",
97
+ file=file,
98
+ flush=True,
99
+ )
100
+
101
+
102
+ class WriteSRT(ResultWriter):
103
+ extension: str = "srt"
104
+
105
+ def write_result(self, result: dict, file: TextIO):
106
+ for i, segment in enumerate(result["segments"], start=1):
107
+ # write srt lines
108
+ print(
109
+ f"{i}\n"
110
+ f"{format_timestamp(segment['start'], always_include_hours=True, decimal_marker=',')} --> "
111
+ f"{format_timestamp(segment['end'], always_include_hours=True, decimal_marker=',')}\n"
112
+ f"{segment['text'].strip().replace('-->', '->')}\n",
113
+ file=file,
114
+ flush=True,
115
+ )
116
+
117
+
118
+ class WriteTSV(ResultWriter):
119
+ """
120
+ Write a transcript to a file in TSV (tab-separated values) format containing lines like:
121
+ <start time in integer milliseconds>\t<end time in integer milliseconds>\t<transcript text>
122
+
123
+ Using integer milliseconds as start and end times means there's no chance of interference from
124
+ an environment setting a language encoding that causes the decimal in a floating point number
125
+ to appear as a comma; also is faster and more efficient to parse & store, e.g., in C++.
126
+ """
127
+ extension: str = "tsv"
128
+
129
+ def write_result(self, result: dict, file: TextIO):
130
+ print("start", "end", "text", sep="\t", file=file)
131
+ for segment in result["segments"]:
132
+ print(round(1000 * segment['start']), file=file, end="\t")
133
+ print(round(1000 * segment['end']), file=file, end="\t")
134
+ print(segment['text'].strip().replace("\t", " "), file=file, flush=True)
135
+
136
+
137
+ class WriteJSON(ResultWriter):
138
+ extension: str = "json"
139
+
140
+ def write_result(self, result: dict, file: TextIO):
141
+ json.dump(result, file)
142
+
143
+
144
+ def get_writer(output_format: str, output_dir: str) -> Callable[[dict, TextIO], None]:
145
+ writers = {
146
+ "txt": WriteTXT,
147
+ "vtt": WriteVTT,
148
+ "srt": WriteSRT,
149
+ "tsv": WriteTSV,
150
+ "json": WriteJSON,
151
+ }
152
+
153
+ if output_format == "all":
154
+ all_writers = [writer(output_dir) for writer in writers.values()]
155
+
156
+ def write_all(result: dict, file: TextIO):
157
+ for writer in all_writers:
158
+ writer(result, file)
159
+
160
+ return write_all
161
+
162
+ return writers[output_format](output_dir)
163
+