Removed the local copy of the whisperspeech library

whisperspeech/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.5.6"

whisperspeech/_modidx.py

@@ -1,615 +0,0 @@
-# Autogenerated by nbdev
-
-d = { 'settings': { 'branch': 'master',
- 'doc_baseurl': '/WhisperSpeech',
- 'doc_host': 'https://collabora.github.io',
- 'git_url': 'https://github.com/collabora/WhisperSpeech',
- 'lib_path': 'whisperspeech'},
- 'syms': { 'whisperspeech.a2wav': { 'whisperspeech.a2wav.Vocoder': ('6. quality-boosting vocoder.html#vocoder', 'whisperspeech/a2wav.py'),
- 'whisperspeech.a2wav.Vocoder.__init__': ( '6. quality-boosting vocoder.html#vocoder.__init__',
- 'whisperspeech/a2wav.py'),
- 'whisperspeech.a2wav.Vocoder.decode': ( '6. quality-boosting vocoder.html#vocoder.decode',
- 'whisperspeech/a2wav.py'),
- 'whisperspeech.a2wav.Vocoder.decode_to_file': ( '6. quality-boosting '
- 'vocoder.html#vocoder.decode_to_file',
- 'whisperspeech/a2wav.py'),
- 'whisperspeech.a2wav.Vocoder.decode_to_notebook': ( '6. quality-boosting '
- 'vocoder.html#vocoder.decode_to_notebook',
- 'whisperspeech/a2wav.py')},
- 'whisperspeech.extract_acoustic': { 'whisperspeech.extract_acoustic.extract_Atoks': ( '1. acoustic token '
- 'extraction.html#extract_atoks',
- 'whisperspeech/extract_acoustic.py'),
- 'whisperspeech.extract_acoustic.extract_acoustic': ( '1. acoustic token '
- 'extraction.html#extract_acoustic',
- 'whisperspeech/extract_acoustic.py'),
- 'whisperspeech.extract_acoustic.load': ( '1. acoustic token extraction.html#load',
- 'whisperspeech/extract_acoustic.py'),
- 'whisperspeech.extract_acoustic.load_model': ( '1. acoustic token '
- 'extraction.html#load_model',
- 'whisperspeech/extract_acoustic.py')},
- 'whisperspeech.extract_semb': { 'whisperspeech.extract_semb.encode_semantic': ( '2c. whisper semantic embedding '
- 'extraction.html#encode_semantic',
- 'whisperspeech/extract_semb.py'),
- 'whisperspeech.extract_semb.extract_semantic': ( '2c. whisper semantic embedding '
- 'extraction.html#extract_semantic',
- 'whisperspeech/extract_semb.py'),
- 'whisperspeech.extract_semb.load_model': ( '2c. whisper semantic embedding '
- 'extraction.html#load_model',
- 'whisperspeech/extract_semb.py')},
- 'whisperspeech.fetch_models': { 'whisperspeech.fetch_models.main': ( '0. download models.html#main',
- 'whisperspeech/fetch_models.py')},
- 'whisperspeech.modules': { 'whisperspeech.modules.Decoder': ('a. neural modules.html#decoder', 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Decoder.__init__': ( 'a. neural modules.html#decoder.__init__',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Decoder.forward': ( 'a. neural modules.html#decoder.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Encoder': ('a. neural modules.html#encoder', 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Encoder.__init__': ( 'a. neural modules.html#encoder.__init__',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Encoder.forward': ( 'a. neural modules.html#encoder.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.LayerNorm': ('a. neural modules.html#layernorm', 'whisperspeech/modules.py'),
- 'whisperspeech.modules.LayerNorm.forward': ( 'a. neural modules.html#layernorm.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.LinearHead': ( 'a. neural modules.html#linearhead',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.MultiHeadAttention': ( 'a. neural modules.html#multiheadattention',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.MultiHeadAttention.__init__': ( 'a. neural '
- 'modules.html#multiheadattention.__init__',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.MultiHeadAttention.forward': ( 'a. neural '
- 'modules.html#multiheadattention.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.MultiHeadAttention.qkv_attention_pth20': ( 'a. neural '
- 'modules.html#multiheadattention.qkv_attention_pth20',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.MultiHeadAttention.qkv_attention_vanilla': ( 'a. neural '
- 'modules.html#multiheadattention.qkv_attention_vanilla',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.MultiHeadAttention.qkv_attention_xformers': ( 'a. neural '
- 'modules.html#multiheadattention.qkv_attention_xformers',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.QueryHead': ('a. neural modules.html#queryhead', 'whisperspeech/modules.py'),
- 'whisperspeech.modules.ResidualAttentionBlock': ( 'a. neural modules.html#residualattentionblock',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.ResidualAttentionBlock.__init__': ( 'a. neural '
- 'modules.html#residualattentionblock.__init__',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.ResidualAttentionBlock.forward': ( 'a. neural '
- 'modules.html#residualattentionblock.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Rotary': ('a. neural modules.html#rotary', 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Rotary.__init__': ( 'a. neural modules.html#rotary.__init__',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.Rotary.forward': ( 'a. neural modules.html#rotary.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.SumDecoder': ( 'a. neural modules.html#sumdecoder',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.SumDecoder.__init__': ( 'a. neural modules.html#sumdecoder.__init__',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.SumDecoder.forward': ( 'a. neural modules.html#sumdecoder.forward',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.apply_rotary_pos_emb': ( 'a. neural modules.html#apply_rotary_pos_emb',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.init_transformer': ( 'a. neural modules.html#init_transformer',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.rotate_half': ( 'a. neural modules.html#rotate_half',
- 'whisperspeech/modules.py'),
- 'whisperspeech.modules.sinusoids': ('a. neural modules.html#sinusoids', 'whisperspeech/modules.py')},
- 'whisperspeech.pipeline': { 'whisperspeech.pipeline.Pipeline': ('7. pipeline.html#pipeline', 'whisperspeech/pipeline.py'),
- 'whisperspeech.pipeline.Pipeline.__init__': ( '7. pipeline.html#pipeline.__init__',
- 'whisperspeech/pipeline.py'),
- 'whisperspeech.pipeline.Pipeline.generate': ( '7. pipeline.html#pipeline.generate',
- 'whisperspeech/pipeline.py'),
- 'whisperspeech.pipeline.Pipeline.generate_atoks': ( '7. pipeline.html#pipeline.generate_atoks',
- 'whisperspeech/pipeline.py'),
- 'whisperspeech.pipeline.Pipeline.generate_to_file': ( '7. pipeline.html#pipeline.generate_to_file',
- 'whisperspeech/pipeline.py'),
- 'whisperspeech.pipeline.Pipeline.generate_to_notebook': ( '7. '
- 'pipeline.html#pipeline.generate_to_notebook',
- 'whisperspeech/pipeline.py')},
- 'whisperspeech.prepare_s2a_dataset': { 'whisperspeech.prepare_s2a_dataset.flac_to_s2a_name': ( '4a. s2a dataset '
- 'preparation.html#flac_to_s2a_name',
- 'whisperspeech/prepare_s2a_dataset.py'),
- 'whisperspeech.prepare_s2a_dataset.prepare_s2a': ( '4a. s2a dataset '
- 'preparation.html#prepare_s2a',
- 'whisperspeech/prepare_s2a_dataset.py'),
- 'whisperspeech.prepare_s2a_dataset.resampler': ( '4a. s2a dataset '
- 'preparation.html#resampler',
- 'whisperspeech/prepare_s2a_dataset.py')},
- 'whisperspeech.prepare_t2s_dataset': { 'whisperspeech.prepare_t2s_dataset.Transcriber': ( '5a. t2s dataset '
- 'preparation.html#transcriber',
- 'whisperspeech/prepare_t2s_dataset.py'),
- 'whisperspeech.prepare_t2s_dataset.Transcriber.__init__': ( '5a. t2s dataset '
- 'preparation.html#transcriber.__init__',
- 'whisperspeech/prepare_t2s_dataset.py'),
- 'whisperspeech.prepare_t2s_dataset.Transcriber.transcribe': ( '5a. t2s dataset '
- 'preparation.html#transcriber.transcribe',
- 'whisperspeech/prepare_t2s_dataset.py'),
- 'whisperspeech.prepare_t2s_dataset.flac_to_t2s_name': ( '5a. t2s dataset '
- 'preparation.html#flac_to_t2s_name',
- 'whisperspeech/prepare_t2s_dataset.py'),
- 'whisperspeech.prepare_t2s_dataset.prepare_t2s': ( '5a. t2s dataset '
- 'preparation.html#prepare_t2s',
- 'whisperspeech/prepare_t2s_dataset.py')},
- 'whisperspeech.s2a_delar_mup_wds': { 'whisperspeech.s2a_delar_mup_wds.CMLMVisual': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.__init__': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual.__init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.add_data': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual.add_data',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.add_table_row': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#cmlmvisual.add_table_row',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.hide': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual.hide',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.on_iter': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual.on_iter',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.plot': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual.plot',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.CMLMVisual.show': ( '4b. semantic to acoustic token '
- 'modeling.html#cmlmvisual.show',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.DelSumDecoder': ( '4b. semantic to acoustic token '
- 'modeling.html#delsumdecoder',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.DelSumDecoder.__init__': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#delsumdecoder.__init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.DelSumDecoder.forward': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#delsumdecoder.forward',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.EmbeddingProjector': ( '4b. semantic to acoustic token '
- 'modeling.html#embeddingprojector',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.MultiHeadAttention': ( '4b. semantic to acoustic token '
- 'modeling.html#multiheadattention',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.MultiHeadAttention.__init__': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#multiheadattention.__init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.MultiHeadAttention.forward': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#multiheadattention.forward',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.MultiHeadAttention.qkv_attention_pth20': ( '4b. semantic '
- 'to acoustic '
- 'token '
- 'modeling.html#multiheadattention.qkv_attention_pth20',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.MultiHeadAttention.qkv_attention_xformers': ( '4b. '
- 'semantic '
- 'to '
- 'acoustic '
- 'token '
- 'modeling.html#multiheadattention.qkv_attention_xformers',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.ResidualAttentionBlock': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#residualattentionblock',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.ResidualAttentionBlock.__init__': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#residualattentionblock.__init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.ResidualAttentionBlock.forward': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#residualattentionblock.forward',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.Rotary': ( '4b. semantic to acoustic token '
- 'modeling.html#rotary',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.Rotary.__init__': ( '4b. semantic to acoustic token '
- 'modeling.html#rotary.__init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.Rotary.forward': ( '4b. semantic to acoustic token '
- 'modeling.html#rotary.forward',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer': ( '4b. semantic to acoustic token '
- 'modeling.html#sadelartransformer',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.__init__': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.__init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.device': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#sadelartransformer.device',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.embed_stoks': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.embed_stoks',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.forward': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#sadelartransformer.forward',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.generate': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.generate',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.get_extra_state': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.get_extra_state',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.get_metrics': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.get_metrics',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.init_transformer': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.init_transformer',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.load_checkpoint': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.load_checkpoint',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.load_frozen_semantic_embeddings': ( '4b. '
- 'semantic '
- 'to '
- 'acoustic '
- 'token '
- 'modeling.html#sadelartransformer.load_frozen_semantic_embeddings',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.load_model': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.load_model',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.save_model': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.save_model',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.set_extra_state': ( '4b. semantic to '
- 'acoustic token '
- 'modeling.html#sadelartransformer.set_extra_state',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.SADelARTransformer.setup': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#sadelartransformer.setup',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.Tunables': ( '4b. semantic to acoustic token '
- 'modeling.html#tunables',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.Tunables.__post_init__': ( '4b. semantic to acoustic '
- 'token '
- 'modeling.html#tunables.__post_init__',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.Tunables.upgrade': ( '4b. semantic to acoustic token '
- 'modeling.html#tunables.upgrade',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds._make_model': ( '4b. semantic to acoustic token '
- 'modeling.html#_make_model',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.apply_rotary_pos_emb': ( '4b. semantic to acoustic token '
- 'modeling.html#apply_rotary_pos_emb',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.load_datasets': ( '4b. semantic to acoustic token '
- 'modeling.html#load_datasets',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.make_model': ( '4b. semantic to acoustic token '
- 'modeling.html#make_model',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.pad_samples': ( '4b. semantic to acoustic token '
- 'modeling.html#pad_samples',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.rand': ( '4b. semantic to acoustic token '
- 'modeling.html#rand',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.random_trunc': ( '4b. semantic to acoustic token '
- 'modeling.html#random_trunc',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.rotate_half': ( '4b. semantic to acoustic token '
- 'modeling.html#rotate_half',
- 'whisperspeech/s2a_delar_mup_wds.py'),
- 'whisperspeech.s2a_delar_mup_wds.speaker_id_extractor': ( '4b. semantic to acoustic token '
- 'modeling.html#speaker_id_extractor',
- 'whisperspeech/s2a_delar_mup_wds.py')},
- 'whisperspeech.t2s_up_wds': { 'whisperspeech.t2s_up_wds.CharTokenizer': ( '5b. text to semantic token '
- 'modeling.html#chartokenizer',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.CharTokenizer.decode': ( '5b. text to semantic token '
- 'modeling.html#chartokenizer.decode',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.CharTokenizer.encode': ( '5b. text to semantic token '
- 'modeling.html#chartokenizer.encode',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Decoder': ( '5b. text to semantic token modeling.html#decoder',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Decoder.__init__': ( '5b. text to semantic token '
- 'modeling.html#decoder.__init__',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Decoder.forward': ( '5b. text to semantic token '
- 'modeling.html#decoder.forward',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.EmbeddingProjector': ( '5b. text to semantic token '
- 'modeling.html#embeddingprojector',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Encoder': ( '5b. text to semantic token modeling.html#encoder',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Encoder.__init__': ( '5b. text to semantic token '
- 'modeling.html#encoder.__init__',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Encoder.forward': ( '5b. text to semantic token '
- 'modeling.html#encoder.forward',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.__init__': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.__init__',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.device': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.device',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.ensure_tokenizer': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.ensure_tokenizer',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.forward': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.forward',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.generate': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.generate',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.generate_batch': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.generate_batch',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.init_transformer': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.init_transformer',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.load_checkpoint': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.load_checkpoint',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.load_frozen_semantic_embeddings': ( '5b. text to '
- 'semantic token '
- 'modeling.html#tsartransformer.load_frozen_semantic_embeddings',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.load_model': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.load_model',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.save_model': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.save_model',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.TSARTransformer.setup': ( '5b. text to semantic token '
- 'modeling.html#tsartransformer.setup',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Tunables': ( '5b. text to semantic token modeling.html#tunables',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.Tunables.__post_init__': ( '5b. text to semantic token '
- 'modeling.html#tunables.__post_init__',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds._make_model': ( '5b. text to semantic token modeling.html#_make_model',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.ar_padder': ( '5b. text to semantic token modeling.html#ar_padder',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.build_speaker_map': ( '5b. text to semantic token '
- 'modeling.html#build_speaker_map',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.char_per_seconder': ( '5b. text to semantic token '
- 'modeling.html#char_per_seconder',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.load_datasets': ( '5b. text to semantic token '
- 'modeling.html#load_datasets',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.make_model': ( '5b. text to semantic token modeling.html#make_model',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.rand': ( '5b. text to semantic token modeling.html#rand',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.speaker_id_extractor': ( '5b. text to semantic token '
- 'modeling.html#speaker_id_extractor',
- 'whisperspeech/t2s_up_wds.py'),
- 'whisperspeech.t2s_up_wds.tokenizer': ( '5b. text to semantic token modeling.html#tokenizer',
- 'whisperspeech/t2s_up_wds.py')},
- 'whisperspeech.train': { 'whisperspeech.train.SimpleVisual': ('b1. training.html#simplevisual', 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.__init__': ( 'b1. training.html#simplevisual.__init__',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.add_data': ( 'b1. training.html#simplevisual.add_data',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.add_table_row': ( 'b1. training.html#simplevisual.add_table_row',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.hide': ( 'b1. training.html#simplevisual.hide',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.on_iter': ( 'b1. training.html#simplevisual.on_iter',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.plot': ( 'b1. training.html#simplevisual.plot',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.SimpleVisual.show': ( 'b1. training.html#simplevisual.show',
- 'whisperspeech/train.py'),
- 'whisperspeech.train.train': ('b1. training.html#train', 'whisperspeech/train.py'),
- 'whisperspeech.train.validate': ('b1. training.html#validate', 'whisperspeech/train.py')},
- 'whisperspeech.train_multi': { 'whisperspeech.train_multi.TrainingTask': ( 'b2. training (lightning).html#trainingtask',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.__init__': ( 'b2. training '
- '(lightning).html#trainingtask.__init__',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.configure_optimizers': ( 'b2. training '
- '(lightning).html#trainingtask.configure_optimizers',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.on_fit_start': ( 'b2. training '
- '(lightning).html#trainingtask.on_fit_start',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.on_validation_epoch_end': ( 'b2. training '
- '(lightning).html#trainingtask.on_validation_epoch_end',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.test_step': ( 'b2. training '
- '(lightning).html#trainingtask.test_step',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.training_step': ( 'b2. training '
- '(lightning).html#trainingtask.training_step',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.TrainingTask.validation_step': ( 'b2. training '
- '(lightning).html#trainingtask.validation_step',
- 'whisperspeech/train_multi.py'),
- 'whisperspeech.train_multi.parse_and_call': ( 'b2. training (lightning).html#parse_and_call',
- 'whisperspeech/train_multi.py')},
- 'whisperspeech.vad': { 'whisperspeech.vad.extract_segments': ( '1b. voice activity detection.html#extract_segments',
- 'whisperspeech/vad.py'),
- 'whisperspeech.vad.fix_dots_in_names': ( '1b. voice activity detection.html#fix_dots_in_names',
- 'whisperspeech/vad.py'),
- 'whisperspeech.vad.flac_to_vad_name': ( '1b. voice activity detection.html#flac_to_vad_name',
- 'whisperspeech/vad.py'),
- 'whisperspeech.vad.load_dataset': ( '1b. voice activity detection.html#load_dataset',
- 'whisperspeech/vad.py'),
- 'whisperspeech.vad.process_shard': ( '1b. voice activity detection.html#process_shard',
- 'whisperspeech/vad.py'),
- 'whisperspeech.vad.segment_audio': ( '1b. voice activity detection.html#segment_audio',
- 'whisperspeech/vad.py')},
- 'whisperspeech.verify_wds': { 'whisperspeech.verify_wds.process_shard': ( '0. verify webdataset archives.html#process_shard',
- 'whisperspeech/verify_wds.py')},
- 'whisperspeech.vq_stoks': { 'whisperspeech.vq_stoks.RQBottleneckTransformer': ( '2b. whisper quantization (semantic token) '
- 'model.html#rqbottlenecktransformer',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.__init__': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.__init__',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.decode_text': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.decode_text',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.dequantize': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.dequantize',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.device': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.device',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.downsample_embeddings': ( '2b. whisper '
- 'quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.downsample_embeddings',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.encode_audio': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.encode_audio',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.encode_mel': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.encode_mel',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.ensure_whisper': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.ensure_whisper',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.extract_teacher': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.extract_teacher',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.forward': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.forward',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.get_metrics': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.get_metrics',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.init_transformer': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.init_transformer',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.load_checkpoint': ( '2b. whisper quantization '
- '(semantic token) '
- 'model.html#rqbottlenecktransformer.load_checkpoint',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.load_model': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.load_model',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.quantize': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.quantize',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.save_model': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.save_model',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.RQBottleneckTransformer.setup': ( '2b. whisper quantization (semantic '
- 'token) '
- 'model.html#rqbottlenecktransformer.setup',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.Tunables': ( '2b. whisper quantization (semantic token) '
- 'model.html#tunables',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.Tunables.__post_init__': ( '2b. whisper quantization (semantic token) '
- 'model.html#tunables.__post_init__',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.Tunables.upgrade': ( '2b. whisper quantization (semantic token) '
- 'model.html#tunables.upgrade',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.add_masks': ( '2b. whisper quantization (semantic token) '
- 'model.html#add_masks',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.derived_dataset': ( '2b. whisper quantization (semantic token) '
- 'model.html#derived_dataset',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.load_datasets': ( '2b. whisper quantization (semantic token) '
- 'model.html#load_datasets',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.logrand': ( '2b. whisper quantization (semantic token) model.html#logrand',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.make_model': ( '2b. whisper quantization (semantic token) '
- 'model.html#make_model',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.merge_in': ( '2b. whisper quantization (semantic token) '
- 'model.html#merge_in',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.rand': ( '2b. whisper quantization (semantic token) model.html#rand',
- 'whisperspeech/vq_stoks.py'),
- 'whisperspeech.vq_stoks.tokenize_text': ( '2b. whisper quantization (semantic token) '
- 'model.html#tokenize_text',
- 'whisperspeech/vq_stoks.py')},
- 'whisperspeech.wer_metrics': { 'whisperspeech.wer_metrics.DfBuilder': ( 'c. word error rate metrics.html#dfbuilder',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.DfBuilder.__init__': ( 'c. word error rate '
- 'metrics.html#dfbuilder.__init__',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.DfBuilder.df': ( 'c. word error rate metrics.html#dfbuilder.df',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.DfBuilder.push': ( 'c. word error rate metrics.html#dfbuilder.push',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.WERStats': ( 'c. word error rate metrics.html#werstats',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.WERStats.__init__': ( 'c. word error rate '
- 'metrics.html#werstats.__init__',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.WERStats.push_sample': ( 'c. word error rate '
- 'metrics.html#werstats.push_sample',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.librispeech_data': ( 'c. word error rate '
- 'metrics.html#librispeech_data',
- 'whisperspeech/wer_metrics.py'),
- 'whisperspeech.wer_metrics.whisper_normalize': ( 'c. word error rate '
- 'metrics.html#whisper_normalize',
- 'whisperspeech/wer_metrics.py')},
- 'whisperspeech.wh_transcribe': { 'whisperspeech.wh_transcribe.chunk_merger': ( '2a. whisper quantization dataset '
- 'preparation.html#chunk_merger',
- 'whisperspeech/wh_transcribe.py'),
- 'whisperspeech.wh_transcribe.flac_to_txt_name': ( '2a. whisper quantization dataset '
- 'preparation.html#flac_to_txt_name',
- 'whisperspeech/wh_transcribe.py'),
- 'whisperspeech.wh_transcribe.merge_in': ( '2a. whisper quantization dataset '
- 'preparation.html#merge_in',
- 'whisperspeech/wh_transcribe.py'),
- 'whisperspeech.wh_transcribe.process_shard': ( '2a. whisper quantization dataset '
- 'preparation.html#process_shard',
- 'whisperspeech/wh_transcribe.py'),
- 'whisperspeech.wh_transcribe.random_cutter': ( '2a. whisper quantization dataset '
- 'preparation.html#random_cutter',
- 'whisperspeech/wh_transcribe.py'),
- 'whisperspeech.wh_transcribe.split_to_chunks': ( '2a. whisper quantization dataset '
- 'preparation.html#split_to_chunks',
- 'whisperspeech/wh_transcribe.py'),
- 'whisperspeech.wh_transcribe.wds_compose': ( '2a. whisper quantization dataset '
- 'preparation.html#wds_compose',
- 'whisperspeech/wh_transcribe.py')}}}

whisperspeech/a2wav.py

@@ -1,45 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/6. Quality-boosting vocoder.ipynb.
-
-# %% auto 0
-__all__ = ['Vocoder']
-
-# %% ../nbs/6. Quality-boosting vocoder.ipynb 1
-from vocos import Vocos
-import torch
-import torchaudio
-
-# %% ../nbs/6. Quality-boosting vocoder.ipynb 2
-class Vocoder:
- def __init__(self, repo_id="charactr/vocos-encodec-24khz"):
- self.vocos = Vocos.from_pretrained(repo_id).cuda()
- 
- def is_notebook(self):
- try:
- return get_ipython().__class__.__name__ == "ZMQInteractiveShell"
- except:
- return False
-
- @torch.no_grad()
- def decode(self, atoks):
- if len(atoks.shape) == 3:
- b,q,t = atoks.shape
- atoks = atoks.permute(1,0,2)
- else:
- q,t = atoks.shape
- 
- features = self.vocos.codes_to_features(atoks)
- bandwidth_id = torch.tensor({2:0,4:1,8:2}[q]).cuda()
- return self.vocos.decode(features, bandwidth_id=bandwidth_id)
- 
- def decode_to_file(self, fname, atoks):
- audio = self.decode(atoks)
- torchaudio.save(fname, audio.cpu(), 24000)
- if self.is_notebook():
- from IPython.display import display, HTML, Audio
- display(HTML(f'<a href="{fname}" target="_blank">Listen to {fname}</a>'))
- 
- def decode_to_notebook(self, atoks):
- from IPython.display import display, HTML, Audio
-
- audio = self.decode(atoks)
- display(Audio(audio.cpu().numpy(), rate=24000))

whisperspeech/extract_acoustic.py

@@ -1,56 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/1. Acoustic token extraction.ipynb.
-
-# %% auto 0
-__all__ = ['load', 'load_model', 'extract_Atoks', 'extract_acoustic']
-
-# %% ../nbs/1. Acoustic token extraction.ipynb 2
-import torch
-import torchaudio
-import gc
-
-from pathlib import Path
-from fastcore.script import *
-from fastprogress import progress_bar, master_bar
-
-# %% ../nbs/1. Acoustic token extraction.ipynb 5
-def load(fname, newsr=24000):
- """Load an audio file to the GPU and resample to `newsr`."""
- x, sr = torchaudio.load(fname)
- _tform = torchaudio.transforms.Resample(sr, newsr)
- return _tform(x).cuda().unsqueeze(0)
-
-# %% ../nbs/1. Acoustic token extraction.ipynb 6
-def load_model():
- "Load the pretrained EnCodec model"
- from encodec.model import EncodecModel
- model = EncodecModel.encodec_model_24khz()
- model.set_target_bandwidth(1.5)
- model.cuda().eval();
- return model
-
-# %% ../nbs/1. Acoustic token extraction.ipynb 7
-def extract_Atoks(model, audio):
- """Extract EnCodec tokens for the given `audio` tensor (or file path)
- using the given `model` (see `load_model`)."""
- if isinstance(audio, (Path, str)):
- audio = load(audio)
- with torch.no_grad():
- frames = torch.cat([model.encode(segment)[0][0]
- for segment in torch.split(audio, 320*20000, dim=-1)], dim=-1)
- return frames
-
-# %% ../nbs/1. Acoustic token extraction.ipynb 8
-@call_parse
-def extract_acoustic(
- srcdir:Path, # source dir, should contain *.flac files
- outdir:Path, # output dir, will get the *.encodec files
- ): 
- "Convert audio files to .encodec files with tensors of tokens"
- model = load_model()
- outdir.mkdir(exist_ok=True, parents=True)
- for name in progress_bar(list(srcdir.rglob('*.flac'))):
- outname = outdir/name.with_suffix('.encodec').name
- tokens = extract_Atoks(model, name)
- torch.save(tokens, outname)
- del tokens
- gc.collect()

whisperspeech/fetch_models.py

@@ -1,17 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/0. Download models.ipynb.
-
-# %% auto 0
-__all__ = []
-
-# %% ../nbs/0. Download models.ipynb 1
-from fastcore.script import call_parse
-import whisperx
-import whisper
-
-# %% ../nbs/0. Download models.ipynb 3
-@call_parse
-def main():
- whisper.load_model('base.en')
- whisper.load_model('small.en')
- whisperx.vad.load_vad_model('cpu')
- whisperx.asr.load_model('medium.en', "cpu", compute_type="float16", language='en')

whisperspeech/languages.py

@@ -1,131 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B. Languages.ipynb.
-
-# %% auto 0
-__all__ = ['to_id']
-
-# %% ../nbs/B. Languages.ipynb 3
-LANGUAGES = {
- "en": "english",
- "zh": "chinese",
- "de": "german",
- "es": "spanish",
- "ru": "russian",
- "ko": "korean",
- "fr": "french",
- "ja": "japanese",
- "pt": "portuguese",
- "tr": "turkish",
- "pl": "polish",
- "ca": "catalan",
- "nl": "dutch",
- "ar": "arabic",
- "sv": "swedish",
- "it": "italian",
- "id": "indonesian",
- "hi": "hindi",
- "fi": "finnish",
- "vi": "vietnamese",
- "he": "hebrew",
- "uk": "ukrainian",
- "el": "greek",
- "ms": "malay",
- "cs": "czech",
- "ro": "romanian",
- "da": "danish",
- "hu": "hungarian",
- "ta": "tamil",
- "no": "norwegian",
- "th": "thai",
- "ur": "urdu",
- "hr": "croatian",
- "bg": "bulgarian",
- "lt": "lithuanian",
- "la": "latin",
- "mi": "maori",
- "ml": "malayalam",
- "cy": "welsh",
- "sk": "slovak",
- "te": "telugu",
- "fa": "persian",
- "lv": "latvian",
- "bn": "bengali",
- "sr": "serbian",
- "az": "azerbaijani",
- "sl": "slovenian",
- "kn": "kannada",
- "et": "estonian",
- "mk": "macedonian",
- "br": "breton",
- "eu": "basque",
- "is": "icelandic",
- "hy": "armenian",
- "ne": "nepali",
- "mn": "mongolian",
- "bs": "bosnian",
- "kk": "kazakh",
- "sq": "albanian",
- "sw": "swahili",
- "gl": "galician",
- "mr": "marathi",
- "pa": "punjabi",
- "si": "sinhala",
- "km": "khmer",
- "sn": "shona",
- "yo": "yoruba",
- "so": "somali",
- "af": "afrikaans",
- "oc": "occitan",
- "ka": "georgian",
- "be": "belarusian",
- "tg": "tajik",
- "sd": "sindhi",
- "gu": "gujarati",
- "am": "amharic",
- "yi": "yiddish",
- "lo": "lao",
- "uz": "uzbek",
- "fo": "faroese",
- "ht": "haitian creole",
- "ps": "pashto",
- "tk": "turkmen",
- "nn": "nynorsk",
- "mt": "maltese",
- "sa": "sanskrit",
- "lb": "luxembourgish",
- "my": "myanmar",
- "bo": "tibetan",
- "tl": "tagalog",
- "mg": "malagasy",
- "as": "assamese",
- "tt": "tatar",
- "haw": "hawaiian",
- "ln": "lingala",
- "ha": "hausa",
- "ba": "bashkir",
- "jw": "javanese",
- "su": "sundanese",
-}
-
-# %% ../nbs/B. Languages.ipynb 4
-# language code lookup by name, with a few language aliases
-TO_LANGUAGE_CODE = {
- **{language: code for code, language in LANGUAGES.items()},
- "burmese": "my",
- "valencian": "ca",
- "flemish": "nl",
- "haitian": "ht",
- "letzeburgesch": "lb",
- "pushto": "ps",
- "panjabi": "pa",
- "moldavian": "ro",
- "moldovan": "ro",
- "sinhalese": "si",
- "castilian": "es",
-}
-
-# %% ../nbs/B. Languages.ipynb 5
-languages = tuple(LANGUAGES.keys())
-
-# %% ../nbs/B. Languages.ipynb 6
-def to_id(lang):
- return languages.index(TO_LANGUAGE_CODE.get(lang, lang))

whisperspeech/modules.py

@@ -1,331 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/A. Neural modules.ipynb.
-
-# %% auto 0
-__all__ = ['LayerNorm', 'LinearHead', 'QueryHead', 'init_transformer', 'sinusoids', 'MultiHeadAttention',
- 'ResidualAttentionBlock', 'BaseDecoder', 'EmbeddingProjector', 'FlexEmbeddings']
-
-# %% ../nbs/A. Neural modules.ipynb 2
-import torch
-import numpy as np
-import math
-
-from torch import Tensor, nn
-import torch.nn.functional as F
-from typing import Dict, Iterable, Optional
-
-# import xformers.ops as xops
-
-# %% ../nbs/A. Neural modules.ipynb 3
-# Code in this file is mostly borrowed from
-# https://github.com/openai/whisper/blob/main/whisper/model.py
-# and is under the MIT License
-
-class LayerNorm(nn.LayerNorm):
- def forward(self, x):
- return super().forward(x.float()).type(x.dtype)
-
-# Used in μP to initialize the weights and configure the optimizer
-# These two layers map the transformer width into a fixed dimension
-class LinearHead(nn.Linear):
- pass
-
-class QueryHead(nn.Linear):
- pass
-
-# based on https://github.com/karpathy/minGPT/blob/master/mingpt/model.py#L163
-def init_transformer(m):
- if isinstance(m, (nn.Linear, nn.Embedding)):
- torch.nn.init.trunc_normal_(m.weight, std=.02)
- if isinstance(m, nn.Linear) and m.bias is not None:
- torch.nn.init.constant_(m.bias, 0)
- elif isinstance(m, nn.LayerNorm):
- torch.nn.init.constant_(m.bias, 0)
- torch.nn.init.constant_(m.weight, 1.0)
-
-# %% ../nbs/A. Neural modules.ipynb 4
-def sinusoids(length, channels, max_timescale=10000):
- """Returns sinusoids for positional embedding"""
- assert channels % 2 == 0
- log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
- inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
- scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
- return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
-
-# %% ../nbs/A. Neural modules.ipynb 5
-class MultiHeadAttention(nn.Module):
- def __init__(self, n_state: int, n_head: int, qk_scale: float = 1, rope: bool = False, cross=False):
- super().__init__()
- self.n_state = n_state
- self.n_head = n_head
- self.sqrt_qk_scale = math.sqrt(qk_scale)
- self.query = QueryHead(n_state, n_state)
- self.key = nn.Linear(n_state, n_state, bias=False)
- self.value = nn.Linear(n_state, n_state)
- self.out = nn.Linear(n_state, n_state)
- self.cross = cross
- self.query_subsampling = 1
- self.key_subsampling = 1
-
- self.cached_kvx = None
- self.register_buffer('k_cache', None)
- self.register_buffer('v_cache', None)
- 
- self.rotary = None
- if rope:
- self.rotary = Rotary(n_state // n_head)
- self.qkv = None
- self.kv = None
-
- def setup_kv_cache(self, max_batch_size, max_seq_len, dtype=torch.float32):
- cache_shape = (max_batch_size, self.n_head, max_seq_len, self.n_state//self.n_head)
- self.k_cache = torch.zeros(cache_shape, dtype=dtype, device=self.key.weight.device)
- self.v_cache = torch.zeros(cache_shape, dtype=dtype, device=self.value.weight.device)
-
- def merge_linears(self, layers, mults):
- bias = [x.bias for x in layers if x.bias is not None][0]
- din, dout = layers[0].weight.shape
- new = nn.Linear(din, len(layers) * dout).to(layers[0].weight.device)
- with torch.no_grad():
- new.weight[:] = torch.cat([x.weight * m for x,m in zip(layers, mults)])
- new.bias[:] = torch.cat([torch.zeros_like(bias) if x.bias is None else x.bias * m for x, m in zip(layers, mults)])
- return new
-
- def convert_for_eval(self):
- if self.qkv or self.kv: raise AttributeError("already converted")
- 
- self.odim = self.key.weight.shape[1]
- if self.cross:
- self.q = self.merge_linears([self.query], [self.sqrt_qk_scale])
- self.kv = self.merge_linears([self.key, self.value],
- [self.sqrt_qk_scale, 1])
- else:
- self.qkv = self.merge_linears([self.query, self.key, self.value],
- [self.sqrt_qk_scale, self.sqrt_qk_scale, 1])
- 
- def split_heads(self, x, x_positions, rope=False, subsampling=1):
- x = x.view(*x.shape[:2], self.n_head, -1)
- if rope:
- x = rope_rotate(x, x_positions * subsampling, *self.rotary(x))
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- qx,
- q_positions,
- kvx,
- kv_positions,
- causal = False,
- mask=None,
- ):
- if self.qkv:
- q,k,v = self.qkv(qx).split(self.odim, dim=-1)
- elif self.kv:
- q = self.q(qx)
- k,v = self.kv(kvx).split(self.odim, dim=-1)
- else:
- q,k,v = None,None,None
- 
- if q is None: q = self.query(qx) * self.sqrt_qk_scale
- q = self.split_heads(q, q_positions, rope = self.rotary, subsampling = self.query_subsampling)
-
- if kvx is not self.cached_kvx:
- if k is None: k = self.key(kvx) * self.sqrt_qk_scale
- k = self.split_heads(k, kv_positions, rope = self.rotary, subsampling = self.key_subsampling)
- if v is None: v = self.value(kvx)
- v = self.split_heads(v, kv_positions)
- if self.k_cache is not None:
- self.k_cache[:,:,kv_positions] = k
- self.v_cache[:,:,kv_positions] = v
-
- if self.k_cache is not None:
- k, v = self.k_cache, self.v_cache
-
- if mask is not None:
- mask = mask[q_positions]
- 
- wv = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0, is_causal=causal)
- 
- return self.out(wv.permute(0, 2, 1, 3).flatten(start_dim=2))
-
-# %% ../nbs/A. Neural modules.ipynb 6
-# modified from https://blog.eleuther.ai/rotary-embeddings/
-
-import torch
-
-class Rotary(torch.nn.Module):
- def __init__(self, dim, base=10000):
- super().__init__()
- inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
- self.register_buffer("inv_freq", inv_freq)
- self.seq_len_cached = None
- self.cos_cached = None
- self.sin_cached = None
-
- def forward(self, x, seq_dim=1):
- seq_len = x.shape[seq_dim]
- if not self.seq_len_cached or seq_len > self.seq_len_cached:
- self.seq_len_cached = 2500
- # self.seq_len_cached = seq_len
- 
- t = torch.arange(self.seq_len_cached, device=x.device).type_as(self.inv_freq)
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
- emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
- self.cos_cached = emb.cos()[None, :, None, :]
- self.sin_cached = emb.sin()[None, :, None, :]
- return self.cos_cached, self.sin_cached
-
-
-# rotary pos emb helpers:
-def rotate_half(x):
- x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
- return torch.cat(
- (-x2, x1), dim=len(x.shape)-1
- )
-
-def rope_rotate(x, positions, cos, sin):
- return x * cos[:,positions] + rotate_half(x) * sin[:,positions]
-
-# %% ../nbs/A. Neural modules.ipynb 7
-class ResidualAttentionBlock(nn.Module):
- def __init__(self, n_state: int, n_head: int, cross_attention: bool = False, rope: bool = False,
- qk_scale: float = 1, ffn_mult: int = 4):
- super().__init__()
- self.attn = MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope)
- self.attn_ln = LayerNorm(n_state)
-
- self.cross_attn = (
- MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope, cross=True) if cross_attention else None
- )
- self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
-
- n_mlp = n_state * ffn_mult
- self.mlp = nn.Sequential(
- nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state)
- )
- self.mlp_ln = LayerNorm(n_state)
- 
- def setup_kv_cache(self, max_batch_size, max_seq_len, max_cross_seq_len=None):
- self.attn.setup_kv_cache(max_batch_size, max_seq_len)
- if self.cross_attn:
- self.cross_attn.setup_kv_cache(max_batch_size, max_cross_seq_len)
- 
- def forward(
- self,
- x: Tensor,
- x_positions: Tensor = None,
- xa: Optional[Tensor] = None,
- xa_positions: Optional[Tensor] = None,
- causal = False,
- mask=None,
- ):
- lnx = self.attn_ln(x)
- x = x + self.attn(lnx, x_positions, lnx, x_positions, causal=causal, mask=mask)
- if self.cross_attn:
- lnx = self.cross_attn_ln(x)
- x = x + self.cross_attn(lnx, x_positions, xa, xa_positions)
- x = x + self.mlp(self.mlp_ln(x))
- return x
-
-# %% ../nbs/A. Neural modules.ipynb 8
-class BaseDecoder(nn.Module):
- def __init__(self, depth=6, n_head=6, width=384, qk_scale=1, ffn_mult=4, length=2250, rope=False):
- super().__init__()
- self.length = length
- self.width = width
- self.layers = nn.ModuleList([
- ResidualAttentionBlock(
- self.width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, cross_attention=True, rope=rope
- ) for _ in range(math.floor(depth))
- ])
-
- self.ln_post = LayerNorm(width)
- 
- mask = torch.empty(length, length).fill_(-torch.inf).triu_(1)
- self.register_buffer("mask", mask, persistent=False)
-
- def forward(self, x, x_positions, xenc, xenc_positions):
- for i,l in enumerate(self.layers):
- x = l(x, x_positions, xenc, xenc_positions, causal=False, mask=self.mask)
-
- x = self.ln_post(x)
-
- return x
-
-# %% ../nbs/A. Neural modules.ipynb 9
-class EmbeddingProjector(nn.Linear):
- pass
-
-class FlexEmbeddings(nn.Module):
- def __init__(self, codes, width, special_codes=None, frozen_width=None, special_embedding=None, unembed=True):
- super().__init__()
- self.codes = codes
- self.special_codes = special_codes
- if frozen_width is None: frozen_width = width
- 
- self.main = nn.Embedding(codes, frozen_width or width)
- self.emb_to_hidden = EmbeddingProjector(frozen_width, width) if frozen_width != width else None
- self.hidden_to_emb = EmbeddingProjector(width, frozen_width) if unembed and frozen_width != width else None
- if special_codes:
- self.special = special_embedding or nn.Embedding(special_codes, width)
- 
- self.register_buffer('merged_in', None)
- self.register_buffer('merged_out', None)
- self.register_buffer('bias_out', None)
- 
- def set_frozen_embeddings(self, values):
- with torch.no_grad():
- self.main.weight[:] = values
- self.main.lr_scale = 0
- 
- @torch.no_grad()
- def convert_for_eval(self):
- if not self.special_codes: return
- # in
- main_w = self.main.weight
- if self.emb_to_hidden is not None: main_w = self.emb_to_hidden(main_w)
- weight = torch.cat([main_w, self.special.weight], dim=0)
- self.merged_in = nn.Embedding(*weight.shape, _weight=weight)
- 
- # out
- weight = self.main.weight
- if self.hidden_to_emb: weight = weight @ self.hidden_to_emb.weight
- self.merged_out = torch.cat([weight.T, self.special.weight.T], dim=1).T.contiguous() # T is for F.linear
- if self.hidden_to_emb:
- self.bias_out = torch.cat([
- self.hidden_to_emb.bias @ self.main.weight.T,
- torch.zeros(self.special.weight.shape[0], device=weight.device, dtype=weight.dtype)
- ], dim=0)
- else:
- self.bias_out = None
-
- def forward(self, toks):
- if not self.training and self.merged_in is not None:
- return self.merged_in(toks)
- 
- if self.special_codes:
- special_mask = toks >= self.codes
- embs = self.main(torch.where(special_mask, 0, toks))
- else:
- embs = self.main(toks)
- 
- if self.emb_to_hidden: embs = self.emb_to_hidden(embs)
- 
- if self.special_codes:
- embs[special_mask] = self.special(toks[special_mask] - self.codes).to(embs.dtype)
- 
- return embs
- 
- def unembed(self, embs):
- if not self.training and self.merged_out is not None:
- return F.linear(embs, self.merged_out, self.bias_out) # embs @ self.merged_out + self.bias_out
-
- orig_embs = embs
- if self.hidden_to_emb: embs = self.hidden_to_emb(embs)
- 
- main_logits = (embs @ self.main.weight.to(embs.dtype).T).float()
- 
- if not self.special_codes:
- return main_logits
- 
- special_logits = (orig_embs @ self.special.weight.to(orig_embs.dtype).T).float()
- return torch.cat([main_logits, special_logits], dim=-1)

whisperspeech/pipeline.py

@@ -1,93 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/7. Pipeline.ipynb.
-
-# %% auto 0
-__all__ = ['Pipeline']
-
-# %% ../nbs/7. Pipeline.ipynb 1
-import torch
-from whisperspeech.t2s_up_wds_mlang_enclm import TSARTransformer
-from whisperspeech.s2a_delar_mup_wds_mlang import SADelARTransformer
-from whisperspeech.a2wav import Vocoder
-import traceback
-from pathlib import Path
-
-# %% ../nbs/7. Pipeline.ipynb 2
-class Pipeline:
- default_speaker = torch.tensor(
- [-0.2929, -0.4503, 0.4155, -0.1417, 0.0473, -0.1624, -0.2322, 0.7071,
- 0.4800, 0.5496, 0.0410, 0.6236, 0.4729, 0.0587, 0.2194, -0.0466,
- -0.3036, 0.0497, 0.5028, -0.1703, 0.5039, -0.6464, 0.3857, -0.7350,
- -0.1605, 0.4808, 0.5397, -0.4851, 0.1774, -0.8712, 0.5789, 0.1785,
- -0.1417, 0.3039, 0.4232, -0.0186, 0.2685, 0.6153, -0.3103, -0.5706,
- -0.4494, 0.3394, -0.6184, -0.3617, 1.1041, -0.1178, -0.1885, 0.1997,
- 0.5571, -0.2906, -0.0477, -0.4048, -0.1062, 1.4779, 0.1639, -0.3712,
- -0.1776, -0.0568, -0.6162, 0.0110, -0.0207, -0.1319, -0.3854, 0.7248,
- 0.0343, 0.5724, 0.0670, 0.0486, -0.3813, 0.1738, 0.3017, 1.0502,
- 0.1550, 0.5708, 0.0366, 0.5093, 0.0294, -0.7091, -0.8220, -0.1583,
- -0.2343, 0.1366, 0.7372, -0.0631, 0.1505, 0.4600, -0.1252, -0.5245,
- 0.7523, -0.0386, -0.2587, 1.0066, -0.2037, 0.1617, -0.3800, 0.2790,
- 0.0184, -0.5111, -0.7291, 0.1627, 0.2367, -0.0192, 0.4822, -0.4458,
- 0.1457, -0.5884, 0.1909, 0.2563, -0.2035, -0.0377, 0.7771, 0.2139,
- 0.3801, 0.6047, -0.6043, -0.2563, -0.0726, 0.3856, 0.3217, 0.0823,
- -0.1302, 0.3287, 0.5693, 0.2453, 0.8231, 0.0072, 1.0327, 0.6065,
- -0.0620, -0.5572, 0.5220, 0.2485, 0.1520, 0.0222, -0.2179, -0.7392,
- -0.3855, 0.1822, 0.1042, 0.7133, 0.3583, 0.0606, -0.0424, -0.9189,
- -0.4882, -0.5480, -0.5719, -0.1660, -0.3439, -0.5814, -0.2542, 0.0197,
- 0.4942, 0.0915, -0.0420, -0.0035, 0.5578, 0.1051, -0.0891, 0.2348,
- 0.6876, -0.6685, 0.8215, -0.3692, -0.3150, -0.0462, -0.6806, -0.2661,
- -0.0308, -0.0050, 0.6756, -0.1647, 1.0734, 0.0049, 0.4969, 0.0259,
- -0.8949, 0.0731, 0.0886, 0.3442, -0.1433, -0.6804, 0.2204, 0.1859,
- 0.2702, 0.1699, -0.1443, -0.9614, 0.3261, 0.1718, 0.3545, -0.0686]
- )
- 
- def __init__(self, t2s_ref=None, s2a_ref=None, optimize=True, torch_compile=False):
- args = dict()
- try:
- if t2s_ref:
- args["ref"] = t2s_ref
- self.t2s = TSARTransformer.load_model(**args).cuda()
- if optimize: self.t2s.optimize(torch_compile=torch_compile)
- except:
- print("Failed to load the T2S model:")
- print(traceback.format_exc())
- try:
- if s2a_ref:
- args["ref"] = s2a_ref
- self.s2a = SADelARTransformer.load_model(**args).cuda()
- if optimize: self.s2a.optimize(torch_compile=torch_compile)
- except:
- print("Failed to load the S2A model:")
- print(traceback.format_exc())
- self.vocoder = Vocoder()
- self.encoder = None
-
- def extract_spk_emb(self, fname):
- """Extracts a speaker embedding from the first 30 seconds of the give audio file.
- """
- import torchaudio
- if self.encoder is None:
- from speechbrain.pretrained import EncoderClassifier
- self.encoder = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
- savedir="~/.cache/speechbrain/",
- run_opts={"device": "cuda"})
- samples, sr = torchaudio.load(fname)
- samples = self.encoder.audio_normalizer(samples[0,:30*sr], sr)
- spk_emb = self.encoder.encode_batch(samples)
- return spk_emb[0,0]
- 
- def generate_atoks(self, text, speaker=None, lang='en', cps=15, step_callback=None):
- if speaker is None: speaker = self.default_speaker
- elif isinstance(speaker, (str, Path)): speaker = self.extract_spk_emb(speaker)
- text = text.replace("\n", " ")
- stoks = self.t2s.generate(text, cps=cps, lang=lang, step=step_callback)
- atoks = self.s2a.generate(stoks, speaker.unsqueeze(0), step=step_callback)
- return atoks
- 
- def generate(self, text, speaker=None, lang='en', cps=15, step_callback=None):
- return self.vocoder.decode(self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=step_callback))
- 
- def generate_to_file(self, fname, text, speaker=None, lang='en', cps=15, step_callback=None):
- self.vocoder.decode_to_file(fname, self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=None))
- 
- def generate_to_notebook(self, text, speaker=None, lang='en', cps=15, step_callback=None):
- self.vocoder.decode_to_notebook(self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=None))

