OpenMOSS-Team/SpeechTokenizer

1# SpeechTokenizer: Unified Speech Tokenizer for Speech Large Language Models2 3<a href='https://github.com/ZhangXInFD/SpeechTokenizer'><img src='https://img.shields.io/badge/Project-Page-Green'></a>  <a href='https://arxiv.org/abs/2308.16692'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>4 5## Introduction6This is the code for the SpeechTokenizer presented in the [SpeechTokenizer: Unified Speech Tokenizer for Speech Large Language Models](https://arxiv.org/abs/2308.16692). SpeechTokenizer is a unified speech tokenizer for speech large language models, which adopts the Encoder-Decoder architecture with residual vector quantization (RVQ). Unifying semantic and acoustic tokens, SpeechTokenizer disentangles different aspects of speech information hierarchically across different RVQ layers. Specifically, The code indices that the first quantizer of RVQ outputs can be considered as semantic tokens and the output of the remaining quantizers can be regarded as acoustic tokens, which serve as supplements for the information lost by the first quantizer. We provide our models:7* A model operated at 16khz on monophonic speech trained on Librispeech with average representation across all HuBERT layers as semantic teacher.8 9<br>10<p align="center">11    <img src="images/overview.png" width="95%"> <br>12    Overview13</p>14<p align="center">15    <img src="images/speechtokenizer_framework.jpg" width="95%"> <br>16    The SpeechTokenizer framework.17</p>18<br>19 20 21Welcome to try our [SLMTokBench](https://github.com/0nutation/SLMTokBench) 22 and we will also open source our  [USLM](https://github.com/0nutation/USLM) !!23 24 25 26## Samples27 28Samples are provided on [our demo page](https://0nutation.github.io/SpeechTokenizer.github.io/).29 30## Installation31 32SpeechTokenizer requires Python>=3.8, and a reasonly recent version of PyTorch.33To install SpeechTokenizer, you can run from this repository:34```bash35pip install -U speechtokenizer36 37# or you can clone the repo and install locally38git clone https://github.com/ZhangXInFD/SpeechTokenizer.git39cd SpeechTokenizer40pip install .41```42## Usage43### Model storage44| Model |Discription|45|:----|:----|46|[speechtokenizer_hubert_avg](https://huggingface.co/fnlp/SpeechTokenizer/tree/main/speechtokenizer_hubert_avg)|Adopt average representation across all HuBERT layers as semantic teacher |47 48### load model49```python50from speechtokenizer import SpeechTokenizer51 52config_path = '/path/config.json'53ckpt_path = '/path/SpeechTokenizer.pt'54model = SpeechTokenizer.load_from_checkpoint(config_path, ckpt_path)55model.eval()56```57### Extracting discrete representions58```python59import torchaudio60import torch61 62# Load and pre-process speech waveform63wav, sr = torchaudio.load('<SPEECH_FILE_PATH>')64if sr != model.sample_rate:65    wav = torchaudio.functional.resample(wav, sr, model.sample_rate)66wav = wav.unsqueeze(0)67 68# Extract discrete codes from SpeechTokenizer69with torch.no_grad():70    codes = model.encode(wav) # codes: (n_q, B, T)71 72semantic_tokens = codes[0, :, :]73acoustic_tokens = codes[1:, :, :]74```75 76### Decoding discrete representions77```python78# Decoding from the first quantizers to ith quantizers79wav = model.decode(codes[:(i + 1)]) # wav: (B, 1, T)80 81# Decoding from ith quantizers to jth quantizers82wav = model.decode(codes[i: (j + 1)], st=i) 83 84# Cancatenating semantic tokens and acoustic tokens and then decoding85semantic_tokens = ... # (..., B, T)86acoustic_tokens = ... # (..., B, T)87wav = model.decode(torch.cat([semantic_tokens, acoustic_tokens], axis=0))88```89 90## Citation91If you use this code or result in your paper, please cite our work as:92```tex93@misc{zhang2023speechtokenizer,94      title={SpeechTokenizer: Unified Speech Tokenizer for Speech Large Language Models}, 95      author={Xin Zhang and Dong Zhang and Shimin Li and Yaqian Zhou and Xipeng Qiu},96      year={2023},97      eprint={2308.16692},98      archivePrefix={arXiv},99      primaryClass={cs.CL}100}101```102 103## License104The code in this repository is released under the Apache 2.0 license as found in the105[LICENSE](LICENSE) file.106