Views
No views yet
📖 See the TinyWave paper (arXiv:2506.23670) and demo site for more details.
1git clone https://github.com/facebookresearch/spiritlm
2cd spiritlm
3pip install -e '.[eval]'1from spiritlm.speech_tokenizer import spiritlm_base
2speech_tokenizer = spiritlm_base()1from transformers import LlamaForCausalLM, AutoTokenizer
2import torchaudio
3import torch
4
5# Load model and tokenizer
6MODEL_PATH = "tinywave/speech-base-2b"
7tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
8model = LlamaForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16)
9
10# Load base speech tokenizer
11speech_tokenizer = spiritlm_base()
12
13def get_inference(audio_path):
14 audio, _ = torchaudio.load(audio_path)
15 input_values = audio.view(1, 1, -1).to(speech_tokenizer.hubert_model.device).float()
16 tokens = speech_tokenizer.encode_string(input_values)
17 input_ids = tokenizer(tokens, return_tensors="pt").input_ids.to(model.device)
18 output = model.generate(input_ids, max_new_tokens=256, top_p=0.9, temperature=0.9, do_sample=True)
19 return tokenizer.decode(output[0])1import numpy as np
2from scipy.io.wavfile import write
3
4def save_array_to_wav_int16(audio_array: np.ndarray, sampling_rate=16000, filename="output.wav"):
5 scaled = np.int16(audio_array / np.max(np.abs(audio_array)) * 32767)
6 write(filename, sampling_rate, scaled)
7
8decoded_audio = speech_tokenizer.decode(generated_output.replace(" ", "").replace("<s>", "").replace("</s>", ""), speaker_id=2)
9save_array_to_wav_int16(decoded_audio, filename="generated.wav")simple_prompt.wav
Output: Semantically consistent speech continuation without expressive variation.| Feature | Description |
|---|---|
| Architecture | 2B parameter distilled transformer |
| Tokenizer | SPIRIT-LM Base (HuBERT phonetic tokens) |
| Input Type | Discrete HuBERT tokens only (speech-only) |
| Output Type | Discrete audio tokens |
| Teacher Model | SPIRIT-LM-Base 7B |
| Tasks | Speech continuation |
| Distillation Method | Layer-aligned (hidden states, attention, logits) |
1@article{nouriborji2025tinywave,
2 title={Efficient Interleaved Speech Modeling through Knowledge Distillation},
3 author={Nouriborji, Mohammadmahdi and Rohanian, Morteza},
4 journal={arXiv preprint arXiv:2506.23670},
5 year={2025}
6}