whisperspeech/prepare_s2a_dataset.py

@@ -1,112 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4A. S2A dataset preparation.ipynb.
-
-# %% auto 0
-__all__ = ['flac_to_s2a_name']
-
-# %% ../nbs/4A. S2A dataset preparation.ipynb 2
-import sys
-import os
-import itertools
-from pathlib import Path
-
-import numpy as np
-import torch
-import torchaudio
-import torch.nn.functional as F
-from torch.profiler import profile, record_function, ProfilerActivity
-
-from fastprogress import progress_bar
-from fastcore.script import *
-
-import whisper
-from . import vad, wh_transcribe, vq_stoks, extract_acoustic
-import webdataset as wds
-
-# %% ../nbs/4A. S2A dataset preparation.ipynb 4
-def flac_to_s2a_name(input):
- if '-flac-' in input:
- return input.rsplit("/", 1)[1].replace('flac', 's2a') + ".gz"
- else:
- return input.rsplit("/", 1)[1].replace('raw', 's2a') + ".gz"
-
-# %% ../nbs/4A. S2A dataset preparation.ipynb 6
-def resampler(newsr = 24000, key = 'samples_24k'):
- _last_sr = None
- tform = None
- 
- def _resample(samples):
- for s in samples:
- sr = s['sample_rate']
- if sr != newsr:
- if sr != _last_sr: tform = torchaudio.transforms.Resample(sr, newsr)
- s[key] = tform(s['samples'])
- else:
- s[key] = s['samples']
- yield s
- 
- return _resample
-
-# %% ../nbs/4A. S2A dataset preparation.ipynb 9
-@call_parse
-def prepare_s2a(
- input:str, # FLAC webdataset file path (or - to read the names from stdin)
- proc_dataset_path:Path, # processed VAD files path
- output:str=None, # output file name
- vq_model:str="collabora/spear-tts-pytorch:whisper-vq-stoks.model", # the model path (use repo_id:filename to download it from hugginface)
- n_samples:int=None, # process a limited amount of samples
- batch_size:int=1, # process several segments at once
- fix_dots:bool=False, # fix dots in file names
-):
- if ":" in vq_model:
- repo, fname = vq_model.split(":", 1)
- vq_model = vq_stoks.RQBottleneckTransformer.load_model(repo, fname).cuda()
- else:
- vq_model = vq_stoks.RQBottleneckTransformer.load_model(local_filename=vq_model).cuda()
- amodel = extract_acoustic.load_model()
- amodel.set_target_bandwidth(3)
-
- if input == "-":
- input = [f.strip() for f in sys.stdin.readlines()]
- assert output, "please provide the output shard name"
- else:
- if output is None: output = flac_to_s2a_name(input)
- input = [input]
- 
- total = n_samples//batch_size if n_samples else 'noinfer'
-
- ds = wds.WebDataset(input, shardshuffle=True, rename_files=vad.fix_dots_in_names if fix_dots else None).compose(
- wds.decode(wds.torch_audio),
- wds.select(lambda x: 'wav' in x or 'flac' in x),
- vq_stoks.merge_in(vq_stoks.derived_dataset(proc_dataset_path, 'vad')),
- wds.map_dict(**{"vad.npy":wh_transcribe.chunk_merger}),
- lambda x: wh_transcribe.split_to_chunks(x),
- resampler(),
- resampler(16000, 'samples_16k'),
- wds.to_tuple('__key__', 'rpad_s', 'samples_16k', 'samples_24k'),
- wds.batched(64),
- )
-
- dl = wds.WebLoader(ds, num_workers=4, batch_size=None).unbatched().shuffle(2000).batched(batch_size)
-
- speakers = set()
- tmp = output+".tmp"
- with wds.TarWriter(tmp) as sink:
- for keys, rpad_ss, samples, samples24k in progress_bar(dl, total=total):
- with record_function('to_cuda'):
- samples, samples24k = samples.cuda(), samples24k.unsqueeze(1).cuda()
- with record_function('encodec'):
- atoks = amodel.encode(samples24k)[0][0]
- with record_function('vq_stoks'):
- stoks = vq_model.encode_audio(samples)
- with record_function('from_cuda'):
- atoks, stoks = atoks.cpu().numpy().astype(np.int16), stoks.cpu().numpy().astype(np.int16)
- for key, rpad_s, _atoks, _stoks in zip(keys, rpad_ss, atoks, stoks):
- speakers.add(key.split('/')[1])
- sink.write({
- "__key__": key,
- "atoks.npy": _atoks[:,:int(-rpad_s * 75)],
- "stoks.npy": _stoks[:int(-rpad_s * 25)],
- })
- with open(output+".speakers.txt", "w") as f: f.write("\n".join(speakers))
- if not n_samples:
- os.rename(tmp, output)

whisperspeech/prepare_t2s_dataset.py

@@ -1,111 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/5A. T2S dataset preparation.ipynb.
-
-# %% auto 0
-__all__ = []
-
-# %% ../nbs/5A. T2S dataset preparation.ipynb 2
-import sys
-import os
-import itertools
-from pathlib import Path
-
-import numpy as np
-import torch
-import torchaudio
-import torch.nn.functional as F
-from torch.profiler import profile, record_function, ProfilerActivity
-
-from fastprogress import progress_bar
-from fastcore.script import *
-
-import whisper, whisperx
-from . import vad, wh_transcribe, vq_stoks, extract_acoustic
-import webdataset as wds
-
-# %% ../nbs/5A. T2S dataset preparation.ipynb 4
-def flac_to_t2s_name(input):
- return input.rsplit("/", 1)[1].replace('flac', 't2s') + ".gz"
-
-# %% ../nbs/5A. T2S dataset preparation.ipynb 6
-class Transcriber:
- """
- A helper class to transcribe a batch of 30 second audio chunks.
- """
- def __init__(self, model_size, lang=False):
- self.model = whisperx.asr.load_model(model_size, "cuda", compute_type="float16", language=lang)
- # without calling vad_model at least once the rest segfaults for some reason...
- self.model.vad_model({"waveform": torch.zeros(1, 16000), "sample_rate": 16000})
- 
- def transcribe(self, batch):
- batch = whisper.log_mel_spectrogram(batch)
- embs = self.model.model.encode(batch.cpu().numpy())
- return self.model.tokenizer.tokenizer.decode_batch([x.sequences_ids[0] for x in 
- self.model.model.model.generate(
- embs,
- [self.model.model.get_prompt(self.model.tokenizer, [], without_timestamps=True)]*len(batch),
- )])
-
-# %% ../nbs/5A. T2S dataset preparation.ipynb 7
-@call_parse
-def prepare_t2s(
- input:str, # FLAC webdataset file path (or - to read the names from stdin)
- proc_dataset_path:Path, # processed VAD files path
- output:str=None, # output file name
- vq_model:str="collabora/spear-tts-pytorch:whisper-vq-stoks.model", # the model path (use repo_id:filename to download it from hugginface)
- n_samples:int=None, # process a limited amount of samples
- batch_size:int=1, # process several segments at once
- transcription_model:str="small.en",
-):
- if ":" in vq_model:
- repo, fname = vq_model.split(":", 1)
- vq_model = vq_stoks.RQBottleneckTransformer.load_model(repo, fname).cuda()
- else:
- vq_model = vq_stoks.RQBottleneckTransformer.load_model(local_filename=vq_model).cuda()
- transcriber = Transcriber(transcription_model)
- 
- if input == "-":
- input = [f.strip() for f in sys.stdin.readlines()]
- assert output, "please provide the output shard name"
- else:
- if output is None: output = flac_to_t2s_name(input)
- input = [input]
- 
- total = n_samples//batch_size if n_samples else 'noinfer'
- if n_samples: print(f"Benchmarking run of {n_samples} samples ({total} batches)")
-
- ds = wds.WebDataset(input, shardshuffle=True, rename_files=vad.fix_dots_in_names).compose(
- wds.decode(wds.torch_audio),
- vq_stoks.merge_in(vq_stoks.derived_dataset(proc_dataset_path, 'vad')),
- wds.map_dict(**{"vad.npy": lambda s: wh_transcribe.chunk_merger(s, wh_transcribe.random_cutter)}),
- lambda x: wh_transcribe.split_to_chunks(x),
- # drop the first and last segment because they tend to be inaccurate
- # (the transcriptions don't have the "LibriVox" header and "end of chapter" suffix)
- wds.select(lambda x: x['i'] != 0 and x['i'] != x['imax']),
- wds.to_tuple('__key__', 'rpad', 'samples'),
- wds.batched(64),
- )
-
- dl = wds.WebLoader(ds, num_workers=4, batch_size=None).unbatched().shuffle(2000).batched(batch_size)
-
- speakers = set()
- tmp = output+".tmp"
- with wds.TarWriter(tmp) as sink:
- for keys, rpads, samples in progress_bar(dl, total=total):
- with record_function('to_cuda'):
- csamples = samples.cuda()
- with record_function('transcribe'):
- txts = transcriber.transcribe(csamples)
- with record_function('vq_stoks'):
- stoks = vq_model.encode_audio(csamples)
- with record_function('from_cuda'):
- stoks = stoks.cpu().numpy().astype(np.int16)
- for key, rpad, txt, _stoks in zip(keys, rpads, txts, stoks):
- speakers.add(key.split('/')[1])
- sink.write({
- "__key__": key,
- "txt": txt,
- "stoks.npy": _stoks[:int(-rpad/16000 * 25)],
- })
- with open(output+".speakers.txt", "w") as f: f.write("\n".join(speakers))
- if not n_samples:
- os.rename(tmp, output)

whisperspeech/s2a_delar_mup_wds.py

@@ -1,688 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4B. Semantic to acoustic token modeling.ipynb.
-
-# %% auto 0
-__all__ = ['load_datasets', 'CMLMVisual', 'Rotary', 'rotate_half', 'apply_rotary_pos_emb', 'ResidualAttentionBlock',
- 'MultiHeadAttention', 'DelSumDecoder', 'EmbeddingProjector', 'rand', 'Tunables', 'SADelARTransformer']
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 1
-import io
-import time
-import math
-import random
-import dataclasses
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 2
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.profiler import profile, record_function, ProfilerActivity, schedule
-from fastcore.basics import store_attr
-from huggingface_hub import hf_hub_download
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 3
-from pathlib import Path
-import json
-from fastprogress import progress_bar, master_bar
-import webdataset as wds
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 4
-from .train import *
-from .modules import *
-from . import vq_stoks
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 8
-def rand(start, end):
- return random.random() * (end - start) + start
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 9
-def random_trunc(random_trunc_p, atoks_len = 2250, stoks_len = 750):
- atoks_per_second = atoks_len / 30
- def _trunc(samples):
- for s in samples:
- if random.random() < random_trunc_p:
- seconds = rand(0.3, 30)
- s['atoks.npy'] = s['atoks.npy'][:,:math.ceil(seconds * atoks_per_second)]
- s['stoks.npy'] = s['stoks.npy'][:math.ceil(s['atoks.npy'].shape[-1]/atoks_len*stoks_len)]
- yield s
- return _trunc
-
-def pad_samples(atoks_len = 2250, stoks_len = 750, stoks_pad_token = 4096):
- def _pad(samples):
- for s in samples:
- s['stoks.npy'] = F.pad(torch.tensor(s['stoks.npy']), (0, stoks_len - s['stoks.npy'].shape[-1]), value=stoks_pad_token)
- s['atoks.npy'] = F.pad(torch.tensor(s['atoks.npy']), (0, atoks_len - s['atoks.npy'].shape[-1]), value=-100)
- yield s
- return _pad
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 10
-def speaker_id_extractor(speaker_map):
- def _extractor(samples):
- for s in samples:
- s['speaker'] = torch.tensor(speaker_map[s['__key__'].split("/")[1]])
- yield s
- return _extractor
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 14
-def load_datasets(
- input:str, # webdataset folder
- samples:int, # samples per epoch
- subsample:float=1, # use a fraction of the files
- val_samples:int=512,
- random_trunc_p:float=0,# probability of truncating the input to less than 30 seconds
- stoks_pad_token=4096,
- ):
-
- if isinstance(input, (Path, str)):
- path = Path(input)
- if path.is_dir():
- glob = '*-s2a-*.tar.gz'
- else:
- glob = path.name
- path = path.parent
- input = Path(path).glob(glob)
- elif isinstance(input, list):
- pass
- else:
- raise ArgumentError("input should be either a list or a path with an optional glob specifier")
- shards = [str(x) for x in input]
-
- speakers = set()
- for shard in shards:
- with open(shard+'.speakers.txt') as f: speakers = speakers.union(set(x.strip() for x in f.readlines()))
- speakers = {id:i for i,id in enumerate(sorted(speakers))}
-
- def ds(shards, length):
- ds = wds.WebDataset(wds.ResampledShards(shards)).compose(
- wds.decode(),
- speaker_id_extractor(speakers),
- random_trunc(random_trunc_p) if random_trunc_p > 0 else lambda x: x,
- pad_samples(stoks_pad_token=stoks_pad_token),
- wds.to_tuple('stoks.npy', 'atoks.npy', 'speaker'),
- wds.batched(64),
- )
- ds.speakers = speakers
- ds.total_samples = length
- return ds.compose(wds.slice(length // 64)).with_epoch(length // 64).with_length(length // 64)
- 
- return (
- ds(shards[1:], samples),
- ds(shards[:1], val_samples),
- )
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 33
-import pylab as plt
-import fastprogress
-import IPython
-import numpy as np
-
-class CMLMVisual:
- """Visualize training progress"""
- def __init__ (self, model, masterbar, total_steps):
- self.model = model
- self.masterbar = masterbar
- self.total_steps = total_steps
- self.epochs = total_steps // masterbar.main_bar.total
- 
- gs = plt.GridSpec(3, 1, height_ratios=[2,2,1])
- graph_fig = plt.figure(figsize=(10,6))
- self.graph_fig = graph_fig
- self.loss_p = graph_fig.add_subplot(gs[0])
- self.acc_p = graph_fig.add_subplot(gs[1], sharex=self.loss_p)
- self.acc_p.tick_params('x', labelbottom=False)
- self.lr_p = graph_fig.add_subplot(gs[2], sharex=self.loss_p)
- self.lr_p.tick_params('x', labelbottom=False)
- self.graph_out = None
- 
- self.its = []
- self.train_losses = []
- self.val_losses = []
- self.lr_history = []
- self.acc = np.nan
- self.acc_history = []
- self.pacc_history = []
- 
- def show(self):
- self.start_t = time.time()
- self.masterbar.write(["samples", "train", "val", "time"], table=True)
- self.graph_out = display(self.graph_fig, display_id=True)
- self.acc_out = display(IPython.display.HTML(''), display_id=True)
- 
- def hide(self):
- if self.graph_out is not None:
- self.graph_out.update(IPython.display.HTML(''))
- 
- def plot(self):
- loss_p, acc_p, lr_p = self.loss_p, self.acc_p, self.lr_p
- loss_p.clear()
- loss_p.plot(self.its, self.train_losses)
- loss_p.plot(self.its, self.val_losses)
- loss_p.set_xlim(0, self.total_steps)
- loss_p.set_yscale('log')
- acc_p.clear()
- for k in self.acc_history[-1].keys():
- acc_p.plot(self.its, [x[k] for x in self.acc_history], ':')
-# acc_p.plot(self.its, np.stack(self.pacc_history), label=range(len(self.pacc_history[0])))
- lr_p.clear()
- lrs = np.array(self.lr_history)
- lr_p.plot(self.its, lrs)
- self.graph_out.update(self.graph_fig)
- 
- def add_data(self, it, lr, train_loss, val_los):
- self.its.append(it)
- self.train_losses.append(train_loss)
- self.val_losses.append(val_los)
- self.lr_history.append(lr)
- metrics = self.model.get_metrics()
- self.acc_history.append(metrics)
-# self.acc_out.update(f"Accuracy: {self.entropy_history[-1]:.2f}")
-# self.pacc_history.append((self.model.pval_true / self.model.pval_total).cpu().numpy())
-# if self.acc_history:
- html = "<h5>Accuracies:</h5><table>"
- html += "<thead>"+(''.join([f"<td>{k}<td>" for k,x in metrics.items()]))+"</thead>"
- html += "<tr>"+(''.join([f"<td>{x*100:.1f}%<td>" for k,x in metrics.items()]))+"</tr>"
- html += "</table>"
- self.acc_out.update(IPython.display.HTML(html))
- self.plot()
-
- def add_table_row(self, it, avg_train_loss, val_loss):
- elapsed_t = time.time() - self.start_t
- self.masterbar.write([it, f"{avg_train_loss:.5f}", f"{val_loss:.5f}", fastprogress.core.format_time(elapsed_t)], table=True)
- 
- def on_iter(self, bar, it, avg_train_loss, val_loss):
- epoch = math.ceil(it / self.total_steps * self.epochs)
- bar.comment = f"#{epoch}/{self.epochs} loss: {avg_train_loss:.3f} / {val_loss:.3f}"
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 34
-# modified from https://blog.eleuther.ai/rotary-embeddings/
-import torch
-
-class Rotary(torch.nn.Module):
- def __init__(self, dim, base=10000):
- super().__init__()
- inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
- self.register_buffer("inv_freq", inv_freq)
- self.seq_len_cached = None
- self.cos_cached = None
- self.sin_cached = None
-
- def forward(self, x, seq_dim=1):
- seq_len = x.shape[seq_dim]
- if seq_len != self.seq_len_cached:
- self.seq_len_cached = seq_len
- t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
- emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
- self.cos_cached = emb.cos()[None, :, None, :]
- self.sin_cached = emb.sin()[None, :, None, :]
- return self.cos_cached, self.sin_cached
-
-
-# rotary pos emb helpers:
-def rotate_half(x):
- x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
- return torch.cat(
- (-x2, x1), dim=-1
- )
-
-#@torch.jit.script
-def apply_rotary_pos_emb(q, k, cos, sin):
- return (q * cos[:,:q.shape[1]]) + (rotate_half(q) * sin[:,:q.shape[1]]), (k * cos) + (rotate_half(k) * sin)
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 35
-from torch import Tensor, nn
-import torch.nn.functional as F
-from typing import Dict, Iterable, Optional
-
-class ResidualAttentionBlock(nn.Module):
- def __init__(self, n_state: int, n_head: int, cross_attention: bool = False, rope: bool = False,
- qk_scale: float = 1, ffn_mult: int = 4):
- super().__init__()
-
- self.attn = MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope)
- self.attn_ln = LayerNorm(n_state)
-
- self.cross_attn = (
- MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope) if cross_attention else None
- )
- self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
-
- n_mlp = n_state * ffn_mult
- self.mlp = nn.Sequential(
- nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state)
- )
- self.mlp_ln = LayerNorm(n_state)
- 
- def forward(
- self,
- x: Tensor,
- xa: Optional[Tensor] = None,
- causal = False,
- kv_cache: Optional[dict] = None,
- ):
- x = x + self.attn(self.attn_ln(x), causal=causal, kv_cache=kv_cache)[0]
- if self.cross_attn:
- x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
- x = x + self.mlp(self.mlp_ln(x))
- return x
- 
-class MultiHeadAttention(nn.Module):
- def __init__(self, n_state: int, n_head: int, qk_scale: float = 1, rope: bool = False):
- super().__init__()
- self.n_head = n_head
- self.sqrt_qk_scale = math.sqrt(qk_scale)
- self.query = QueryHead(n_state, n_state)
- self.key = nn.Linear(n_state, n_state, bias=False)
- self.value = nn.Linear(n_state, n_state)
- self.out = nn.Linear(n_state, n_state)
- 
- self.rotary = None
- if rope:
- self.rotary = Rotary(n_state // n_head)
-
- def forward(
- self,
- x: Tensor,
- xa: Optional[Tensor] = None,
- causal = False,
- kv_cache: Optional[dict] = None,
- ):
- q = self.query(x)
-
- if kv_cache is None or xa is None or self.key not in kv_cache:
- # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
- # otherwise, perform key/value projections for self- or cross-attention as usual.
- k = self.key(x if xa is None else xa)
- v = self.value(x if xa is None else xa)
- else:
- # for cross-attention, calculate keys and values once and reuse in subsequent calls.
- k = kv_cache[self.key]
- v = kv_cache[self.value]
-
- if self.sqrt_qk_scale != 1:
- q *= self.sqrt_qk_scale
- k *= self.sqrt_qk_scale
- 
- wv, qk = self.qkv_attention_pth20(q, k, v, causal)
-# wv, qk = self.qkv_attention_xformers(q, k, v, causal)
- 
- return self.out(wv), qk
-
- def qkv_attention_pth20(
- self, q: Tensor, k: Tensor, v: Tensor, causal = False
- ):
- n_batch, n_ctx, n_state = q.shape
- q = q.view(*q.shape[:2], self.n_head, -1)
- k = k.view(*k.shape[:2], self.n_head, -1)
- v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
- 
- #print('before rot:', q.shape, k.shape)
- if self.rotary:
- q, k = apply_rotary_pos_emb(q, k, *self.rotary(k))
- #print(' after rot:', q.shape, k.shape)
-
- k = k.permute(0, 2, 1, 3)
- q = q.permute(0, 2, 1, 3)
- # modified for better performance under PyTorch 2.0
- wv = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0, is_causal=causal)
-
- # previously we've returned q@k which we don't have now
- # since it's not actually used anywhere else, let's just keep two return values for compatibility
- return wv.permute(0, 2, 1, 3).flatten(start_dim=2), None
-
- def qkv_attention_xformers(
- self, q: Tensor, k: Tensor, v: Tensor, causal = False
- ):
- n_batch, n_ctx, n_state = q.shape
- q = q.view(*q.shape[:2], self.n_head, -1)
- k = k.view(*k.shape[:2], self.n_head, -1)
- v = v.view(*v.shape[:2], self.n_head, -1)
-
- if self.rotary:
- q, k = apply_rotary_pos_emb(q, k, *self.rotary(k))
-
- bias = xops.LowerTriangularMask() if causal else None
- wv = xops.memory_efficient_attention(q,k,v, attn_bias=bias)
-
- # previously we've returned q@k which we don't have now
- # since it's not actually used anywhere else, let's just keep two return values for compatibility
- return wv.flatten(start_dim=2), None
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 36
-class DelSumDecoder(nn.Module):
- def __init__(self, depth=6, n_head=6, head_width=64, qk_scale=1, ffn_mult=4, length=2250, codes=1024, quantizers=8, linear_heads=True, rope=False, pos_embs=None):
- super().__init__()
- self.length = length
- width = n_head * head_width
- self.width = width
- self.codes = codes
- self.quantizers = quantizers
- self.linear_heads = linear_heads
-
- self.embeddings = nn.ModuleList([nn.Embedding(codes+1, width) for _ in range(quantizers)])
- if pos_embs is not None:
- self.register_buffer("positional_embedding", pos_embs)
- 
- self.layers = nn.ModuleList([
- ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, cross_attention=True, rope=rope) for _ in range(math.floor(depth))
- ])
-
- self.ln_post = LayerNorm(width)
-
- if self.linear_heads:
- self.heads = LinearHead(width, (codes+1) * quantizers, bias=False)
- else:
- self.splitter = nn.Sequential(
- nn.Linear(width, width * quantizers),
- nn.GELU(),
- )
- self.heads = nn.ModuleList([
- LinearHead(width, codes+1, bias=True) for _ in range(quantizers)
- ])
-
- def forward(self, toks, xenc):
- b,_,n = toks.shape
- newn = min(n+1, self.length)
- embs = torch.zeros((b,newn,self.width), dtype=xenc.dtype, device=xenc.device)
- for i in range(self.quantizers):
- embs[:,:i+1] += self.embeddings[i](torch.tensor([self.codes], device=xenc.device))
- if i < n:
- embs[:,i+1:] += self.embeddings[i](toks[:,i,:newn-i-1])
-
- x = embs.to(xenc.dtype)
- 
- for l in self.layers:
- x = l(x, xenc, causal=True)
- x = self.ln_post(x)
-
- if self.linear_heads:
- logits = self.heads(x).view(b,newn,self.quantizers,self.codes+1).permute(0,2,1,3)
- else:
- split = self.splitter(x).view(b,newn,self.quantizers,self.width)
- logits = torch.stack([self.heads[q](split[:,:,q]) for q in range(self.quantizers)], dim=1)
-
- return logits
- 
-class EmbeddingProjector(nn.Linear):
- pass
-
-def rand(start, end):
- return random.random() * (end - start) + start
- 
-@dataclasses.dataclass
-class Tunables:
- init_std :float = 9
- embeddings_std :float = 0.2
- embeddings_lr_scale: float = 10
- output_mult :float = 5.6
- # FIXME: try separate mults for self and cross attention
- query_mult :float = .3
- encoder_depth_ratio :float = 0.25
- linear_heads :bool = False
- rope :bool = True
- 
- lr0 :float = 3e-3
- clip_gradient_norm :float = 2
- weight_decay :float = 1e-3
- warmup_steps :float = 2000
-
- random :bool = False
-
- def __post_init__(self):
- # randomize the hyperparams if requested
- if self.random:
- self.init_std = 2*10**rand(0,1)
- self.embeddings_std = 10**rand(-1.7,-0.22)
- self.embeddings_lr_scale = 2**rand(2,4)
- self.output_mult = 2**rand(1.5,3)
- self.query_mult = 2**rand(-3,-1.3)
- self.encoder_depth_ratio = random.choice([0.25,0.5])
- self.linear_heads = False
- self.rope = True
- 
- self.lr0 = 3e-3
- self.clip_gradient_norm = 10**rand(-1,1)
- self.warmup_steps = 100*(10**rand(1.18,1.3))
- 
- @staticmethod
- def upgrade(args):
- args = {k:v for k,v in args.items()}
- def old_default(name, value):
- if name not in args: args[name] = value
- old_default('rope', False)
- old_default('linear_heads', True)
- return args
- 
-class SADelARTransformer(nn.Module):
- def __init__(self, depth=3, ctx_n=2250, stoks_len=750, stoks_codes=4097, stoks_width=None, spk_width=None, n_head=3, head_width=64, ffn_mult=4,
- quantizers=8, speaker_map={"1":0}, tunables=Tunables()):
- super().__init__()
- self.quantizers = quantizers
- width = n_head * head_width
- store_attr("depth,ctx_n,stoks_len,stoks_codes,stoks_width,spk_width,n_head,head_width,ffn_mult,quantizers,speaker_map")
- self.width = width
- self.base_width = 3 * head_width
- self.tunables = tunables
- 
- if stoks_width is None: stoks_width = width
- if spk_width is None: spk_width = width
- self.emb_factor = width != stoks_width
- self.spk_factor = width != spk_width
-
- if tunables.rope:
- self.positional_embeddings = None
- else:
- self.register_buffer('positional_embeddings', sinusoids(ctx_n, width))
- 
- self.speaker_embedding = nn.Embedding(len(speaker_map), width)
- self.semantic_embedding = nn.Embedding(stoks_codes, stoks_width)
- if self.emb_factor:
- self.emb_to_hidden = nn.Linear(stoks_width, width)
- 
- if self.spk_factor:
- self.spk_to_hidden = EmbeddingProjector(spk_width, width)
-
- qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)
- 
- encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
- decoder_depth = depth * 2 - encoder_depth
- self.encoder = nn.Sequential(*[
- ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(encoder_depth)
- ])
- self.ln_post = LayerNorm(width)
- 
- self.decoder = DelSumDecoder(pos_embs=self.positional_embeddings, qk_scale=qk_scale,
- length=ctx_n, n_head=n_head, head_width=head_width, ffn_mult=ffn_mult,
- depth=decoder_depth, quantizers=quantizers,
- linear_heads=tunables.linear_heads, rope=tunables.rope)
-
- self.register_buffer('val_true', torch.zeros(self.quantizers).cuda())
- self.register_buffer('val_total', torch.zeros(self.quantizers).cuda())
- self.apply(self.init_transformer)
-
- def setup(self, device):
- pass
- 
- def load_frozen_semantic_embeddings(self, vqmodel):
- with torch.no_grad():
- self.semantic_embedding.weight[:] = vqmodel.rq.layers[0]._codebook.embed[0]
- self.semantic_embedding.lr_scale = 0
- 
- def init_transformer(self, m):
- if isinstance(m, LinearHead):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, QueryHead):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, nn.Embedding):
- m.no_weight_decay = True
- m.lr_scale = self.tunables.embeddings_lr_scale
- std = self.tunables.embeddings_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, EmbeddingProjector):
- m.lr_scale = self.tunables.embeddings_lr_scale/2
- std = self.tunables.init_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.Linear):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- std = self.tunables.init_std / m.weight.shape[1]
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- if m.bias is not None:
- torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.LayerNorm):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.bias, 0)
- torch.nn.init.constant_(m.weight, 1)
-
- def embed_stoks(self, Stoks):
- b,n = Stoks.shape
- if self.stoks_len == 1500:
- # converts 50 toks/s to 75 toks/s by adding padding between every two tokens
- x = Stoks.reshape(b,n//2,2)
- x = x.repeat_interleave(2, -1)[:,:,:3]
- x[:,:,1] = 1024
- x = x.reshape(b,n//2*3)
- else:
- # it's a lot easier with 25 toks/s
- x = Stoks.repeat_interleave(3, -1)
- # embed semantic tokens
- Sembs = self.semantic_embedding(x.to(torch.long))
- if self.emb_factor:
- Sembs = self.emb_to_hidden(Sembs)
- return Sembs
-
- def forward(self, Stoks, Atoks, speakers, noloss=False):
- Atoks = Atoks.to(torch.long)
- semb = self.embed_stoks(Stoks)
- with record_function("encoder"):
- if self.positional_embeddings is not None: semb = semb + self.positional_embeddings
- xenc = self.ln_post(self.encoder(semb))
-# xenc = torch.zeros_like(xenc)
- with record_function("decoder"):
- Atoks_gt = Atoks.clone()
- Atoks_gt[Atoks == -100] = 1024
- # we can randomize speaker ids during validation to measure
- # the importance of the speaker embedding vs. just the acoustic prompt/prefix
-# if not self.training: speakers = speakers[torch.randperm(speakers.nelement())]
- spk_embs = self.speaker_embedding(speakers)
- if self.spk_factor: spk_embs = self.spk_to_hidden(spk_embs)
- logits = self.decoder(Atoks_gt, xenc + spk_embs.unsqueeze(1))
- logits *= self.tunables.output_mult / (self.width / self.base_width)
- 
- if noloss:
- return logits
-
- with record_function("loss"):
- N = Atoks.shape[-1]
- loss = 0
- for i in range(self.quantizers):
- loss += F.cross_entropy(logits[:,i,i:].reshape(-1,logits.shape[-1]), Atoks[:,i,:N-i].reshape(-1))
- loss /= self.quantizers
-
- if not self.training:
- for i in range(self.quantizers):
- Atoks_i = Atoks[:,i,:N-i]
- valid_Atoks = Atoks_i != -100
- self.val_true[i] += (logits[:,i,i:].argmax(-1)[valid_Atoks] == Atoks_i[valid_Atoks]).float().sum()
- self.val_total[i] += valid_Atoks.float().sum()
-
- return logits, loss
-
- def get_metrics(self):
- metrics = {
- f'acc_{i}':x.item() for i,x in enumerate(self.val_true / self.val_total)
- }
- self.val_true[:] = 0
- self.val_total[:] = 0
- return metrics
-
- #
- # inference
- #
- @classmethod
- def load_model(cls, repo_id="collabora/whisperspeech", filename="s2a_up_wds.model", local_filename=None):
- if not local_filename:
- local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
- spec = torch.load(local_filename)
- if '_extra_state' not in spec['state_dict']: spec['state_dict']['_extra_state'] = { 'speaker_map': spec['config']['speaker_map'] }
- model = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec['tunables'])))
- model.load_state_dict(spec['state_dict'])
- model.eval()
- return model
- 
- def get_extra_state(self):
- return { 'speaker_map': self.speaker_map }
- 
- def set_extra_state(self, st):
- self.speaker_map = st['speaker_map']
-
- def load_checkpoint(self, local_filename):
- spec = torch.load(local_filename, map_location='cpu')
- assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'
- state_dict = {k.replace('model.', ''):v
- for k,v in spec['state_dict'].items()}
- self.load_state_dict(state_dict)
- return self
- 
- def save_model(self, fname):
- torch.save(dict(config = self.__stored_args__,
- tunables = dataclasses.asdict(self.tunables),
- state_dict = self.state_dict()), fname)
-
- @property
- def device(self):
- return next(self.parameters()).device
- 
- @torch.no_grad()
- def generate(self, stoks, speakers, N=None, T=0.7, top_k=None, show_progress_bar=True):
- dev = self.device
- if self.stoks_len == 1500:
- N = N or len(stoks) * 3 // 2
- else:
- N = N or len(stoks) * 3
- stoks = F.pad(stoks.to(dev), (0, self.stoks_len - len(stoks)), value=self.stoks_codes-1).unsqueeze(0)
- speakers = torch.tensor([self.speaker_map[spk] for spk in speakers], device=dev)
- toks = torch.zeros((1,self.quantizers,N), dtype=torch.long, device=dev)
- it = range(0,N)
- if show_progress_bar: it = progress_bar(it)
- for i in it:
- p = self(stoks, toks[:,:,:i], speakers, noloss=True)
- last_p = p[0,:,-1]
- if top_k:
- last_p[last_p < torch.topk(last_p, top_k).values[:,-1,None]] = -torch.inf
- for j,tok in enumerate(torch.multinomial((last_p / float(T)).softmax(-1), 1)):
- toks[0,j,max(0,i-j)] = tok
- if toks[0,0,i] == 1024: return toks[0,:,:i]
- return toks[0]
-
-# %% ../nbs/4B. Semantic to acoustic token modeling.ipynb 37
-def _make_model(size:str, quantizers:int=4, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None, **kwargs):
- assert(dataset is not None)
- kwargs = dict(speaker_map=dataset.speakers, quantizers=quantizers, tunables=tunables, **kwargs)
- if size == 'micro':
- return SADelARTransformer(depth=4, n_head=3, ffn_mult=2, **kwargs)
- if size == 'tiny-narrow':
- return SADelARTransformer(depth=4, n_head=6, ffn_mult=1, **kwargs)
- if size == 'tiny':
- return SADelARTransformer(depth=4, n_head=6, **kwargs)
- if size == 'base':
- return SADelARTransformer(depth=6, n_head=8, **kwargs)
- if size == 'base-deep':
- return SADelARTransformer(depth=9, n_head=8, **kwargs)
- if size == 'base-wide':
- return SADelARTransformer(depth=6, n_head=12, **kwargs)
- if size == 'small/2':
- return SADelARTransformer(depth=9, n_head=12, **kwargs)
- if size == 'small':
- return SADelARTransformer(depth=12, n_head=12, **kwargs)
- if size == 'medium':
- return SADelARTransformer(depth=24, n_head=16, **kwargs)
-
-def make_model(size:str, quantizers:int=4, frozen_embeddings_model:str=None, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
- if frozen_embeddings_model:
- vqmodel = vq_stoks.RQBottleneckTransformer.load_model(frozen_embeddings_model)
- model = _make_model(size, quantizers, tunables, dataset, stoks_codes=vqmodel.vq_codes+1, stoks_width=vqmodel.rq.layers[0]._codebook.embed[0].shape[-1])
- model.load_frozen_semantic_embeddings(vqmodel)
- else:
- model = _make_model(size, quantizers, tunables, dataset)
- return model

whisperspeech/s2a_delar_mup_wds_mlang.py

@@ -1,564 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb.
-
-# %% auto 0
-__all__ = ['load_dataset', 'DelSumEmbedding', 'DelSumHead', 'rand', 'Tunables', 'SADelARTransformer']
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 1
-import io
-import time
-import math
-import random
-import dataclasses
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 2
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from torch.profiler import profile, record_function, ProfilerActivity, schedule
-from fastcore.basics import store_attr
-from huggingface_hub import hf_hub_download
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 3
-from pathlib import Path
-import json
-from fastprogress import progress_bar, master_bar
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 4
-from .modules import *
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 8
-def rand(start, end):
- return random.random() * (end - start) + start
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 9
-def random_trunc(random_trunc_p, atoks_len = 2250, stoks_len = 750):
- atoks_per_second = atoks_len / 30
- def _trunc(samples):
- for s in samples:
- if random.random() < random_trunc_p:
- seconds = rand(0.3, 30)
- s['atoks.npy'] = s['atoks.npy'][:,:math.ceil(seconds * atoks_per_second)]
- s['stoks.npy'] = s['stoks.npy'][:math.ceil(s['atoks.npy'].shape[-1]/atoks_len*stoks_len)]
- yield s
- return _trunc
-
-def pad_samples(atoks_len = 2250, stoks_len = 750, stoks_pad_token = 4096):
- def _pad(samples):
- for s in samples:
- s['stoks.npy'] = F.pad(torch.tensor(s['stoks.npy']), (1, stoks_len - s['stoks.npy'].shape[-1]-1), value=stoks_pad_token)
- s['out_stoks'] = F.pad(torch.tensor(s['stoks.npy']), (0, stoks_len - s['stoks.npy'].shape[-1]), value=stoks_pad_token)
- s['atoks.npy'] = F.pad(torch.tensor(s['atoks.npy']), (0, atoks_len - s['atoks.npy'].shape[-1]), value=-100)
- yield s
- return _pad
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 10
-def make_speaker_map(shards):
- speakers = set()
- for shard in shards:
- with open(shard+'.speakers.txt') as f: speakers = speakers.union(set(x.strip() for x in f.readlines()))
- return {id:i for i,id in enumerate(sorted(speakers))}
-
-def speaker_id_extractor(speaker_map):
- def _extractor(samples):
- for s in samples:
- s['speaker'] = torch.tensor(speaker_map[s['__key__'].split("/")[1]])
- yield s
- return _extractor
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 27
-def load_dataset(
- atoks_shard_spec:str, # webdataset folder
- stoks_shard_dir:str, # stoks webdataset base dir
- samples:int, # samples per epoch
- random_trunc_p:float=0,# probability of truncating the input to less than 30 seconds
- vq_codes:int=4096,
- language:str='en',
- weight:float=1,
- validation:bool=False,
- exclude_files:str=None,
- randomize_speakers:bool=False,
- ):
- import webdataset as wds
- from whisperspeech import utils
-
- shards = utils.shard_glob(atoks_shard_spec)
- excludes = {x for file in exclude_files.split() for x in utils.readlines(file)} if exclude_files else set()
- 
- def check_for_nan(s):
- if torch.tensor(s['spk_emb.npy']).isnan().any(): print("found NaN:", s['__key__'])
- return s
- 
- def set_language(x):
- x['language'] = language
- return x
- 
- same_on_all_nodes = lambda urls: urls # will only be used for validation
- ds = wds.WebDataset(shards, resampled=not validation, nodesplitter=same_on_all_nodes).compose(
- wds.decode(),
- utils.merge_in(utils.derived_dataset('maxvad-stoks', base='atoks-3kbps', suffix='', dir=stoks_shard_dir)),
- wds.map(check_for_nan),
- wds.select(lambda s: s['__key__'] not in excludes),
- wds.map_dict(**{'spk_emb.npy':np.nan_to_num}), # remove nans from the speaker embedding model
- random_trunc(random_trunc_p) if random_trunc_p > 0 else lambda x: x,
- pad_samples(stoks_pad_token=vq_codes-1),
- wds.map(set_language),
- wds.to_tuple('stoks.npy', 'atoks.npy', 'spk_emb.npy', 'language', 'out_stoks'),
- wds.shuffle(20000, initial=20000),
- wds.batched(64),
- )
- if randomize_speakers:
- rng = np.random.default_rng()
- ds = ds.compose(
- wds.map_tuple(None, None, lambda x: rng.permutation(x), None),
- )
- if validation:
- ds = ds.slice(samples // 64)
- ds.total_samples = samples
- ds.weight = weight
- 
- return ds
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 37
-class DelSumEmbedding(nn.Module):
- def __init__(self, n_head=6, head_width=64, atoks_width=None, length=2250, codes=1024, quantizers=8, pos_embs=None):
- super().__init__()
- self.length = length
- width = n_head * head_width
- if atoks_width is None: atoks_width = width
- self.width = width
- self.quantizers = quantizers
-
- emb = None
- embs = []
- for _ in range(quantizers):
- emb = FlexEmbeddings(codes, width, special_codes=2, frozen_width=atoks_width,
- special_embedding=emb and emb.special)
- embs.append(emb)
- self.embeddings = nn.ModuleList(embs)
- if pos_embs is not None:
- self.register_buffer("positional_embedding", pos_embs)
-
- def forward(self, toks, xenc):
- with record_function("embeddings"):
- b,_,n = toks.shape
- newn = min(n, self.length)
-
- embs = torch.zeros((b,newn,self.width), dtype=xenc.dtype, device=xenc.device)
- for i in range(self.quantizers):
- embs[:, :] += self.embeddings[i](toks[:,i,:])
- 
- x = embs.to(xenc.dtype)
- return x
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 38
-class DelSumHead(nn.Module):
- def __init__(self, quantizers=8, n_head=6, head_width=64):
- super().__init__()
- self.width = n_head * head_width
- self.quantizers = quantizers
- self.splitter = nn.Sequential(
- nn.Linear(self.width, self.width * quantizers),
- nn.GELU(),
- )
-
- def forward(self, x, embeddings=None):
- b, newn, _ = x.shape
- with record_function("splitter"):
- split = self.splitter(x).view(b,newn,self.quantizers,self.width)
- with record_function("unembed"):
- logits = torch.stack([embeddings[q].unembed(split[:,:,q]) for q in range(self.quantizers)], dim=1)
- return logits
- 
-def rand(start, end):
- return random.random() * (end - start) + start
- 
-@dataclasses.dataclass
-class Tunables:
- init_std :float = 9
- embeddings_std :float = 0.2
- embeddings_lr_scale: float = 10
- output_mult :float = 5.6
- # FIXME: try separate mults for self and cross attention
- query_mult :float = .3
- encoder_depth_ratio :float = 0.25
- linear_heads :bool = False
- rope :bool = True
- 
- lr0 :float = 3e-3
- clip_gradient_norm :float = 2
- weight_decay :float = 1e-3
- warmup_steps :float = 2000
-
- random :bool = False
-
- def __post_init__(self):
- # randomize the hyperparams if requested
- if self.random:
- self.init_std = 2*10**rand(0,1)
- self.embeddings_std = 10**rand(-1.7,-0.22)
- self.embeddings_lr_scale = 2**rand(2,4)
- self.output_mult = 2**rand(1.5,3)
- self.query_mult = 2**rand(-3,-1.3)
- self.encoder_depth_ratio = random.choice([0.25,0.5])
- self.linear_heads = False
- self.rope = True
- 
- self.lr0 = 3e-3
- self.clip_gradient_norm = 10**rand(-1,1)
- self.warmup_steps = 100*(10**rand(1.18,1.3))
- 
- @staticmethod
- def upgrade(args):
- args = {k:v for k,v in args.items()}
- def old_default(name, value):
- if name not in args: args[name] = value
- old_default('rope', False)
- old_default('linear_heads', True)
- return args
- 
-class SADelARTransformer(nn.Module):
- def __init__(self, depth=3, ctx_n=2250,
- stoks_len=750, stoks_codes=4097, stoks_width=None,
- spk_width=None,
- atoks_width=None,
- n_head=3, head_width=64, ffn_mult=4,
- quantizers=8, speaker_map={"1":0}, tunables=Tunables()):
- super().__init__()
- self.quantizers = quantizers
- self.codes = 1024
- width = n_head * head_width
- store_attr("depth,ctx_n,stoks_len,stoks_codes,stoks_width,spk_width,atoks_width,n_head,head_width,ffn_mult,quantizers,speaker_map")
- self.width = width
- self.base_width = 3 * head_width
- self.tunables = tunables
- 
- if stoks_width is None: stoks_width = width
- if spk_width is None: spk_width = width
- self.emb_factor = width != stoks_width
- self.spk_factor = width != spk_width
-
- if tunables.rope:
- self.positional_embeddings = None
- else:
- self.register_buffer('positional_embeddings', sinusoids(ctx_n, width))
- 
-# self.speaker_embedding = nn.Embedding(len(speaker_map), spk_width)
- self.semantic_embedding = nn.Embedding(stoks_codes, stoks_width)
- if self.emb_factor:
- self.emb_to_hidden = nn.Linear(stoks_width, width)
- self.hidden_to_emb = nn.Linear(width, stoks_width)
- 
- if self.spk_factor:
- self.spk_to_hidden = nn.Linear(spk_width, width)
-
- qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)
- 
- encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
- decoder_depth = depth * 2 - encoder_depth
- self.encoder = nn.Sequential(*[
- ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(encoder_depth)
- ]) # FIXME: enclm requires causal attention here
- self.ln_post = LayerNorm(width)
-
- self.embds = DelSumEmbedding(
- pos_embs=self.positional_embeddings, length=ctx_n,
- n_head=n_head, head_width=head_width, atoks_width=atoks_width,
- quantizers=quantizers,
- )
- self.decoder = BaseDecoder(qk_scale=qk_scale, length=ctx_n,
- n_head=n_head, width=n_head * head_width, 
- ffn_mult=ffn_mult, depth=decoder_depth,
- rope=tunables.rope)
- self.head = DelSumHead(n_head=n_head, head_width=head_width, quantizers=quantizers)
- for l in self.decoder.layers:
- l.cross_attn.key_subsampling = 3
-# for l in self.encoder:
-# l.attn.key_subsampling = 3
-# l.attn.query_subsampling = 3
- 
- self.register_buffer('val_true', torch.zeros(self.quantizers).cuda())
- self.register_buffer('val_total', torch.zeros(self.quantizers).cuda())
- self.apply(self.init_transformer)
-
- def setup(self, device):
- pass
- 
- def load_frozen_semantic_embeddings(self, vqmodel):
- with torch.no_grad():
- self.semantic_embedding.weight[:] = vqmodel.rq.layers[0]._codebook.embed[0]
- self.semantic_embedding.lr_scale = 0
-
- def load_frozen_acoustic_embeddings(self, amodel):
- for i in range(self.quantizers):
- self.decoder.embeddings[i].set_frozen_embeddings(amodel.quantizer.vq.layers[i].codebook)
- 
- def init_transformer(self, m):
- if isinstance(m, LinearHead):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, QueryHead):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, nn.Embedding):
- m.no_weight_decay = True
- m.lr_scale = self.tunables.embeddings_lr_scale
- std = self.tunables.embeddings_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
-# elif isinstance(m, EmbeddingProjector):
-# m.lr_scale = self.tunables.embeddings_lr_scale #1/(m.weight.shape[1] / self.base_width)
-# m.lr_scale = 2/(m.weight.shape[1] / self.base_width)
-# std = self.tunables.init_std / m.weight.shape[1]
-# torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.Linear):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- std = self.tunables.init_std / m.weight.shape[1]
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- if m.bias is not None:
- torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.LayerNorm):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.bias, 0)
- torch.nn.init.constant_(m.weight, 1)
-
- def embed_stoks(self, Stoks):
- b,n = Stoks.shape
- if self.stoks_len == 1500:
- # converts 50 toks/s to 75 toks/s by adding padding between every two tokens
- x = Stoks.reshape(b,n//2,2)
- x = x.repeat_interleave(2, -1)[:,:,:3]
- x[:,:,1] = 1024
- x = x.reshape(b,n//2*3)
- else:
- # it's a lot easier with 25 toks/s
-# x = Stoks.repeat_interleave(3, -1)
- x = Stoks
- # embed semantic tokens
- Sembs = self.semantic_embedding(x.to(torch.long))
- if self.emb_factor:
- Sembs = self.emb_to_hidden(Sembs)
- return Sembs
-
- def _encoder(self, semb, positions):
- x = semb
- for l in self.encoder: x = l(x, positions)
- return self.ln_post(x)
- 
- def run_encoder(self, Stoks, speakers):
- semb = self.embed_stoks(Stoks)
- with record_function("encoder"):
- if self.positional_embeddings is not None: semb = semb + self.positional_embeddings
- positions = torch.arange(0, semb.shape[1], device=semb.device)
- xenc = self._encoder(semb, positions)
- if self.training:
- enc_logits = (self.hidden_to_emb(xenc) @ self.semantic_embedding.weight.to(xenc.dtype).T).float()
- enc_logits = enc_logits * self.tunables.output_mult / (self.width / self.base_width)
- else:
- enc_logits = None
-# print(xenc.shape, speakers.shape)
- spk_embs = F.normalize(speakers, dim=-1) # use extracted embeddings
- if self.spk_factor: spk_embs = self.spk_to_hidden(spk_embs)
- return xenc + spk_embs.unsqueeze(1), positions, enc_logits
-
- def forward(self, Stoks, Atoks, speakers, langs=None, out_stoks=None, noloss=False, xenc=None, xenc_positions=None, atoks_positions=None):
- if xenc is None:
- Atoks = Atoks.to(torch.long)
- out_stoks = out_stoks.to(torch.long)
- Atoks_gt = Atoks.clone()
- Atoks_gt[Atoks == -100] = 1024
- xenc, enc_logits = self.run_encoder(Stoks, speakers)
- else:
- Atoks_gt = Atoks
- with record_function("decoder"):
- embs = self.embds(Atoks, xenc)
- if atoks_positions is None: atoks_positions = torch.arange(0, embs.shape[1], device=embs.device)
- x = self.decoder(embs, atoks_positions, xenc, xenc_positions)
- logits = self.head(x, embeddings=self.embds.embeddings)
- logits *= self.tunables.output_mult / (self.width / self.base_width)
- 
- if noloss:
- return logits
-
- with record_function("loss"):
- N = Atoks.shape[-1]
- loss = 0
- for i in range(self.quantizers):
- loss += F.cross_entropy(logits[:,i,i:].reshape(-1,logits.shape[-1]), Atoks[:,i,:N-i].reshape(-1))
- if self.training and i == 0:
- loss *= 5
- loss /= self.quantizers
- if self.training:
- loss += 0.1 * F.cross_entropy(enc_logits.transpose(-1,-2), out_stoks)
-
- if not self.training:
- for i in range(self.quantizers):
- Atoks_i = Atoks[:,i,:N-i]
- valid_Atoks = Atoks_i != -100
- self.val_true[i] += (logits[:,i,i:].argmax(-1)[valid_Atoks] == Atoks_i[valid_Atoks]).float().sum()
- self.val_total[i] += valid_Atoks.float().sum()
-
- return logits, loss
-
- def get_metrics(self):
- metrics = {
- f'acc_{i}':x.item() for i,x in enumerate(self.val_true / self.val_total)
- }
- self.val_true[:] = 0
- self.val_total[:] = 0
- return metrics
-
- #
- # inference
- #
- @classmethod
- def load_model(cls, ref="collabora/whisperspeech:s2a-q4-small-en+pl.model",
- repo_id=None, filename=None, local_filename=None):
- if repo_id is None and filename is None and local_filename is None:
- if ":" in ref:
- repo_id, filename = ref.split(":", 1)
- else:
- local_filename = ref
- if not local_filename:
- local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
- spec = torch.load(local_filename)
- if '_extra_state' not in spec['state_dict']: spec['state_dict']['_extra_state'] = { 'speaker_map': spec['config']['speaker_map'] }
- model = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec['tunables'])))
- model.load_state_dict(spec['state_dict'])
- model.eval()
- return model
- 
- def get_extra_state(self):
- return { 'speaker_map': self.speaker_map }
- 
- def set_extra_state(self, st):
- self.speaker_map = st['speaker_map']
-
- def load_checkpoint(self, local_filename):
- spec = torch.load(local_filename, map_location='cpu')
- assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'
- state_dict = {k.replace('model.', ''):v
- for k,v in spec['state_dict'].items()}
- self.load_state_dict(state_dict)
- return self
- 
- def save_model(self, fname):
- torch.save(dict(config = self.__stored_args__,
- tunables = dataclasses.asdict(self.tunables),
- state_dict = self.state_dict()), fname)
-
- def switch_dtypes(self, dtype=torch.float16):
- self.dtype = dtype
- for n,m in self.named_modules():
- # convert every leaf layer apart from the LayerNorms
- if isinstance(m, (nn.Linear, nn.Embedding)):
- m.to(dtype)
- # take care of buffers ([kv]_cache, masks) that are not in the leaf layers
- for bn,b in m.named_buffers(recurse=False):
- setattr(m,bn,b.to(dtype))
-
- def optimize(self, max_batch_size=1, dtype=torch.float16, torch_compile=True):
- for emb in self.embds.embeddings:
- emb.convert_for_eval()
- for l in self.encoder:
- l.attn.convert_for_eval()
- for l in self.decoder.layers:
- l.attn.convert_for_eval()
- l.cross_attn.convert_for_eval()
- l.setup_kv_cache(max_batch_size, self.ctx_n, self.stoks_len)
- self.switch_dtypes(dtype)
- if torch_compile:
- self.generate_next = torch.compile(self.generate_next, mode="reduce-overhead", fullgraph=True)
-
- @property
- def device(self):
- return next(self.parameters()).device
-
- # from https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
- def multinomial_sample_one_no_sync(self, probs_sort): # Does multinomial sampling without a cuda synchronization
- q = torch.empty_like(probs_sort).exponential_(1)
- return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
-
- def logits_to_probs(self, logits, T=1.0, top_k=None):
- logits = logits / max(T, 1e-5)
-
- if top_k is not None:
- v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
- pivot = v.select(-1, -1).unsqueeze(-1)
- logits = torch.where(logits < pivot, -float("Inf"), logits)
- probs = torch.nn.functional.softmax(logits, dim=-1)
- return probs
-
- def sample(self, logits, T=1.0, top_k=None):
- probs = self.logits_to_probs(logits[0,:,-1], T, top_k)
- idx_next = self.multinomial_sample_one_no_sync(probs)
- return idx_next
-
- def generate_one(self, toks, positions, langs, xenc, xenc_positions, T, top_k):
- probs = self(None, toks, None, langs, noloss=True, xenc=xenc, xenc_positions=xenc_positions, atoks_positions=positions)
- return self.sample(probs, T, top_k)
-
- def generate_next(self, *args, **kwargs):
- return self.generate_one(*args, **kwargs)
- 
- @torch.no_grad()
- def generate(self, stoks, speakers, langs=None, N=None, T=0.7, top_k=None, show_progress_bar=True, step=None, subsample_enc=False):
- dev = self.device
- N = N or len(stoks) * 3
- stoks = F.pad(stoks.to(dev), (1, self.stoks_len - len(stoks)-1), value=self.stoks_codes-1).unsqueeze(0)
- speakers = speakers.to(device=dev, dtype=self.dtype)
- toks = torch.full((1,self.quantizers,2250), self.codes+1, dtype=torch.long, device=dev)
- it = range(1,min(N,2250-1))
- if show_progress_bar: it = progress_bar(it)
- with record_function("encode"):
- xenc, xenc_positions, _ = self.run_encoder(stoks, speakers)
- toks_positions = torch.arange(N, device=dev)
- with record_function("prefill"):
- toks[0,0,1] = self.generate_one(toks[:,:,:1], toks_positions[:1], langs, xenc, xenc_positions, T, top_k)[0,0]
- with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
- for i in it:
- with record_function("generate_one"):
- toks[0,:i+1,i+1] = self.generate_next(toks[:,:,i:i+1], toks_positions[i:i+1], langs, xenc, xenc_positions, T, top_k)[:i+1,0]
-
- # for profiling, debugging or early exit
- if step is not None: step()
- # shift tokens
- toks = toks[:,:,1:N]
- for j in range(self.quantizers):
- toks[0, j] = torch.roll(toks[0, j], -j)
- return toks[0]
-
-# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 39
-def _make_model(size:str, quantizers:int=4, tunables:Tunables=Tunables(), **kwargs):
- kwargs = dict(quantizers=quantizers, tunables=tunables, **kwargs)
- if size == 'micro':
- return SADelARTransformer(depth=4, n_head=3, ffn_mult=2, **kwargs)
- if size == 'tiny-narrow':
- return SADelARTransformer(depth=4, n_head=6, ffn_mult=1, **kwargs)
- if size == 'tiny':
- return SADelARTransformer(depth=4, n_head=6, **kwargs)
- if size == 'base':
- return SADelARTransformer(depth=6, n_head=8, **kwargs)
- if size == 'base-deep':
- return SADelARTransformer(depth=9, n_head=8, **kwargs)
- if size == 'base-wide':
- return SADelARTransformer(depth=6, n_head=12, **kwargs)
- if size == 'small/2':
- return SADelARTransformer(depth=9, n_head=12, **kwargs)
- if size == 'small':
- return SADelARTransformer(depth=12, n_head=12, **kwargs)
- if size == 'medium':
- return SADelARTransformer(depth=24, n_head=16, **kwargs)
-
-def make_model(size:str, quantizers:int=4, frozen_embeddings_model:str=None, frozen_acoustic_embeddings:bool=False, spk_width:int=None, tunables:Tunables=Tunables(), dataset=None):
- from encodec.model import EncodecModel
- from whisperspeech import vq_stoks
-
- amodel = EncodecModel.encodec_model_24khz() if frozen_acoustic_embeddings else None
- vqmodel = vq_stoks.RQBottleneckTransformer.load_model(frozen_embeddings_model) if frozen_embeddings_model else None
- model = _make_model(size, quantizers, tunables,
- spk_width=spk_width,
- atoks_width=amodel and amodel.quantizer.vq.layers[0]._codebook.embed.shape[-1],
- stoks_codes=vqmodel.vq_codes+1, stoks_width=vqmodel.rq.layers[0]._codebook.embed[0].shape[-1])
- if vqmodel: model.load_frozen_semantic_embeddings(vqmodel)
- if amodel: model.load_frozen_acoustic_embeddings(amodel)
- return model

whisperspeech/t2s_up_wds.py

@@ -1,442 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/5B. Text to semantic token modeling.ipynb.
-
-# %% auto 0
-__all__ = ['load_datasets', 'rand', 'Tunables', 'Encoder', 'Decoder', 'TSARTransformer', 'make_model']
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 1
-import dataclasses
-import random
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.profiler import record_function
-
-from huggingface_hub import hf_hub_download
-from fastcore.basics import store_attr
-from fastprogress import progress_bar
-
-import webdataset as wds
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 2
-from pathlib import Path
-import pylab as plt
-import pandas as pd
-import numpy as np
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 3
-import whisper
-from whisperspeech.train import *
-from whisperspeech.modules import *
-from whisperspeech import vq_stoks
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 8
-import re
-
-class CharTokenizer:
- """Trivial tokenizer – just use UTF-8 bytes"""
- eot = 0
- 
- def encode(self, txt):
- return list(bytes(txt.strip(), 'utf-8'))
-
- def decode(self, tokens):
- return bytes(tokens).decode('utf-8')
- 
-def tokenizer(ikey, okey, length):
- """Tokenizes a transcript"""
- tok = CharTokenizer()
- def _tokenizer(samples):
- for s in samples:
- toks = torch.tensor(tok.encode(s[ikey]))
- s[okey] = F.pad(toks, (0, length - toks.shape[-1]), value=tok.eot)
- yield s
- return _tokenizer
-
-def ar_padder(ikey, okey, length, pad_token):
- """Pads the tokens for autoregresive training"""
- def _ar_padder(samples):
- for s in samples:
- toks = s[ikey]
- if isinstance(toks, (list, np.ndarray)): toks = torch.tensor(toks)
- toks = toks.to(torch.long)
- s['in_' +okey] = F.pad(toks, (1, length - toks.shape[-1] - 1), value=pad_token)
- s['out_'+okey] = F.pad(toks, (0, length - toks.shape[-1]), value=pad_token)
- yield s
- return _ar_padder
-
-def char_per_seconder(txt_key, stoks_key, cps_key, stoks_per_second=25):
- """Adds the characters per second metric to the input data"""
- def _char_per_seconder(samples):
- for s in samples:
- secs = s[stoks_key].shape[-1] / stoks_per_second
- s[cps_key] = len(s[txt_key]) / secs
- yield s
- return _char_per_seconder
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 9
-def build_speaker_map(shards):
- speakers = set()
- for shard in shards:
- with open(shard+'.speakers.txt') as f: speakers = speakers.union(set(x.strip() for x in f.readlines()))
- return {id:i for i,id in enumerate(speakers)}
-
-def speaker_id_extractor(speaker_map):
- def _extractor(samples):
- for s in samples:
- s['speaker'] = torch.tensor(speaker_map[s['__key__'].split("/")[1]])
- yield s
- return _extractor
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 10
-def load_datasets(
- input:str, # webdataset folder or shard list
- samples:int, # samples per epoch
- subsample:float=1, # use a fraction of the files
- val_samples:int=512,
- vq_codes:int=4096,
-):
- if isinstance(input, (Path, str)):
- path = Path(input)
- if path.is_dir():
- glob = '*-t2s-*.tar.gz'
- else:
- glob = path.name
- path = path.parent
- input = Path(path).glob(glob)
- elif isinstance(input, list):
- pass
- else:
- raise ArgumentError("input should be either a list of a path with an optional glob specifier")
- shards = [str(x) for x in input]
-
- speaker_map = build_speaker_map(shards)
- 
- def ds(shards, length):
- ds = wds.WebDataset(wds.ResampledShards(shards)).compose(
- wds.decode(),
- speaker_id_extractor(speaker_map),
- wds.select(lambda s: s['stoks.npy'].shape[-1] > 12), # select samples > .5s
- tokenizer('txt', 'ttoks', length=550),
- ar_padder('stoks.npy', 'stoks', length=750, pad_token=vq_codes-1),
- char_per_seconder('txt', 'stoks.npy', 'cps', stoks_per_second=25),
- wds.to_tuple('ttoks', 'speaker', 'cps', 'in_stoks', 'out_stoks'),
- wds.batched(64)
- )
- ds.speakers = speaker_map
- ds.total_samples = length
- ds.stoks_len = 750
- ds.stoks_codes = vq_codes
- ds.ttoks_len = 550
- return ds.compose(wds.slice(length // 64)).with_epoch(length // 64).with_length(length // 64)
- 
- return (
- ds(shards[1:], samples),
- ds(shards[:1], val_samples),
- )
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 14
-def rand(start, end):
- return random.random() * (end - start) + start
-
-@dataclasses.dataclass
-class Tunables:
- init_std :float = 1
- embeddings_std :float = .01
- embeddings_lr_scale: float = 5
- embedding_projector_lr_scale: float = 2.5
- output_mult :float = .35
- query_mult :float = 1
- encoder_depth_ratio :float = 0.25
- eot_dropout_p :float = .5
- cps_input: bool = True
- cps_bins: int = 32
- 
- lr0 :float = 1.5e-3
- clip_gradient_norm :float = .2
- weight_decay :float = 1e-1
- warmup_steps :float = 4000
-
- random :bool = False
-
- def __post_init__(self):
- # randomize the hyperparams if requested
- if self.random:
- self.init_std = 10**rand(-1,1)
- self.embeddings_std = 10**rand(-3,-.7)
- self.embeddings_lr_scale = rand(2,6)
- self.output_mult = rand(0.25,0.65)
- self.query_mult = 2**rand(-2,3)
- self.encoder_depth_ratio = 0.25
- 
- self.lr0 = rand(1,5)*1e-3
- self.clip_gradient_norm = 10**rand(-3,0)
- self.warmup_steps = 100*(10**rand(1,1.85))
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 15
-class EmbeddingProjector(nn.Linear):
- pass
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 16
-class Encoder(nn.Module):
- def __init__(self, depth=6, width=384, n_head=6, length=1500, codes=1024, emb_width=384, ffn_mult=4, pos_embs=None, tunables=Tunables()):
- super().__init__()
- self.emb_width = emb_width
- 
- self.emb_factor = width != emb_width
- 
- self.embedding = nn.Embedding(codes, emb_width)
- if self.emb_factor:
- self.emb_to_hidden = EmbeddingProjector(emb_width, width)
-
- if pos_embs is None: pos_embs = sinusoids(length, width)
- self.register_buffer("positional_embedding", pos_embs)
-
- self.layers = nn.Sequential(*[
- ResidualAttentionBlock(width, n_head,
- qk_scale=tunables.query_mult*8/math.sqrt(width/n_head), ffn_mult=ffn_mult) for _ in range(depth)
- ])
-
- self.ln_post = LayerNorm(width)
- 
- def forward(self, Stoks):
- xin = self.embedding(Stoks)
- if self.emb_factor:
- xin = self.emb_to_hidden(xin)
- 
- assert xin.shape[1:] == self.positional_embedding.shape, "incorrect semantic token shape"
- xin = (xin + self.positional_embedding).to(xin.dtype)
-
- return self.ln_post(self.layers(xin))
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 17
-class Decoder(nn.Module):
- def __init__(self, depth=6, stoks_width=384, width=384, n_head=6, length=1500, codes=1024, ffn_mult=4, pos_embs=None, tunables=Tunables()):
- super().__init__()
- self.length = length
- self.codes = codes
- self.width = width
- self.stoks_width = stoks_width
- 
- self.emb_factor = width != stoks_width
- 
- # embed semantic tokens
- self.embedding = nn.Embedding(codes, stoks_width)
- if self.emb_factor:
- self.emb_to_hidden = EmbeddingProjector(stoks_width, width)
- self.hidden_to_emb = EmbeddingProjector(width, stoks_width)
-
- if pos_embs is None: pos_embs = sinusoids(length, width)
- self.register_buffer("positional_embedding", pos_embs)
- 
- self.layers = nn.ModuleList([
- ResidualAttentionBlock(width, n_head, cross_attention=True,
- qk_scale=tunables.query_mult*8/math.sqrt(width/n_head), ffn_mult=ffn_mult) for _ in range(depth)
- ])
- self.ln_post = LayerNorm(width)
- 
- def forward(self, Stoks, xenc, cps=None):
- Sembs = self.embedding(Stoks)
- 
- if self.emb_factor:
- Sembs = self.emb_to_hidden(Sembs)
- 
- xin = (Sembs + self.positional_embedding[:Sembs.shape[1]]).to(xenc.dtype)
- if cps is not None: xin = xin + cps
- 
- x = xin
- for l in self.layers: x = l(x, xenc, causal=True)
- 
- x = self.ln_post(x)
- 
- if self.emb_factor:
- x = self.hidden_to_emb(x)
- 
- logits = (x @ self.embedding.weight.to(x.dtype).T).float()
- return logits
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 18
-class TSARTransformer(nn.Module):
- def __init__(self, depth=6, n_head=6, head_width=64, ffn_mult=4, language='en',
- ttoks_len=200, ttoks_codes=50364, ttoks_width=None,
- stoks_len=1500, stoks_codes=1024, stoks_width=None,
- tunables=Tunables()):
- assert language == 'en', "only english is supported right now"
- super().__init__()
- store_attr("depth,n_head,head_width,ffn_mult,stoks_width,ttoks_width,ttoks_len,stoks_len,ttoks_codes,stoks_codes,language")
-
- width = n_head * head_width
- self.width = width
- self.base_width = 3 * head_width
- self.tunables = tunables
- if self.stoks_width is None: self.stoks_width = self.width
- if self.ttoks_width is None: self.ttoks_width = self.width
- 
- if tunables.cps_input:
- self.cps_embeddings = nn.Embedding(tunables.cps_bins, self.width)
- else:
- self.cps_embeddings = None
- 
- encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
- decoder_depth = depth * 2 - encoder_depth
- tformer_args = dict(width=width, n_head=n_head, ffn_mult=ffn_mult, tunables=tunables)
- self.encoder = Encoder(length=ttoks_len, codes=ttoks_codes, emb_width=self.ttoks_width, depth=encoder_depth, **tformer_args)
- self.decoder = Decoder(length=stoks_len, codes=stoks_codes, stoks_width=self.stoks_width, depth=decoder_depth, **tformer_args)
- 
- self.tokenizer = None
- 
- self.apply(self.init_transformer)
-
- def load_frozen_semantic_embeddings(self, vqmodel):
- with torch.no_grad():
- self.decoder.embedding.weight[:] = vqmodel.rq.layers[0]._codebook.embed[0]
- self.decoder.embedding.lr_scale = 0
-
- def setup(self, device):
- pass
-
- def init_transformer(self, m):
- if isinstance(m, LinearHead):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, QueryHead):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, nn.Embedding):
- m.no_weight_decay = True
- m.lr_scale = self.tunables.embeddings_lr_scale
- std = self.tunables.embeddings_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, EmbeddingProjector):
- m.lr_scale = self.tunables.embedding_projector_lr_scale
- std = self.tunables.init_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.Linear):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- std = self.tunables.init_std / m.weight.shape[1]
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- if m.bias is not None:
- torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.LayerNorm):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.bias, 0)
- torch.nn.init.constant_(m.weight, 1)
- 
- def forward(self, Ttoks, speakers, cpss, in_stoks, out_stoks=None, loss=True):
- with record_function("encoder"):
- xenc = self.encoder(Ttoks.to(torch.long))
- with record_function("decoder"):
- if self.cps_embeddings:
- cps_bin = (cpss / 20 * self.tunables.cps_bins).to(torch.long)
- cps_bin[cps_bin >= self.tunables.cps_bins] = self.tunables.cps_bins-1
- cps_embs = self.cps_embeddings(cps_bin).unsqueeze(1)
- else:
- cps_embs = None
- logits = self.decoder(in_stoks, xenc, cps=cps_embs) * self.tunables.output_mult / (self.width / self.base_width)
- if loss is not None:
- with record_function("loss"):
- loss = F.cross_entropy(logits.transpose(-1,-2), out_stoks)#, reduction='none')
- return logits, loss
-
- #
- # inference
- #
- @classmethod
- def load_model(cls, repo_id="collabora/whisperspeech", filename="t2s_up_wds.model", local_filename=None):
- if not local_filename:
- local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
- spec = torch.load(local_filename)
- model = cls(**spec['config'], tunables=Tunables(**spec['tunables']))
- model.load_state_dict(spec['state_dict'])
- model.eval()
- return model
-
- def load_checkpoint(self, local_filename):
- spec = torch.load(local_filename, map_location='cpu')
- assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'
- state_dict = {k.replace('model.', ''):v
- for k,v in spec['state_dict'].items()}
- self.load_state_dict(state_dict)
- return self
-
- def save_model(self, fname):
- torch.save(dict(config = self.__stored_args__,
- tunables = dataclasses.asdict(self.tunables),
- state_dict = self.state_dict()), fname)
-
- def ensure_tokenizer(self):
- assert not self.training
- if self.tokenizer is None: self.tokenizer = CharTokenizer()
- #whisper.tokenizer.get_tokenizer(multilingual=True) 
-
- @property
- def device(self):
- return next(self.parameters()).device
- 
- @torch.no_grad()
- def generate(self, txt, cps=15, N=None, T=0.7, top_k=None, show_progress_bar=True):
- self.ensure_tokenizer()
- N = N or self.stoks_len
- dev = self.device
- ttoks = torch.tensor(self.tokenizer.encode(txt), device=dev)
- ttoks = F.pad(ttoks, (0, self.ttoks_len - len(ttoks)), value=self.tokenizer.eot).unsqueeze(0)
- cpss = torch.tensor([cps], device=dev)
- toks = torch.zeros((1,N), dtype=torch.long, device=dev)
- toks[0,0] = self.stoks_codes-1
- it = range(1,N)
- if show_progress_bar: it = progress_bar(it)
- for i in it:
- p, _ = self(ttoks, None, cpss, toks[:,:i], loss=None)
- last_p = p[0,-1]
- if top_k:
- last_p[last_p < torch.topk(last_p, top_k).values[-1,None]] = -torch.inf
- tok = torch.multinomial((last_p / float(T)).softmax(-1), 1)
- toks[0,i] = tok
- if toks[0,i] == self.stoks_codes-1: return toks[0,1:i]
- return toks[0,1:]
- 
- @torch.no_grad()
- def generate_batch(self, txts, N=None, T=1.1, top_k=7, show_progress_bar=True):
- self.ensure_tokenizer()
- N = self.stoks_len
- dev = self.device
- ttoks = []
- for txt in txts:
- ttoks_ = torch.tensor(self.tokenizer.encode(txt), device=dev)
- ttoks_ = F.pad(ttoks_, (0, self.ttoks_len - len(ttoks_)), value=self.tokenizer.eot).unsqueeze(0)
- ttoks.append(ttoks_)
- ttoks = torch.cat(ttoks, dim=0)
- toks = torch.zeros((len(ttoks),N), dtype=torch.long, device=dev)
- it = range(N)
- if show_progress_bar: it = progress_bar(it)
- for i in it:
- p, _ = self(ttoks, toks[:,:i], loss=None)
- last_p = p[:,-1]
- if top_k:
- last_p[last_p < torch.topk(last_p, top_k).values[:,-1,None]] = -torch.inf
- tok = torch.multinomial((last_p / float(T)).softmax(-1), 1)
- toks[:,i] = tok[:,0]
- if (toks[:,i] == self.stoks_codes-1).all(): return toks[:,:i]
- return toks
-
-# %% ../nbs/5B. Text to semantic token modeling.ipynb 19
-def _make_model(size:str, tunables:Tunables=Tunables(), dataset=None, **kwargs):
- kwargs = dict(stoks_len = dataset.stoks_len, ttoks_len = dataset.ttoks_len, tunables=tunables, **kwargs)
- if 'stoks_codes' not in kwargs: kwargs['stoks_codes'] = dataset.stoks_codes
- if size == 'micro':
- return TSARTransformer(depth=2, n_head=3, ffn_mult=1, **kwargs)
- if size == 'tiny':
- return TSARTransformer(depth=4, n_head=6, **kwargs)
- if size == 'base':
- return TSARTransformer(depth=6, n_head=8, **kwargs)
- if size == 'small':
- return TSARTransformer(depth=12, n_head=16, **kwargs)
-
-def make_model(size:str, frozen_embeddings_model:str=None, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
- if frozen_embeddings_model:
- vqmodel = vq_stoks.RQBottleneckTransformer.load_model(frozen_embeddings_model)
- model = _make_model(size, tunables, dataset, stoks_codes=vqmodel.vq_codes+1, stoks_width=vqmodel.rq.layers[0]._codebook.embed[0].shape[-1])
- model.load_frozen_semantic_embeddings(vqmodel)
- else:
- model = _make_model(size, quantizers, tunables, dataset)
- return model

whisperspeech/t2s_up_wds_mlang_enclm.py

@@ -1,519 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/5B. Multi-lang text to semantic token modeling.ipynb.
-
-# %% auto 0
-__all__ = ['load_dataset', 'rand', 'Tunables', 'T2SEmbedding', 'Encoder', 'TSARTransformer', 'make_model']
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 1
-import dataclasses
-import random
-import math
-import itertools
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.profiler import record_function
-
-from huggingface_hub import hf_hub_download
-from fastcore.basics import store_attr
-from fastprogress import progress_bar
-
-from pathlib import Path
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 2
-from whisperspeech.modules import *
-from whisperspeech import languages
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 6
-import re
-
-class CharTokenizer:
- """Trivial tokenizer – just use UTF-8 bytes"""
- eot = 0
- 
- def encode(self, txt):
- return list(bytes(txt.strip(), 'utf-8'))
-
- def decode(self, tokens):
- return bytes(tokens).decode('utf-8')
- 
-def tokenizer(ikey, okey, length):
- """Tokenizes a transcript"""
- tok = CharTokenizer()
- def _tokenizer(samples):
- for s in samples:
- toks = torch.tensor(tok.encode(s[ikey]))
- s[okey] = F.pad(toks, (0, length - toks.shape[-1]), value=tok.eot)
- yield s
- return _tokenizer
-
-def ar_padder(ikey, okey, length, pad_token):
- """Pads the tokens for autoregresive training"""
- import numpy as np
-
- def _ar_padder(samples):
- for s in samples:
- toks = s[ikey]
- if isinstance(toks, (list, np.ndarray)): toks = torch.tensor(toks)
- toks = toks.to(torch.long)
- s['in_' +okey] = F.pad(toks, (1, length - toks.shape[-1] - 1), value=pad_token)
- s['out_'+okey] = F.pad(toks, (0, length - toks.shape[-1]), value=pad_token)
- yield s
- return _ar_padder
-
-def char_per_seconder(txt_key, stoks_key, cps_key, stoks_per_second=25):
- """Adds the characters per second metric to the input data"""
- def _char_per_seconder(samples):
- for s in samples:
- secs = s[stoks_key].shape[-1] / stoks_per_second
- s[cps_key] = len(s[txt_key]) / secs
- yield s
- return _char_per_seconder
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 7
-def load_dataset(
- txt_shard_spec:str, # transcription webdataset shards
- stoks_shard_dir:str, # stoks webdataset base dir
- samples:int, # samples per epoch
- txt_kind:str='small.en-txt',
- vq_codes:int=4096,
- language:str='en',
- weight:float=1,
- validation:bool=False,
- exclude_files:str=None,
-):
- import webdataset as wds
- from whisperspeech import utils
-
- shards = utils.shard_glob(txt_shard_spec)
- excludes = {x for file in exclude_files.split() for x in utils.readlines(file)} if exclude_files else set()
- 
- language = languages.to_id(language)
- 
- def set_language(x):
- x['language'] = language
- return x
- 
- same_on_all_nodes = lambda urls: urls # will only be used for validation
- ds = wds.WebDataset(shards, resampled=not validation, nodesplitter=same_on_all_nodes).compose(
- wds.decode(),
- utils.merge_in(utils.derived_dataset('eqvad-stoks', base=txt_kind, suffix='', dir=stoks_shard_dir)),
- # discard validation samples, select samples > .5s
- wds.select(lambda s: s['__key__'] not in excludes and s['stoks.npy'].shape[-1] > 12),
- tokenizer('txt', 'ttoks', length=550),
- ar_padder('stoks.npy', 'stoks', length=750, pad_token=vq_codes-1),
- ar_padder('ttoks', 'ttoks', length=550, pad_token=CharTokenizer.eot),
- char_per_seconder('txt', 'stoks.npy', 'cps', stoks_per_second=25),
- wds.map(set_language),
- wds.to_tuple('in_ttoks', 'out_ttoks', 'language', 'cps', 'in_stoks', 'out_stoks'),
- wds.shuffle(20000, initial=20000),
- wds.batched(64)
- )
- if validation:
- ds = ds.slice(samples // 64)
- ds.total_samples = samples
- ds.stoks_len = 750
- ds.stoks_codes = vq_codes
- ds.ttoks_len = 550
- ds.weight = weight
-
- return ds
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 14
-def rand(start, end):
- return random.random() * (end - start) + start
-
-@dataclasses.dataclass
-class Tunables:
- init_std :float = 1
- embeddings_std :float = .01
- embeddings_lr_scale: float = 5
- embedding_projector_lr_scale: float = 2.5
- output_mult :float = .35
- query_mult :float = 1
- encoder_depth_ratio :float = 0.25
- eot_dropout_p :float = .5
- cps_input: bool = True
- cps_bins: int = 32
- 
- lr0 :float = 1.5e-3
- clip_gradient_norm :float = .2
- weight_decay :float = 1e-1
- warmup_steps :float = 4000
-
- random :bool = False
-
- def __post_init__(self):
- # randomize the hyperparams if requested
- if self.random:
- self.init_std = 10**rand(-1,1)
- self.embeddings_std = 10**rand(-3,-.7)
- self.embeddings_lr_scale = rand(2,6)
- self.output_mult = rand(0.25,0.65)
- self.query_mult = 2**rand(-2,3)
- self.encoder_depth_ratio = 0.25
- 
- self.lr0 = rand(1,5)*1e-3
- self.clip_gradient_norm = 10**rand(-3,0)
- self.warmup_steps = 100*(10**rand(1,1.85))
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 15
-class T2SEmbedding(nn.Module):
- def __init__(self, length=1500, codes=1024, width=384, pos_embs=None, stoks_width=384):
- super().__init__()
- self.embedding = FlexEmbeddings(codes, width, special_codes=1, frozen_width=stoks_width)
- if pos_embs is None: pos_embs = sinusoids(length, width)
- self.register_buffer("positional_embedding", pos_embs)
- 
- def forward(self, Stoks, xenc, cps=None, offset=0):
- Sembs = self.embedding(Stoks)
- xin = (Sembs + self.positional_embedding[offset : offset + Sembs.shape[1]]).to(xenc.dtype)
- if cps is not None: xin = xin + cps
- return xin, offset
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 16
-class Encoder(nn.Module):
- def __init__(self, depth=6, width=384, n_head=6, length=1500, codes=1024, emb_width=384, ffn_mult=4, pos_embs=None, tunables=Tunables()):
- super().__init__()
- self.emb_width = emb_width
- 
- self.embedding = FlexEmbeddings(codes, width, frozen_width=emb_width)
-
- if pos_embs is None: pos_embs = sinusoids(length, width)
- self.register_buffer("positional_embedding", pos_embs)
-
- self.layers = nn.ModuleList([
- ResidualAttentionBlock(width, n_head,
- qk_scale=tunables.query_mult*8/math.sqrt(width/n_head), ffn_mult=ffn_mult) for _ in range(depth)
- ])
-
- self.ln_post = LayerNorm(width)
- 
- mask = torch.empty(length, length).fill_(-torch.inf).triu_(1)
- self.register_buffer("mask", mask, persistent=False)
- 
- def forward(self, Stoks, positions, lang_emb=None):
- xin = self.embedding(Stoks)
-
- if lang_emb is not None: xin += lang_emb
- 
-# assert xin.shape[1:] == self.positional_embedding.shape, "incorrect semantic token shape"
- x = (xin +
- self.positional_embedding[positions]).to(xin.dtype)
-
- for l in self.layers: x = l(x, positions, causal=False, mask=self.mask)
- 
- return self.ln_post(x)
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 17
-class TSARTransformer(nn.Module):
- def __init__(self, depth=6, n_head=6, head_width=64, ffn_mult=4,
- ttoks_len=200, ttoks_codes=256, ttoks_width=None,
- stoks_len=1500, stoks_codes=1024, stoks_width=None,
- tunables=Tunables()):
- super().__init__()
- store_attr("depth,n_head,head_width,ffn_mult,stoks_width,ttoks_width,ttoks_len,stoks_len,ttoks_codes,stoks_codes")
-
- width = n_head * head_width
- self.width = width
- self.base_width = 3 * head_width
- self.tunables = tunables
- if self.stoks_width is None: self.stoks_width = self.width
- if self.ttoks_width is None: self.ttoks_width = self.width
- 
- self.lang_embeddings = nn.Embedding(len(languages.languages), width)
- if tunables.cps_input:
- self.cps_embeddings = nn.Embedding(tunables.cps_bins, self.width)
- else:
- self.cps_embeddings = None 
- 
- encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
- decoder_depth = depth * 2 - encoder_depth
- tformer_args = dict(width=width, n_head=n_head, ffn_mult=ffn_mult, tunables=tunables)
- self.encoder = Encoder(length=ttoks_len, codes=ttoks_codes, emb_width=self.ttoks_width, depth=encoder_depth, **tformer_args)
- self.embeddings = T2SEmbedding(length=stoks_len, codes=stoks_codes, width=width, stoks_width=self.stoks_width)
-
- self.decoder = BaseDecoder(
- length=stoks_len, 
- depth=decoder_depth,
- qk_scale=tunables.query_mult*8/math.sqrt(width/n_head),
- width=width, n_head=n_head, ffn_mult=ffn_mult,
- )
- self.tokenizer = None
- 
- self.apply(self.init_transformer)
-
- def load_frozen_semantic_embeddings(self, vqmodel):
- self.embeddings.embedding.set_frozen_embeddings(vqmodel.rq.layers[0]._codebook.embed[0])
-
- def setup(self, device):
- pass
-
- def init_transformer(self, m):
- if isinstance(m, LinearHead):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, QueryHead):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, nn.Embedding):
- m.no_weight_decay = True
- m.lr_scale = self.tunables.embeddings_lr_scale
- std = self.tunables.embeddings_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, EmbeddingProjector):
- m.lr_scale = self.tunables.embedding_projector_lr_scale
- std = self.tunables.init_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.Linear):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- std = self.tunables.init_std / m.weight.shape[1]
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- if m.bias is not None:
- torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.LayerNorm):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.bias, 0)
- torch.nn.init.constant_(m.weight, 1)
- 
- def _embed_cps(self, cpss):
- if self.cps_embeddings is None: return None
-
- cps_bin = (cpss / 20 * self.tunables.cps_bins).to(torch.long)
- cps_bin[cps_bin >= self.tunables.cps_bins] = self.tunables.cps_bins-1
- return self.cps_embeddings(cps_bin).unsqueeze(1)
-
- def run_encoder(self, in_ttoks, languages, cpss):
- if len(languages.shape) != 3: lang_embs = self.lang_embeddings(languages)
- else: lang_embs = languages
- if len(lang_embs.shape) == 2: lang_embs = lang_embs.unsqueeze(1)
- 
- cps_emb = self._embed_cps(cpss)
-
- with record_function("encoder"):
- positions = torch.arange(0, in_ttoks.shape[1], device=in_ttoks.device)
- xenc = self.encoder(in_ttoks.to(torch.long), positions, lang_emb=lang_embs)
-
- return xenc, positions, cps_emb
- 
- def forward(self, in_ttoks, out_ttoks, languages, cpss, in_stoks, in_stoks_positions, out_stoks=None, loss=True, offset=None, xenc=None, xenc_positions=None, cps_emb=None):
- if xenc is None:
- xenc, cps_emb = self.run_encoder(in_ttoks, languages, cpss)
-
- with record_function("decoder"):
- x = (self.embeddings.embedding(in_stoks) + 
- self.embeddings.positional_embedding[in_stoks_positions] +
- cps_emb).to(xenc[0].dtype)
- x = self.decoder(x, in_stoks_positions, xenc, xenc_positions)
- logits = self.embeddings.embedding.unembed(x)
- logits = logits * self.tunables.output_mult / (self.width / self.base_width)
-
- if loss is not None:
- enc_logits = self.encoder.embedding.unembed(xenc[0])
- enc_logits = enc_logits * self.tunables.output_mult / (self.width / self.base_width)
- with record_function("loss"):
- loss = F.cross_entropy(logits.transpose(-1,-2), out_stoks)
- if self.training:
- loss += 0.1 * F.cross_entropy(enc_logits.transpose(-1,-2), out_ttoks)
- 
- return logits, loss
-
- #
- # inference
- #
- @classmethod
- def load_model(cls, ref="collabora/whisperspeech:t2s-small-en+pl.model",
- repo_id=None, filename=None, local_filename=None):
- if repo_id is None and filename is None and local_filename is None:
- if ":" in ref:
- repo_id, filename = ref.split(":", 1)
- else:
- local_filename = ref
- if not local_filename:
- local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
- spec = torch.load(local_filename)
- model = cls(**spec['config'], tunables=Tunables(**spec['tunables']))
- model.load_state_dict(spec['state_dict'])
- model.eval()
- return model
-
- def load_checkpoint(self, local_filename):
- spec = torch.load(local_filename, map_location='cpu')
- assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'
- state_dict = {k.replace('model.', ''):v
- for k,v in spec['state_dict'].items()}
- self.load_state_dict(state_dict)
- return self
-
- def save_model(self, fname):
- torch.save(dict(config = self.__stored_args__,
- tunables = dataclasses.asdict(self.tunables),
- state_dict = self.state_dict()), fname)
-
- def ensure_tokenizer(self):
- assert not self.training
- if self.tokenizer is None: self.tokenizer = CharTokenizer()
-
- def switch_dtypes(self, dtype=torch.float16):
- self.dtype = dtype
- for n,m in self.named_modules():
- # convert every leaf layer apart from the LayerNorms
- if isinstance(m, (nn.Linear, nn.Embedding)):
- m.to(dtype)
- # take care of buffers ([kv]_cache, masks) that are not in the leaf layers
- for bn,b in m.named_buffers(recurse=False):
- setattr(m,bn,b.to(dtype))
-
- def optimize(self, max_batch_size=1, dtype=torch.float16, torch_compile=True):
- for emb in [self.embeddings.embedding, self.embeddings.embedding]:
- emb.convert_for_eval()
- for l in self.encoder.layers:
- l.attn.convert_for_eval()
- for l in self.decoder.layers:
- l.attn.convert_for_eval()
- l.cross_attn.convert_for_eval()
- l.setup_kv_cache(max_batch_size, self.stoks_len, self.ttoks_len)
- self.switch_dtypes(dtype)
- if torch_compile:
- self.generate_next = torch.compile(self.generate_next, mode="reduce-overhead", fullgraph=True)
-
- @property
- def device(self):
- return next(self.parameters()).device
- 
- # from https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
- def multinomial_sample_one_no_sync(self, probs_sort): # Does multinomial sampling without a cuda synchronization
- q = torch.empty_like(probs_sort).exponential_(1)
- return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
-
- def logits_to_probs(self, logits, T=1.0, top_k=None):
- logits = logits / max(T, 1e-5)
-
- logits[self.embeddings.embedding.codes:] = -torch.inf
- if top_k is not None:
- v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
- pivot = v.select(-1, -1).unsqueeze(-1)
- logits = torch.where(logits < pivot, -float("Inf"), logits)
-
- probs = torch.nn.functional.softmax(logits, dim=-1)
- return probs
-
- def sample(self, logits, T=1.0, top_k=None):
- probs = self.logits_to_probs(logits[0,-1], T, top_k)
- idx_next = self.multinomial_sample_one_no_sync(probs)
- return idx_next
-
- def generate_one(self, toks, toks_positions, cps_emb, xenc, xenc_positions, T, top_k):
- probs, _ = self(None, None, None, None, toks, toks_positions, loss=None, xenc=xenc, xenc_positions=xenc_positions, cps_emb=cps_emb)
- return self.sample(probs, T, top_k)
-
- def generate_next(self, *args, **kwargs):
- return self.generate_one(*args, **kwargs)
-
- @torch.no_grad()
- def prep(self, txt, cps=15, lang="en"):
- dev = self.device
- ttoks = torch.tensor(self.tokenizer.encode(txt), device=dev)
- ttoks = F.pad(ttoks, (0, self.ttoks_len - len(ttoks)), value=self.tokenizer.eot).unsqueeze(0)
- cpss = torch.tensor([cps], device=dev)
- langs = torch.tensor([languages.to_id(lang)], device=dev)
- return ttoks, cpss, langs
- 
- @torch.no_grad()
- def generate(self, txt, cps=15, lang="en", N=None, T=0.7, top_k=None, step=None, show_progress_bar=True):
- self.ensure_tokenizer()
- N = N or self.stoks_len
- dev = self.device
- ttoks = []
- langs = []
- if isinstance(lang, list):
- lang0 = lang[0]
- assert isinstance(txt, list), "lang and txt have to be both lists or strings"
- for txt, lang in zip(txt, lang):
- tt = self.tokenizer.encode(txt)
- ttoks += tt
- langs += [languages.to_id(lang)] * len(tt)
- elif isinstance(lang, torch.Tensor):
- langs = lang
- ttoks = self.tokenizer.encode(txt)
- else:
- lang0 = lang
- ttoks = self.tokenizer.encode(txt)
- langs = torch.tensor([languages.to_id(lang)], device=dev).unsqueeze(0)
- ttoks = torch.tensor(ttoks, device=dev)
- ttoks = F.pad(ttoks, (1, self.ttoks_len - len(ttoks) - 1), value=self.tokenizer.eot).unsqueeze(0)
- cpss = torch.tensor([cps], device=dev)
- if not isinstance(langs, torch.Tensor):
- langs = torch.tensor(langs, device=dev)
- langs = F.pad(langs, (1, self.ttoks_len - len(langs) - 1), value=languages.to_id(lang0)).unsqueeze(0)
- it = range(0,N-1)
- if show_progress_bar: it = progress_bar(it)
-
- toks = torch.zeros((1,N), dtype=torch.long, device=dev)
- toks[:,0] = self.stoks_codes-1
- toks_positions = torch.arange(N, device=dev)
- with record_function("encode"):
- xenc, xenc_positions, cps_emb = self.run_encoder(ttoks, langs, cpss)
- toks_positions = torch.arange(N+1, device=dev)
- # contrary to S2A this model works without prefill and is actually a tiny bit faster
- # with record_function("prefill"):
- # toks[0,1] = self.generate_one(toks[:,:1], toks_positions[:1], cps_emb, xenc, xenc_positions, T, top_k)
- with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
- for i in it:
- toks[0,i+1] = self.generate_next(toks[:,i:i+1], toks_positions[i:i+1], cps_emb, xenc, xenc_positions, T, top_k)
- if i % 25 == 0 and toks[0,i+1] == self.stoks_codes-1: return toks[0,:i+1]
-
- # for profiling, debugging or early exit
- if step is not None: step()
- return toks[0,:]
- 
- @torch.no_grad()
- def generate_batch(self, txts, N=None, T=1.1, top_k=7, show_progress_bar=True):
- self.ensure_tokenizer()
- N = self.stoks_len
- dev = self.device
- ttoks = []
- for txt in txts:
- ttoks_ = torch.tensor(self.tokenizer.encode(txt), device=dev)
- ttoks_ = F.pad(ttoks_, (0, self.ttoks_len - len(ttoks_)), value=self.tokenizer.eot).unsqueeze(0)
- ttoks.append(ttoks_)
- ttoks = torch.cat(ttoks, dim=0)
- toks = torch.zeros((len(ttoks),N), dtype=torch.long, device=dev)
- it = range(N)
- if show_progress_bar: it = progress_bar(it)
- for i in it:
- p, _ = self(ttoks, toks[:,:i], loss=None)
- last_p = p[:,-1]
- if top_k:
- last_p[last_p < torch.topk(last_p, top_k).values[:,-1,None]] = -torch.inf
- tok = torch.multinomial((last_p / float(T)).softmax(-1), 1)
- toks[:,i] = tok[:,0]
- if (toks[:,i] == self.stoks_codes-1).all(): return toks[:,:i]
- return toks
-
-# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 18
-def _make_model(size:str, tunables:Tunables=Tunables(), dataset=None, **kwargs):
- kwargs = dict(stoks_len = dataset.stoks_len, ttoks_len = dataset.ttoks_len, tunables=tunables, **kwargs)
- if 'stoks_codes' not in kwargs: kwargs['stoks_codes'] = dataset.stoks_codes
- if size == 'micro':
- return TSARTransformer(depth=2, n_head=3, ffn_mult=1, **kwargs)
- if size == 'tiny':
- return TSARTransformer(depth=4, n_head=6, **kwargs)
- if size == 'base':
- return TSARTransformer(depth=6, n_head=8, **kwargs)
- if size == 'small':
- return TSARTransformer(depth=12, n_head=12, **kwargs)
- if size == 'small+':
- return TSARTransformer(depth=12, n_head=16, **kwargs)
- if size == 'medium':
- return TSARTransformer(depth=24, n_head=16, **kwargs)
-
-def make_model(size:str, frozen_embeddings_model:str=None, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
- from whisperspeech import vq_stoks
-
- if frozen_embeddings_model:
- vqmodel = vq_stoks.RQBottleneckTransformer.load_model(frozen_embeddings_model)
- model = _make_model(size, tunables, dataset, stoks_codes=vqmodel.vq_codes+1, stoks_width=vqmodel.rq.layers[0]._codebook.embed[0].shape[-1])
- model.load_frozen_semantic_embeddings(vqmodel)
- else:
- model = _make_model(size, tunables, dataset, mode=mode)
- return model

whisperspeech/train.py

@@ -1,271 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B1. Training.ipynb.
-
-# %% auto 0
-__all__ = ['SimpleVisual', 'validate', 'train']
-
-# %% ../nbs/B1. Training.ipynb 2
-import io
-import time
-import random
-from pathlib import Path
-
-from fastprogress import progress_bar, master_bar
-import fastprogress
-
-import numpy as np
-import pylab as plt
-import math
-
-import IPython
-
-import torch
-import torch.nn as nn
-from torch.utils.data.dataloader import DataLoader
-from torch.profiler import record_function
-
-import webdataset as wds
-
-torch.backends.cudnn.benchmark = True
-torch.backends.cudnn.enabled = True
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.set_float32_matmul_precision('medium')
-
-# %% ../nbs/B1. Training.ipynb 3
-class SimpleVisual:
- def __init__ (self, model, masterbar, total_steps):
- self.model = model
- self.masterbar = masterbar
- self.total_steps = total_steps
- self.epochs = total_steps // masterbar.main_bar.total
- 
- gs = plt.GridSpec(2, 1, height_ratios=[3,1])
- graph_fig = plt.figure(figsize=(10,6))
- self.graph_fig = graph_fig
- self.loss_p = graph_fig.add_subplot(gs[0])
- self.lr_p = graph_fig.add_subplot(gs[1], sharex=self.loss_p)
- self.lr_p.tick_params('x', labelbottom=False)
- self.graph_out = None
- 
- self.its = []
- self.train_losses = []
- self.val_losses = []
- self.lr_history = []
- 
- def show(self):
- self.start_t = time.time()
- self.masterbar.write(["samples", "train", "val", "time"], table=True)
- self.graph_out = display(self.graph_fig, display_id=True, clear=True)
- 
- def hide(self):
- if self.graph_out is not None:
- self.graph_out.update(IPython.display.HTML(''))
- 
- def plot(self):
- loss_p, lr_p = self.loss_p, self.lr_p
- loss_p.clear()
- loss_p.plot(self.its, self.train_losses)
- loss_p.plot(self.its, self.val_losses)
- loss_p.set_xlim(0, self.total_steps)
- loss_p.set_yscale('log')
- lr_p.clear()
- lrs = np.array(self.lr_history)
- lr_p.plot(self.its, lrs)
- self.graph_out.update(self.graph_fig)
- 
- def add_data(self, it, lr, train_loss, val_los):
- self.its.append(it)
- self.train_losses.append(train_loss)
- self.val_losses.append(val_los)
- self.lr_history.append(lr)
- self.plot()
-
- def add_table_row(self, it, avg_train_loss, val_loss):
- elapsed_t = time.time() - self.start_t
- self.masterbar.write([it, f"{avg_train_loss:.5f}", f"{val_loss:.5f}", fastprogress.core.format_time(elapsed_t)], table=True)
- 
- def on_iter(self, bar, it, avg_train_loss, val_loss):
- epoch = math.ceil(it / self.total_steps * self.epochs)
- bar.comment = f"#{epoch}/{self.epochs} loss: {avg_train_loss:.3f} / {val_loss:.3f}"
-
-# %% ../nbs/B1. Training.ipynb 4
-# FIXME: we need to keep this synchronised with the validation code below...
-def validate(model, val, half=True, bs=16, drop_last=False, dl_workers=8, device="cuda"):
- if isinstance(val, torch.utils.data.IterableDataset):
- val_loader = wds.WebLoader(val, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \
- .unbatched().shuffle(1024).batched(bs)
- else:
- val_loader = DataLoader(val, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last)
- 
- with torch.no_grad():
- val_loss = 0
- val_samples = 0
- for args in val_loader:
- args = [x.to(device, non_blocking=True) for x in args]
- with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):
- ps, loss = model(*args)
- N = args[0].shape[0]
- val_loss += loss.mean().item() * N
- val_samples += N
- val_loss = val_loss / val_samples
-
- return val_loss
-
-# %% ../nbs/B1. Training.ipynb 5
-def train(checkpoint_path, model, train, val, half=True, bs=16, lr=1e-4, drop_last=False,
- weight_decay=0.1, warmup_steps=10000, epochs=10, clip_gradient_norm=None,
- dl_workers=8, visual_class = SimpleVisual, profiler=None,
- run_valid_every_iters=8000, table_row_every_iters=80000, chkpt_every_iters=None,
- device="cuda", trainable_params=None):
- if chkpt_every_iters is None:
- chkpt_every_iters = table_row_every_iters
- 
- mb = master_bar(range(epochs))
- if isinstance(train, torch.utils.data.IterableDataset):
- pct_start = min(0.3, warmup_steps / (epochs * (train.total_samples//bs)))
- visual = visual_class(model, mb, epochs * train.total_samples)
-# pct_start = min(0.3, warmup_steps / (epochs * len(train)))
-# visual = visual_class(model, mb, epochs*len(train)*bs)
- else:
- pct_start = min(0.3, warmup_steps / (epochs * len(train) / bs))
- visual = visual_class(model, mb, epochs*len(train))
- model.visual = visual
- 
- Path(checkpoint_path).mkdir(exist_ok=True)
-
- if isinstance(train, torch.utils.data.IterableDataset):
-# train_loader = DataLoader(train, batch_size=None, num_workers=dl_workers, pin_memory=True, drop_last=False, shuffle=False)
-# val_loader = DataLoader(val, batch_size=None, num_workers=dl_workers, pin_memory=True, drop_last=False)
- train_loader = wds.WebLoader(train, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \
- .unbatched().shuffle(1024).batched(bs, partial=False)
- val_loader = wds.WebLoader(val, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \
- .unbatched().shuffle(1024).batched(bs)
- else:
- train_loader = DataLoader(train, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last, shuffle=True)
- val_loader = DataLoader(val, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last)
- 
- val_loss = torch.nan
- avg_train_loss = torch.nan
- 
- if hasattr(model, 'setup'):
- model.setup(device)
- 
- try:
- scheduler = None
-
- if trainable_params is None: trainable_params = model.parameters()
- all_params = set(trainable_params)
- customized_params = set()
- groups = []
- group_map = {}
- for name,m in model.named_modules():
- if hasattr(m, 'no_weight_decay') or hasattr(m, 'lr_scale'):
- m_trainable = [x for x in m.parameters() if x in all_params]
- if not m_trainable: continue
- customized_params |= set(m_trainable)
- m_wd = 0 if hasattr(m, 'no_weight_decay') else weight_decay
- m_lr = lr * getattr(m, 'lr_scale', 1)
- group = group_map.get((m_wd, m_lr), None)
- if not group:
- group = {"params": [], "names": [], "weight_decay": m_wd, "lr": m_lr}
- groups.append(group)
- group_map[(m_wd, m_lr)] = group
- group['params'] += m_trainable
- group['names'].append(name)
- 
- other_params = all_params - customized_params
- 
- if other_params:
- groups = groups + [
- {"names": ["other"], "params": list(other_params), "weight_decay": weight_decay },
- ]
-
- optimizer = torch.optim.AdamW(lr=lr, betas=(0.9, 0.95), fused=device!='cpu', params=groups)
- model._optimizer = optimizer
- scaler = torch.cuda.amp.GradScaler(enabled=half)
- scheduler = torch.optim.lr_scheduler.OneCycleLR(
- optimizer, pct_start=pct_start, steps_per_epoch=math.ceil(train.total_samples/bs), epochs=epochs,
- max_lr=[pg.get('lr', lr) for pg in groups],
- final_div_factor=25)
- 
- it = 0
- next_val_it = it + 50
- next_chkpt_it = chkpt_every_iters
- next_table_it = table_row_every_iters
- 
- visual.show()
-
- running_loss = [0]
- 
- for epoch in mb:
- bar = progress_bar(train_loader, total=train.total_samples//bs, parent=mb)
- for args in bar:
- with record_function("forward"):
- args = [x.to(device, non_blocking=True) for x in args]
-
- # zero the parameter gradients
- optimizer.zero_grad(set_to_none=True)
-
- with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):
- ps, loss = model(*args)
- loss = loss.mean()
-
- with record_function("backward"):
- scaler.scale(loss).backward()
-
- if clip_gradient_norm:
- scaler.unscale_(optimizer)
- # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
- torch.nn.utils.clip_grad_norm_(model.parameters(), clip_gradient_norm)
-
- scaler.step(optimizer)
- scaler.update()
-
- scheduler.step()
-
- if profiler is not None: profiler.step()
-
- with record_function("running_loss"):
- running_loss.append(loss.item())
- running_loss = running_loss[-5:]
- avg_train_loss = sum(running_loss)/len(running_loss)
-
- if it >= next_chkpt_it:
- with record_function("checkpoint"):
- next_chkpt_it += chkpt_every_iters
- torch.save(model.state_dict(), f'{checkpoint_path}/{it:08d}.pt')
- 
- if it >= next_val_it:
- next_val_it += run_valid_every_iters
- with record_function("validation"):
- with record_function("model.eval"):
- model.eval()
- with torch.no_grad():
- val_loss = 0
- val_samples = 0
- for args in val_loader:
- args = [x.to(device, non_blocking=True) for x in args]
- with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):
- ps, loss = model(*args)
- N = args[0].shape[0]
- val_loss += loss.mean().item() * N
- val_samples += N
- val_loss = val_loss / val_samples
- with record_function("model.train"):
- model.train()
- with record_function("plotting"):
- visual.add_data(it, scheduler.get_last_lr(), avg_train_loss, val_loss)
- 
- if it >= next_table_it:
- visual.add_table_row(it, avg_train_loss, val_loss)
- next_table_it += table_row_every_iters
-
- it += bs
- visual.on_iter(bar, it, avg_train_loss, val_loss)
- except KeyboardInterrupt:
- mb.write(f"interrupted")
- mb.show()
- pass
- finally:
- visual.add_table_row(it, avg_train_loss, val_loss)
- mb.show()
- visual.hide()

whisperspeech/train_multi.py

@@ -1,263 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B2. Training (Lightning).ipynb.
-
-# %% auto 0
-__all__ = []
-
-# %% ../nbs/B2. Training (Lightning).ipynb 2
-import io
-import time
-import random
-from pathlib import Path
-
-from fastprogress import progress_bar, master_bar
-import fastprogress
-import wandb
-
-import numpy as np
-import pylab as plt
-
-import torch
-import torch.nn as nn
-from torch.utils.data.dataloader import DataLoader
-from torch.profiler import record_function
-
-# %% ../nbs/B2. Training (Lightning).ipynb 3
-import lightning.pytorch as pl
-import math
-
-class TrainingTask(pl.LightningModule):
- def __init__(self, model, model_hparams=None):
- super().__init__()
- self.model = model
- self.model_hparams = model_hparams
- 
- def on_fit_start(self):
- if getattr(self.model, 'setup'):
- self.model.setup(self.device)
- 
- def configure_optimizers(self):
- """ Initialize AdamW optimizer"""
- lr = self.model_hparams['lr0']
- weight_decay = self.model_hparams['weight_decay']
- 
- all_params = set(model.parameters())
- customized_params = set()
- groups = []
- group_map = {}
- for name,m in model.named_modules():
- if hasattr(m, 'no_weight_decay') or hasattr(m, 'lr_scale'):
- customized_params |= set(m.parameters())
- m_wd = 0 if hasattr(m, 'no_weight_decay') else weight_decay
- m_lr = lr * getattr(m, 'lr_scale', 1)
- group = group_map.get((m_wd, m_lr), None)
- if not group:
- group = {"params": [], "names": [], "weight_decay": m_wd, "lr": m_lr}
- groups.append(group)
- group_map[(m_wd, m_lr)] = group
- group['params'] += m.parameters()
- group['names'].append(name)
- 
- other_params = all_params - customized_params
- 
- param_groups = groups + [
- {"names": ["other"], "params": list(other_params), "weight_decay": weight_decay },
- ]
-
- optimizer = torch.optim.AdamW(lr=lr, betas=(0.9, 0.95), params=param_groups)
- 
- # modified from https://github.com/Lightning-AI/lightning/issues/5449#issuecomment-1501597319
- def num_steps_per_epoch() -> int:
- """Get number of steps"""
- # Accessing _data_source is flaky and might break
- dataset = self.trainer.fit_loop._data_source.dataloader()
- dataset_size = len(dataset)
- # math.ceil so always overestimate (underestimating throws exceptions)
- num_steps = math.ceil(dataset_size / self.trainer.accumulate_grad_batches)
- return num_steps
- 
- total_steps = self.model_hparams['epochs'] * num_steps_per_epoch()
- self.model_hparams['pct_start'] = min(0.3, self.model_hparams['warmup_steps'] / total_steps)
-
- print(f"{self.model_hparams['epochs']=} epochs x {num_steps_per_epoch()=} steps")
- 
- lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
- optimizer,
- pct_start=self.model_hparams['pct_start'],
- max_lr=[pg.get('lr', lr) for pg in param_groups],
- steps_per_epoch=num_steps_per_epoch(),
- epochs=int(self.model_hparams['epochs']),
- final_div_factor=25
- )
-
- return [optimizer], [{'scheduler': lr_scheduler, 'interval': 'step'}]
- 
- def training_step(self, train_batch, batch_idx):
- train_logits, train_loss = self.model.forward(*train_batch)
-
- self.log("train_loss", train_loss, sync_dist=True)
- return train_loss
- 
- def validation_step(self, val_batch, batch_idx):
- val_logits, val_loss = self.model.forward(*val_batch)
-
- self.log("val_loss", val_loss, sync_dist=True)
- return val_loss
-
- def on_validation_epoch_end(self):
- if hasattr(self.model, 'get_metrics'):
- self.log_dict({'metrics/'+k:v for k,v in self.model.get_metrics().items()}, sync_dist=True)
- 
- def test_step(self, val_batch, batch_idx):
- test_logits, test_loss = self.model.forward(*val_batch)
-
- self.log("test_loss", test_loss, sync_dist=True)
- return test_loss
-
-# %% ../nbs/B2. Training (Lightning).ipynb 4
-from fastcore.script import anno_parser
-import shlex
-
-# watch out: we can only pass Python values as keyword arguments (not positional)
-# everything else has to be a string
-def parse_and_call(name, fun, args, kwargs={}, log_to_wandb=True):
- p = anno_parser(fun)
- args = p.parse_args(args).__dict__
- args.pop('xtra'); args.pop('pdb')
- args.update({k:v for k, v in kwargs.items()})
- if log_to_wandb and type(wandb_logger.experiment.config) == wandb.sdk.wandb_config.Config:
- wandb_logger.experiment.config[name] = {k:v for k,v in args.items() if k not in ['dataset', 'tunables']}
- return fun(**args)
-
-# %% ../nbs/B2. Training (Lightning).ipynb 8
-import argparse
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--task', type=str, help='Task to train')
-parser.add_argument('--seed', type=int, default=0, help='Global training seed')
-parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
-parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
-parser.add_argument('--input-dir', type=str, default='', help='input data path') # fixed in the model for now
-parser.add_argument("--checkpoint-dir", type=str, default="./checkpoints/", help="directory to save the checkpoints")
-parser.add_argument('--epochs', type=int, default=10, help='total training epochs')
-parser.add_argument('--validate-every-n-steps', type=int, default=500, help='how training steps to run between validations')
-parser.add_argument('--weight-decay', type=float, default=1e-2, help='optimizer weight decay')
-parser.add_argument('--lr0', type=float, default=1e-4, help='optimizer initial learning rate')
-parser.add_argument('--clip-gradient-norm', type=float, default=None, help='enable gradient norm clipping')
-parser.add_argument('--accumulate-grad-batches', type=int, default=1, help='perform the optimizer step only after going through several batches of samples')
-parser.add_argument('--precision', type=str, default="16-mixed", help="floating point precision")
-parser.add_argument('--warmup-steps', type=int, default=10000, help='total number steps during which the learning rate rises (defaults to 10k updates)')
-parser.add_argument('--tunables', type=str, default="", help='tunable hyperparameters')
-parser.add_argument('--resume-from', type=Path, default=None, help='resume training from the given checkpoint')
-parser.add_argument('--strategy', type=str, default='ddp', help='distributed training strategy')
-parser.add_argument('--wandb-suffix', type=str, default=None, help='W&B project name suffix')
-parser.add_argument('--wandb-task-name', type=str, default=None, help='Task name for the W&B project name')
-
-args = parser.parse_args().__dict__
-
-task_args: list = shlex.split(args.pop("task"))
-task_name, task_args = task_args[0], task_args[1:]
-input_args: list = shlex.split(args.pop("input_dir"))
-checkpoint_dir: str = args.pop("checkpoint_dir")
-num_workers: int = args.pop("workers")
-batch_size: int = args.pop("batch_size")
-epochs: int = args.pop("epochs")
-tunables_args: list = shlex.split(args.pop("tunables"))
-
-hyp_params = {}
-hyp_params['batch_size'] = batch_size
-hyp_params['warmup_steps'] = args['warmup_steps']
-hyp_params['weight_decay'] = args['weight_decay']
-hyp_params['clip_gradient_norm'] = args['clip_gradient_norm']
-hyp_params['accumulate_grad_batches'] = args['accumulate_grad_batches']
-hyp_params['precision'] = args['precision']
-hyp_params['lr0'] = args['lr0']
-hyp_params['epochs'] = epochs
-hyp_params['strategy'] = args['strategy']
-
-# %% ../nbs/B2. Training (Lightning).ipynb 9
-from lightning.pytorch.loggers import WandbLogger
-from lightning.pytorch.callbacks import LearningRateMonitor
-import datetime
-import webdataset as wds
-import importlib
-
-torch.set_float32_matmul_precision('medium')
-
-project = f"WhisperSpeech-{args['wandb_task_name'] or task_name}"
-if args['wandb_suffix']:
- project += "-"+args['wandb_suffix']
-
-wandb_logger = WandbLogger(project=project)
-
-ckpt_callback = pl.callbacks.ModelCheckpoint(
- dirpath=f'{task_name}-{epochs}e',
- filename=task_name+"-{epoch}-{step}-{val_loss:.2f}",
- monitor="val_loss",
- save_top_k=4,
- train_time_interval=datetime.timedelta(minutes=5),
- )
-
-lr_monitor_callback = LearningRateMonitor(logging_interval='step')
-
-from torch.utils.data import DataLoader
-
-task = importlib.import_module("whisperspeech."+task_name)
-
-train_ds, val_ds = parse_and_call('dataset', task.load_datasets, input_args)
-
-tunables = None
-if hasattr(task, "Tunables"):
- import dataclasses
- tunables = parse_and_call('tunables', task.Tunables, tunables_args, log_to_wandb=False)
- if type(wandb_logger.experiment.config) == wandb.sdk.wandb_config.Config:
- wandb_logger.experiment.config['tunables'] = dataclasses.asdict(tunables)
-
- for name in ["lr0", "clip_gradient_norm", "weight_decay", "warmup_steps"]:
- val = getattr(tunables, name, None)
- if val is not None: hyp_params[name] = val
-
-if isinstance(train_ds, torch.utils.data.IterableDataset):
- dl_batch_size, dl_shuffle = None, False
- pin_memory = False
-else:
- dl_batch_size, dl_shuffle = batch_size, True
- pin_memory = True
-
-val_loader = wds.WebLoader(val_ds,
- batch_size=dl_batch_size,
- num_workers=num_workers,
- drop_last=False,
- pin_memory=pin_memory).unbatched().shuffle(1024).batched(batch_size).with_length(val_ds.total_samples // batch_size)
-
-train_loader = wds.WebLoader(train_ds,
- batch_size=dl_batch_size,
- num_workers=num_workers,
- drop_last=False,
- shuffle=dl_shuffle,
- pin_memory=pin_memory).unbatched().shuffle(1024).batched(batch_size).with_length(train_ds.total_samples // batch_size)
-
-model_kwargs = dict(dataset=train_ds)
-if tunables is not None: model_kwargs['tunables'] = tunables
-model = parse_and_call('model', task.make_model, task_args, model_kwargs)
-
-task = TrainingTask(model, model_hparams=hyp_params)
-
-trainer = pl.Trainer(strategy=hyp_params['strategy'],
- max_epochs=hyp_params['epochs'],
- accelerator="gpu",
- profiler="simple",
- precision=hyp_params['precision'],
- gradient_clip_val=hyp_params['clip_gradient_norm'],
- accumulate_grad_batches=hyp_params['accumulate_grad_batches'],
- val_check_interval=args.pop("validate_every_n_steps"),
- enable_checkpointing=True,
- logger=wandb_logger,
- callbacks=[ckpt_callback, lr_monitor_callback])
-
-if type(wandb_logger.experiment.config) == wandb.sdk.wandb_config.Config:
- wandb_logger.experiment.config.update(hyp_params)
- 
-kwargs = {}
-if 'resume_from' in args:
- kwargs['ckpt_path'] = args['resume_from']
-trainer.fit(model=task, train_dataloaders=train_loader, val_dataloaders=val_loader, **kwargs)

whisperspeech/utils.py

@@ -1,159 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/D. Common dataset utilities.ipynb.
-
-# %% auto 0
-__all__ = ['shard_glob', 'join_datasets', 'resampler', 'derived_name', 'derived_dataset', 'merge_in', 'AtomicTarWriter',
- 'readlines']
-
-# %% ../nbs/D. Common dataset utilities.ipynb 1
-import os
-import torch
-import torchaudio
-from pathlib import Path
-import webdataset as wds
-from contextlib import contextmanager
-
-import torch.nn.functional as F
-
-# %% ../nbs/D. Common dataset utilities.ipynb 2
-def shard_glob(input):
- if '{' in input:
- return wds.shardlists.expand_urls(input)
- if isinstance(input, (Path, str)):
- path = Path(input)
- if path.is_dir():
- glob = '*.tar.gz'
- else:
- glob = path.name
- path = path.parent
- input = Path(path).glob(glob)
- else:
- raise ArgumentError("input should be either a list or a path with an optional glob specifier")
- return [str(x) for x in input]
-
-# %% ../nbs/D. Common dataset utilities.ipynb 3
-class join_datasets(torch.utils.data.IterableDataset):
- def __init__(self, datasets):
- self.datasets = datasets
- 
- def __iter__(self):
- probs = torch.tensor([getattr(ds, 'weight', 1) for ds in self.datasets], dtype=torch.float)
- its = [iter(ds) for ds in self.datasets]
- while True:
- try:
- yield next(its[torch.multinomial(probs, 1)])
- except StopIteration:
- return 
- 
- def __len__(self):
- return sum([ds.total_samples for ds in self.datasets])
-
-# %% ../nbs/D. Common dataset utilities.ipynb 5
-def resampler(newsr = 24000, key = 'samples_24k'):
- _last_sr = None
- tform = None
- 
- def _resample(samples):
- for s in samples:
- sr = s['sample_rate']
- if sr != newsr:
- if sr != _last_sr: tform = torchaudio.transforms.Resample(sr, newsr)
- s[key] = tform(s['samples'])
- else:
- s[key] = s['samples']
- yield s
- 
- return _resample
-
-# %% ../nbs/D. Common dataset utilities.ipynb 6
-def derived_name(input, kind, base="audio", suffix=".gz", dir=None):
- dir = Path(dir) if dir else Path(input).parent
- return str(dir/(Path(input).name.replace(f"-{base}-", f"-{kind}-") + suffix))
-
-# %% ../nbs/D. Common dataset utilities.ipynb 7
-def derived_dataset(kind, base='audio', suffix=".gz", decoders=[], dir=None):
- def deriver(url):
- url = str(derived_name(url, kind, base=base, suffix=suffix, dir=dir))
- return wds.WebDataset(
- wds.SimpleShardList([url])
- ).decode(*decoders)
- return deriver
-
-# %% ../nbs/D. Common dataset utilities.ipynb 8
-def merge_in(dataset_fun):
- """Merge a dataset into the current one returning samples with the union of keys. Pass in a function
- that takes a URL of a sample and returns a dataset for it (called everytime the URL changes).
- 
- It requires (and validates) that both datasets have the same ordering of keys so you have
- to use it before any sample shuffling. Shard shuffling is ok.
- """
- def merge_loop(main_samples):
- #print("new merge loop:", dataset_fun)
- merged_samples = None
- cur_url = None
- i = None
- for s in main_samples:
- url = s['__url__']
- if url != cur_url:
- # this will open a new file when we get the first sample with a new __url__
- merged_samples = iter(dataset_fun(url))
- cur_url = url
- try:
- merge_s = next(merged_samples)
- except StopIteration:
- # if the original shard got repeated we won't observe a __url__ change
- # in this case restart the dataset from the beginning
- merged_samples = iter(dataset_fun(url))
- merge_s = next(merged_samples)
- assert merge_s['__key__'] == s['__key__'], f"sample keys don't match: {merge_s['__key__']}, {s['__key__']} in file {s['__url__']}"
- news = {}
- news.update(merge_s)
- news.update(s)
- yield news
- return merge_loop
-
-# %% ../nbs/D. Common dataset utilities.ipynb 9
-def split_to_chunks(stream, ikey='vad.npy', metakeys=[], pad_to_seconds=30, random_shift=False):
- for s in stream:
- audio, sr = s['audio']
- imax = len(s[ikey]) - 1
- for i,(ts,te) in enumerate(s[ikey]):
- samples = audio[0,int(ts*sr):int(te*sr)]
- if pad_to_seconds is not None:
- padding = pad_to_seconds*sr-samples.shape[-1]
- lpad = random.randint(0, padding) if random_shift else 0
- samples = F.pad(samples, (lpad, padding-lpad))
- subs = {"__key__": s['__key__'] + f"_{i:03d}",
- "src_key": s['__key__'],
- "__url__": s['__url__'],
- "i": i, "imax": imax,
- "tstart": ts, "tend": te, "total_seconds": audio.shape[-1]/sr,
- "lpad": lpad, "rpad": padding-lpad,
- "lpad_s": lpad/sr, "rpad_s": (padding-lpad)/sr,
- "samples": samples, "sample_rate": sr}
- for k in metakeys:
- subs[k] = s[k][i]
- yield subs
-
-# %% ../nbs/D. Common dataset utilities.ipynb 10
-def vad_dataset(shards, ikey='vad.npy', kind='vad'):
- return wds.WebDataset(shards).compose(
- wds.decode(wds.torch_audio),
- merge_in(derived_dataset(kind)),
- wds.select(lambda x: 'wav' in x or 'flac' in x or 'mp3' in x or 'ogg' in x), # skip samples without audio
- wds.rename(audio="flac;mp3;wav;ogg"),
- lambda x: split_to_chunks(x, ikey=ikey),
- )
-
-# %% ../nbs/D. Common dataset utilities.ipynb 11
-@contextmanager
-def AtomicTarWriter(name, throwaway=False):
- tmp = name+".tmp"
- with wds.TarWriter(tmp, compress=name.endswith('gz')) as sink:
- yield sink
- if not throwaway:
- os.rename(tmp, name)
-
-# %% ../nbs/D. Common dataset utilities.ipynb 12
-def readlines(fname):
- with open(fname) as file:
- return [line.rstrip() for line in file]

whisperspeech/vad.py

@@ -1,71 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/1B. Voice activity detection.ipynb.
-
-# %% auto 0
-__all__ = []
-
-# %% ../nbs/1B. Voice activity detection.ipynb 3
-import os
-import torch
-import torchaudio
-
-from pathlib import Path
-from fastprogress import progress_bar
-from fastcore.script import call_parse
-
-import whisperx
-import random
-import numpy as np
-import webdataset as wds
-
-# %% ../nbs/1B. Voice activity detection.ipynb 5
-# some of the original file names have a dot in their name
-# webdataset does not like it so let's patch it
-def fix_dots_in_names(name):
- name, ext = name.rsplit('.', 1)
- return ".".join((name.replace('.', '_'), ext))
-
-def load_dataset(url, decode=True, rename_files=None):
- ds = wds.WebDataset(url, rename_files=rename_files)
- if not decode: return ds
- return ds.decode(wds.torch_audio)
-
-# %% ../nbs/1B. Voice activity detection.ipynb 7
-def extract_segments(vad_result, max_duration):
- binarize = whisperx.vad.Binarize(max_duration=max_duration)
- segments = binarize(vad_result)
- return [(x.start, x.end) for x in segments.get_timeline()]
-
-def segment_audio(vad_model, audio, sr=16000):
- vad_result = vad_model({"waveform": audio, "sample_rate": sr})
- return extract_segments(vad_result, 30)
-
-# %% ../nbs/1B. Voice activity detection.ipynb 13
-def flac_to_vad_name(input):
- if '-flac-' in input:
- return input.rsplit("/", 1)[1].replace('flac', 'vad') + ".gz"
- else:
- return input.rsplit("/", 1)[1].replace('raw', 'vad') + ".gz"
-
-@call_parse
-def process_shard(
- input:str, # input shard URL/path
- output:str=None, # output shard URL/path
- fix_dots:bool=False, # fix dots in LibriLight filenames
-):
- if output is None: output = flac_to_vad_name(input)
- 
- ds = torch.utils.data.DataLoader(load_dataset(input, rename_files=fix_dots_in_names if fix_dots else None), num_workers=2, batch_size=None)
- vad_model = whisperx.vad.load_vad_model('cuda')
- 
- tmp = output+".tmp"
- with wds.TarWriter(tmp) as sink:
- for s in progress_bar(ds, total='noinfer'):
- audio, sr = s.get('flac', s.get('wav', (None, None)))
- if audio is None:
- print(f"warning: '{s['__key__']}' does not contain an audio file")
- continue
- sink.write({
- "__key__": s['__key__'],
- "vad.npy": np.array(segment_audio(vad_model, audio, sr=sr), dtype=np.float16)
- })
- os.rename(tmp, output)

whisperspeech/vq_stoks.py

@@ -1,493 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2B. Whisper quantization (semantic token) model.ipynb.
-
-# %% auto 0
-__all__ = ['RQBottleneckTransformer', 'make_model']
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 2
-import io
-import sys
-import time
-import torch
-import torchaudio
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 3
-from pathlib import Path
-import json
-from fastprogress import progress_bar, master_bar
-import fastprogress
-import numpy as np
-import pylab as plt
-import pandas as pd
-import random
-
-import whisper
-from huggingface_hub import hf_hub_download
-from fastcore.basics import store_attr
-
-from torch import nn
-import torch.optim as optim
-import torch.nn.functional as F
-from torch.utils.data.dataloader import DataLoader
-import webdataset as wds
-from . import utils
-
-from vector_quantize_pytorch import ResidualVQ
-
-from fastcore.script import *
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 9
-def merge_in(dataset_fun):
- """Merge a dataset into the current one returning samples with the union of keys. Pass in a function
- that takes a URL of a sample and returns a dataset for it (called everytime the URL changes).
- 
- It requires (and validates) that both datasets have the same ordering of keys so you have
- to use it before any sample shuffling. Shard shuffling is ok.
- """
- def merge_loop(main_samples):
- #print("new merge loop:", dataset_fun)
- merged_samples = None
- cur_url = None
- i = None
- for s in main_samples:
- url = s['__url__']
- if url != cur_url:
- # this will open a new file when we get the first sample with a new __url__
- merged_samples = iter(dataset_fun(url))
- cur_url = url
- try:
- merge_s = next(merged_samples)
- except StopIteration:
- # if the original shard got repeated we won't observe a __url__ change
- # in this case restart the dataset from the beginning
- merged_samples = iter(dataset_fun(url))
- merge_s = next(merged_samples)
- assert merge_s['__key__'] == s['__key__'], f"sample keys don't match: {merge_s['__key__']}, {s['__key__']} in file {s['__url__']}"
- news = {}
- news.update(merge_s)
- news.update(s)
- yield news
- return merge_loop
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 10
-def derived_dataset(kind, key='audio'):
- def deriver(url):
- url = str(Path(url).parent/(Path(url).name.replace(key, kind) + ".gz"))
- return wds.WebDataset(
- wds.SimpleShardList([url])
- ).decode()
- return deriver
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 17
-def add_masks(samples):
- for s in samples:
- seconds = s['tend'] - s['tstart']
- # a mask (downsampled to the Whisper encoder token rate of 50/s) is used
- # to teach the model the concept of padding
- # this let's us decode shorter sequences later
- mask = torch.zeros(30*16000//320, dtype=torch.bool)
- mask[:int(seconds * 16000) // 320] = 1
- s['mask'] = mask
- yield s
-
-def tokenize_text(samples, ttoks_size=200, model="base.en", language="en"):
- multilingual = not model.endswith(".en")
- tokenizer = whisper.tokenizer.get_tokenizer(multilingual, language=language, task="transcribe")
- for s in samples:
- ttoks = tokenizer.encode(s['txt'])
- tokens = list(tokenizer.sot_sequence) + ttoks
- rpad = ttoks_size - len(tokens)
- s['in_ttoks'] = F.pad(torch.tensor(tokens), (0, rpad), value=tokenizer.eot)
- s['out_ttoks'] = F.pad(torch.tensor(tokens[1:] + [tokenizer.eot]), (0, rpad), value=-100)
- yield s
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 22
-def load_dataset(
- shard_spec:str,
- proc_dataset_path:Path, # processed VAD and txt files
- samples:int, # set the per-GPU sample count
- txt_label:str="base.en-txt", # the label of the files containing transcriptions
- model:str="base.en",
- key:str="flac",
- language:str=None,
- validation:bool=False, 
- ):
- from . import wh_transcribe
- shards = utils.shard_glob(shard_spec)
- 
- if not language and model.endswith('en'): language = 'en'
- assert language, "please provide the dataset language for multilang models"
- 
- same_on_all_nodes = lambda urls: urls # will only be used for validation
- ds = wds.WebDataset(shards, resampled=not validation, nodesplitter=same_on_all_nodes).compose(
- wds.decode(wds.torch_audio),
- wds.select(lambda x: 'wav' in x or 'flac' in x or 'mp3' in x or 'ogg' in x), # skip samples without audio
- wds.rename(audio="flac;mp3;wav;ogg"),
- merge_in(derived_dataset(proc_dataset_path, 'vad', key=key)),
- wds.map_dict(**{"vad.npy":wh_transcribe.chunk_merger}),
- wh_transcribe.split_to_chunks,
- utils.resampler(16000, 'samples_16k'),
- merge_in(derived_dataset(proc_dataset_path, txt_label, key=key)),
- )
- if 'librilight' in shards[0]:
- ds = ds.compose(
- # drop the first and last segment because they tend to be inaccurate
- # (the transcriptions don't have the "LibriVox" headers and "end of chapter" suffixes)
- wds.select(lambda x: x['i'] != 0 and x['i'] != x['imax']),
- )
- ds = ds.compose(
- add_masks,
- lambda x: tokenize_text(x, model=model, language=language),
- wds.to_tuple('samples_16k', 'mask', 'in_ttoks', 'out_ttoks'),
- wds.batched(32),
- )
- ds.total_samples = samples
- 
- return ds
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 28
-from whisperspeech.train import *
-from whisperspeech.modules import *
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 29
-import dataclasses
-
-def rand(start, end):
- return random.random() * (end - start) + start
-
-def logrand(start, end):
- return 10**rand(math.log10(start), math.log10(end))
-
-@dataclasses.dataclass
-class Tunables:
- init_std :float = 1.5
- embeddings_std :float = 4.5e-2
- embeddings_lr_scale: float = 1
- output_mult :float = 1
- query_mult :float = 2
- rope :bool = True
- mask_embs :bool = True # force embeddings corresponding to the input audio padding to a constant value
- downsample_conv: bool = False
- downsample_mean: bool = True
- 
- codebook_dim: int = 32
- codebook_decay: float = 0.9
- 
- lr0 :float = .9e-3
- clip_gradient_norm :float = 2
- weight_decay :float = 1e-3
- warmup_steps :float = 850
-
- random :bool = False
-
- def __post_init__(self):
- # randomize the hyperparams if requested
- if self.random:
- self.init_std = logrand(1, 2)
- self.embeddings_std = logrand(3e-2,6e-2)
- self.embeddings_lr_scale = 2**rand(0,3)
- self.output_mult = 2**rand(-3,3)
- self.query_mult = logrand(1,8)
- self.codebook_dim = int(logrand(30,50))
- self.codebook_decay = logrand(0.86,0.95)
- self.rope = True
- self.mask_embs = True
- self.downsample_mean = True
- 
- self.lr0 = logrand(.8e-3,1e-3)
- self.clip_gradient_norm = 10**rand(-1,1)
- self.warmup_steps = logrand(700,1000)
- 
- @staticmethod
- def upgrade(args):
- args = {k:v for k,v in args.items()}
- def old_default(name, value):
- if name not in args: args[name] = value
- old_default('output_mult', 1)
- old_default('query_mult', 1)
- old_default('rope', False)
- old_default('mask_embs', False)
- old_default('downsample_conv', False)
- old_default('downsample_mean', False)
- if 'encoder_depth_ratio' in args: del args['encoder_depth_ratio']
- if 'vq_codes' in args: del args['vq_codes']
- return args
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 30
-import math
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 31
-class RQBottleneckTransformer(nn.Module):
- def __init__(self, vq_codes=512, q_depth=12, depth=1, n_head=2, head_width=64, ffn_mult=4,
- codebook_dim=2, threshold_ema_dead_code=2, use_cosine_sim = False, kl_loss_mul=1,
- downsample=1,
- whisper_model_name='tiny.en', tunables=Tunables()):
- super().__init__()
- width = n_head * head_width
- store_attr("codebook_dim,vq_codes,q_depth,n_head,head_width,ffn_mult,depth,use_cosine_sim,downsample,whisper_model_name")
- self.width = width
- self.base_width = 3 * head_width
- self.vq_codes = vq_codes
- self.tunables = tunables
- self.stoks_len = 1500//downsample
- self.stoks_per_sec = self.stoks_len//30
- 
- qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)
- 
- self.kl_loss_mul = kl_loss_mul
- 
- n_mlp = width * ffn_mult
- self.mlp = nn.Sequential(
- nn.Linear(width, n_mlp), nn.GELU(), nn.Linear(n_mlp, width)
- )
- self.mlp_ln = LayerNorm(width)
-
- if tunables.downsample_conv:
- self.downsample_conv = nn.Conv1d(width, width, kernel_size=3, stride=downsample, padding=1)
- else:
- self.downsample_conv = None
- 
- if tunables.mask_embs: vq_codes = vq_codes + 1
- self.rq = ResidualVQ(
- dim = width,
- codebook_size = vq_codes, # codebook size
- decay = tunables.codebook_decay, # the exponential moving average decay, lower means the dictionary will change faster
- commitment_weight = 1., # the weight on the commitment loss
- threshold_ema_dead_code = threshold_ema_dead_code,
- use_cosine_sim = use_cosine_sim,
- codebook_dim = codebook_dim,
- num_quantizers= 1,
- )
- 
- self.ce_lossf = nn.CrossEntropyLoss(ignore_index=-100)
- self.kl_lossf = nn.KLDivLoss(reduction='batchmean')
-
- self.positional_embedding = nn.Embedding(1500, width) # FIXME: should be self.stoks_len
- 
- self.out_blocks = nn.Sequential(*[
- ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(depth)
- ])
- self.ln_post = LayerNorm(width)
- 
- self.whmodel = None
-
- self.apply(self.init_transformer)
- self.register_buffer('val_true', torch.zeros(1).cuda())
- self.register_buffer('val_total', torch.zeros(1).cuda())
- 
- def setup(self, device):
- self.ensure_whisper(device)
- 
- def init_transformer(self, m):
- if isinstance(m, LinearHead):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, QueryHead):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- torch.nn.init.constant_(m.weight, 0)
- elif isinstance(m, nn.Embedding):
- m.no_weight_decay = True
- m.lr_scale = self.tunables.embeddings_lr_scale
- std = self.tunables.embeddings_std
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.Linear):
- m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
- std = self.tunables.init_std / m.weight.shape[1]
- torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
- if m.bias is not None:
- torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)
- elif isinstance(m, nn.LayerNorm):
- m.no_weight_decay = True
- torch.nn.init.constant_(m.bias, 0)
- torch.nn.init.constant_(m.weight, 1)
-
- @property
- def device(self):
- return next(self.parameters()).device
- 
- #
- # training
- #
- @torch.no_grad()
- def extract_teacher(self, samples, input_toks, output_toks):
- embs = self.whmodel[0].encoder(whisper.log_mel_spectrogram(samples))
- teacher_logits = self.whmodel[0].decoder(input_toks, embs)
- # set teacher logits to 0 for padding positions so KLDivLoss ignores them
- teacher_logits[output_toks == -100] = 0
- return embs, teacher_logits
- 
- def downsample_embeddings(self, x):
- if self.downsample_conv is not None:
- return x[:,::self.downsample] + self.downsample_conv(x.transpose(-1,-2)).transpose(-2,-1)
- elif self.tunables.downsample_mean:
- bs,slen,depth = x.shape
- return x.reshape(bs,slen//self.downsample,self.downsample,depth).mean(-2)
- else:
- return x[:,::self.downsample]
- 
- def forward(self, samples, mask, input_toks, output_toks):
- embs, teacher_logits = self.extract_teacher(samples, input_toks, output_toks)
- 
- x = self.downsample_embeddings(embs)
- x = x + self.mlp(self.mlp_ln(x))
- # VQ bottleneck
- quantized, self.indices, self.commit_loss = self.rq(x)
- self.commit_loss = self.commit_loss.mean()
-
- x = quantized.repeat_interleave(self.downsample, -2)
- project_out = getattr(self.rq, 'project_out', None) or self.rq.layers[0].project_out
- if self.tunables.mask_embs: x[~mask] = project_out(self.rq.layers[0]._codebook.embed[0,self.vq_codes])
- positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)
- x = x + self.positional_embedding(positions)
- x = self.ln_post(self.out_blocks(x))
- 
- logits = self.whmodel[0].decoder(input_toks, x)
- self.ce_loss = self.ce_lossf(logits.view(-1,logits.shape[-1]), output_toks.view(-1))
- self.kl_loss = self.kl_lossf(F.log_softmax(logits, dim=-1), F.softmax(teacher_logits, dim=-1))
- loss = self.ce_loss + self.kl_loss_mul * self.kl_loss + self.commit_loss
- 
- if not self.training:
- valid_toks = output_toks != -100
- self.val_true += (logits.argmax(-1)[valid_toks] == output_toks[valid_toks]).float().sum()
- self.val_total += valid_toks.float().sum()
-
- return x, loss
- 
- def get_metrics(self):
- metrics = {
- 'acc_0': (self.val_true / self.val_total).item(),
- }
- self.val_true[:] = 0
- self.val_total[:] = 0
- return metrics
- 
- #
- # inference
- #
- @classmethod
- def load_model(cls, ref="collabora/spear-tts-pytorch:whisper-vq-stoks-medium-en+pl.model",
- repo_id=None, filename=None, local_filename=None):
- if repo_id is None and filename is None and local_filename is None:
- if ":" in ref:
- repo_id, filename = ref.split(":", 1)
- else:
- local_filename = ref
- if not local_filename:
- local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
- spec = torch.load(local_filename) 
- vqmodel = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec.get('tunables', {}))))
- vqmodel.load_state_dict(spec['state_dict'])
- vqmodel.eval()
- return vqmodel
- 
- def load_checkpoint(self, local_filename):
- spec = torch.load(local_filename, map_location='cpu')
- assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'
- state_dict = {k.replace('model.', ''):v
- for k,v in spec['state_dict'].items()}
- self.load_state_dict(state_dict)
- return self
- 
- def save_model(self, fname, store_parameters=True):
- torch.save(dict(config = self.__stored_args__,
- tunables = dataclasses.asdict(self.tunables),
- state_dict = self.state_dict() if store_parameters else None), fname)
- 
- def ensure_whisper(self, device):
- # the list wrapper is a hack to make sure the whole of Whisper is not sucked into self.parameters()
- if self.whmodel is None: self.whmodel = [whisper.load_model(self.whisper_model_name, device=device)]
- self.decoding_options = whisper.DecodingOptions()
- multilingual = not self.whisper_model_name.endswith('.en')
- self.tokenizer = whisper.tokenizer.get_tokenizer(multilingual)
- 
- def quantize(self, embs):
- x = self.downsample_embeddings(embs)
- x = x + self.mlp(self.mlp_ln(x))
- _, stoks, _ = self.rq(x)
- if self.q_depth == 1:
- stoks = stoks.squeeze(-1)
- return stoks
-
- def dequantize(self, stoks):
- assert self.q_depth == 1
- assert len(stoks.shape) == 1, "batch processing is not supported"
- if isinstance(stoks, np.ndarray): stoks = torch.tensor(stoks)
- # remove padding
- padding = torch.nonzero(stoks == self.vq_codes)
- if padding.any(): stoks = stoks[:padding[0,0]]
- stoks = F.pad(stoks, (0,self.stoks_len - stoks.shape[-1]), value=self.vq_codes if self.tunables.mask_embs else 0)
- x = self.rq.layers[0]._codebook.embed[0,stoks.to(torch.long).view(-1)]
- x = x.repeat_interleave(self.downsample, -2)
- project_out = getattr(self.rq, 'project_out', None) or self.rq.layers[0].project_out
- x = project_out(x).unsqueeze(0)
- positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)
- x = x + self.positional_embedding(positions)
- return self.ln_post(self.out_blocks(x))
-
- def encode_audio(self, audio):
- if isinstance(audio, str):
- x, sr = torchaudio.load(audio)
- x = torchaudio.transforms.Resample(sr, 16000)(x)[0]
- audio = x.unsqueeze(0)
- return self.encode_mel(whisper.log_mel_spectrogram(audio).to(self.device))
- 
- def encode_mel(self, mel):
- assert len(mel.shape) == 3, "invalid mel spectrogram shape, expect (batch,chn,time)"
- self.ensure_whisper(self.device)
- n = mel.shape[-1]
- if n > whisper.audio.N_FRAMES:
- padding = 0
- padded = mel[:,:,:whisper.audio.N_FRAMES]
- else:
- padding = -n % whisper.audio.N_FRAMES
- padded = F.pad(mel, (0, padding), value=-1.5)
- embs = self.whmodel[0].encoder(padded)#.to(self.whmodel[0].device))#[:,:n//2]
- stoks = self.quantize(embs)
- if self.tunables.mask_embs:
- return stoks[:,:n//2//self.downsample]
- else:
- return stoks
- 
- def decode_text(self, stoks, decoding_options=None):
- self.ensure_whisper(self.device)
- if decoding_options is None: decoding_options = self.decoding_options
- embs = self.dequantize(stoks).to(self.whmodel[0].device)
- return self.whmodel[0].decode(embs, decoding_options)
-
-# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 33
-def make_model(size:str, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
- if size == 'base.en-2d-4096c':
- model = RQBottleneckTransformer(codebook_dim=32, vq_codes=4096, q_depth=1, n_head=8, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- if size == 'base.en-2d-512c':
- model = RQBottleneckTransformer(codebook_dim=32, vq_codes=512, q_depth=1, n_head=8, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- if size == 'base.en-2d-512c-dim64':
- model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, q_depth=1, n_head=8, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- if size == 'base-2d-512c-dim64':
- model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, q_depth=1, n_head=8, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- if size == 'base-2d-1024c-dim64':
- model = RQBottleneckTransformer(codebook_dim=64, vq_codes=1024, q_depth=1, n_head=8, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- if size == 'medium-2d-512c-dim64':
- model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, q_depth=1, n_head=16, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- if size == 'medium-2d-1024c-dim64':
- model = RQBottleneckTransformer(codebook_dim=64, vq_codes=1024, q_depth=1, n_head=16, depth=1,
- downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
- whisper_model_name=size.split("-")[0], tunables=tunables)
- return model
- raise ArgumentError(f"invalid model size: {size}")

whisperspeech/wer_metrics.py

@@ -1,77 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/C. Word error rate metrics.ipynb.
-
-# %% auto 0
-__all__ = ['librispeech_data', 'DfBuilder', 'WERStats']
-
-# %% ../nbs/C. Word error rate metrics.ipynb 2
-import jiwer
-from whisper_normalizer.english import EnglishTextNormalizer
-
-import torchaudio
-from pathlib import Path
-import pandas as pd
-
-# %% ../nbs/C. Word error rate metrics.ipynb 3
-engnorm = EnglishTextNormalizer()
-def whisper_normalize(x):
- if type(x) == list:
- return [engnorm(y) for y in x]
- else:
- return engnorm(x)
-
-default_transform = jiwer.transforms.Compose([
- jiwer.transforms.ToLowerCase(),
- jiwer.transforms.ExpandCommonEnglishContractions(),
- whisper_normalize,
- jiwer.transforms.RemoveMultipleSpaces(),
- jiwer.transforms.Strip(),
- jiwer.transforms.RemovePunctuation(),
- jiwer.transforms.ReduceToListOfListOfWords(),
-])
-
-# %% ../nbs/C. Word error rate metrics.ipynb 5
-def librispeech_data(datadir, sample_rate=16000):
- for file in Path(datadir).rglob('*.txt'):
- for line in file.read_text().split('\n'):
- if not line: continue
- idx, text = line.split(" ", 1)
- x, sr = torchaudio.load((file.parent/idx).with_suffix('.flac'))
- if sr != sample_rate:
- x = torchaudio.transforms.Resample(sr, self.sample_rate)(x)
- yield x, text
-
-# %% ../nbs/C. Word error rate metrics.ipynb 6
-class DfBuilder:
- def __init__(self):
- self.data = {}
- 
- def push(self, **kwargs):
- for k,v in kwargs.items():
- if k not in self.data:
- self.data[k] = [v]
- else:
- self.data[k].append(v)
- 
- def df(self):
- return pd.DataFrame(self.data)
-
-# %% ../nbs/C. Word error rate metrics.ipynb 7
-class WERStats(DfBuilder):
- def __init__(self, transform=default_transform):
- super().__init__()
- self.reference_transform = transform
- self.hypothesis_transform = transform
- 
- def push_sample(self, snd, gt_text, text, idx=None):
- if snd is not None: self.push(secs = snd.shape[-1]/16000)
- diff = jiwer.process_words(gt_text, text, reference_transform=self.reference_transform, hypothesis_transform=self.hypothesis_transform)
- self.push(
- idx = idx,
- gt_text = gt_text,
- text = text,
- wer = diff.wer,
- mer = diff.mer,
- wil = diff.wil,
- wip = diff.wip,
- )
- return diff

whisperspeech/wh_transcribe.py

@@ -1,146 +0,0 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2A. Whisper quantization dataset preparation.ipynb.
-
-# %% auto 0
-__all__ = []
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 3
-import os
-import io
-import time
-import torch
-import torchaudio
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 4
-from pathlib import Path
-import json
-from fastprogress import progress_bar, master_bar
-import numpy as np
-import random
-
-import whisper
-
-from torch import nn
-import torch.nn.functional as F
-from torch.utils.data.dataloader import DataLoader
-
-from fastcore.script import *
-
-from . import vad
-import webdataset as wds
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 9
-# let's make it a bit more conservative
-# with full 30 second chunks it sometimes misses a small part of the transcript
-def random_cutter(dur):
- if random.random() < 0.5:
- return dur > 28 * (random.random()*0.95+0.05)
- else:
- return dur > 28
-
-def chunk_merger(segments, should_cut=lambda x: x > 28):
- if len(segments) == 0: return segments
- curr_start = segments[0][0]
- curr_end = 0
- merged = []
-
- for ts,te in segments:
- if should_cut(te - curr_start) and curr_end - curr_start > 0:
- merged.append((curr_start, curr_end))
- curr_start = ts
- curr_end = te
- merged.append((curr_start, curr_end))
- return merged
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 18
-def merge_in(*datasets):
- """Merge multiple datasets into the current one returning samples with the union of keys.
- 
- It requires (and validates) all datasets to have the same ordering of keys so you have
- to use it before any sample shuffling. Shard shuffling is ok.
- """
- def merge_loop(main_samples):
- for samples in zip(*[main_samples]+[iter(x) for x in datasets]):
- key = samples[0]['__key__']
- news = {}
- for s in samples:
- assert s['__key__'] == key
- news.update(s)
- yield news
- return merge_loop
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 19
-import copy
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 20
-# a workaround for https://github.com/webdataset/webdataset/issues/297
-# should be possible to use ds.compose here
-def wds_compose(ds, *args):
- ds = copy.copy(ds)
- ds.pipeline = copy.copy(ds.pipeline)
- for f in args:
- ds.append(f)
- return ds
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 24
-def split_to_chunks(stream, pad_to_seconds=30, random_shift=False):
- for s in stream:
- audio, sr = s.get('flac', s.get('wav', (None, None)))
- if audio is None:
- print(f"warning: '{s['__key__']}' does not contain an audio file")
- continue
- imax = len(s['vad.npy']) - 1
- for i,(ts,te) in enumerate(s['vad.npy']):
- samples = audio[0,int(ts*sr):int(te*sr)]
- if pad_to_seconds is not None:
- padding = pad_to_seconds*sr-samples.shape[-1]
- lpad = random.randint(0, padding) if random_shift else 0
- samples = F.pad(samples, (lpad, padding-lpad))
- yield {"__key__": s['__key__'] + f"_{i:03d}",
- "__url__": s['__url__'],
- "i": i, "imax": imax,
- "tstart": ts, "tend": te, "total_seconds": audio.shape[-1]/sr,
- "lpad": lpad, "rpad": padding-lpad,
- "lpad_s": lpad/sr, "rpad_s": (padding-lpad)/sr,
- "samples": samples, "sample_rate": sr}
-
-# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 38
-def flac_to_txt_name(input, model_size):
- return input.rsplit("/", 1)[1].replace('flac', f'{model_size}-txt') + ".gz"
-
-@call_parse
-def process_shard(
- input:str, # input shard URL/path
- output:str=None, # output shard URL/path
- bs:int=None, # batch size (16 uses around 11GB of VRAM)
- n_samples:int=None, # limit the number of samples (useful for quick benchmarking)
- whisper_model:str="base.en" # Whisper model size
-):
- if output is None: output = flac_to_txt_name(input, whisper_model)
- if bs is None: bs = 16
- if n_samples is None: n_samples = 'noinfer'
- else: n_samples = n_samples // bs
-
- ds = wds_compose(vad.load_dataset(input),
- merge_in(wds.WebDataset(vad.flac_to_vad_name(input)).decode()),
- wds.map_dict(**{"vad.npy":chunk_merger}),
- split_to_chunks,
- wds.to_tuple('__key__', 'samples'),
- wds.batched(bs),
- )
- dl = DataLoader(ds, num_workers=2, batch_size=None)
- 
- whmodel = whisper.load_model(whisper_model)
- decoding_options = whisper.DecodingOptions(language='en')
- 
- tmp = output+".tmp"
- with wds.TarWriter(tmp) as sink:
- for keys, samples in progress_bar(dl, total=n_samples):
- with torch.no_grad():
- embs = whmodel.encoder(whisper.log_mel_spectrogram(samples).cuda())
- decs = whmodel.decode(embs, decoding_options)
- for key, dec in zip(keys, decs):
- sink.write({
- "__key__": key,
- "txt": dec.text,
- })
- os.rename(tmp, output)