Add inference code

inference.py

@@ -0,0 +1,139 @@
+import torchaudio
+import torch
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
+
+from speechtokenizer import SpeechTokenizer
+from audiotools import AudioSignal
+
+
+def decode_tts(tokens, quantizer, n_codebooks, n_original_tokens, start_audio_token_id, end_audio_token_id):
+    # find start and end indices of audio tokens
+    start = torch.nonzero(tokens == start_audio_token_id)
+    end = torch.nonzero(tokens == end_audio_token_id)
+
+    start = start[0, -1] + 1 if len(start) else 0
+    end = end[0, -1] if len(end) else tokens.shape[-1]
+
+    # subtract length of original vocabulary -> tokens in range [0, 1024)
+    audio_tokens = tokens[start:end] % n_original_tokens
+    reminder = audio_tokens.shape[-1] % n_codebooks
+
+    if reminder:
+        # pad if last frame is incomplete
+        pad_tokens = torch.zeros(n_codebooks - reminder, device="cuda")
+        audio_tokens = torch.cat([audio_tokens, pad_tokens[reminder:n_codebooks]], dim=0)
+
+    transposed = audio_tokens.view(-1, n_codebooks).t()
+    codes = transposed.view(n_codebooks, 1, -1).to(device)
+
+    audio = quantizer.decode(codes).squeeze(0)
+
+    del tokens
+    del audio_tokens
+    torch.cuda.empty_cache()
+
+    return AudioSignal(audio.detach().cpu().numpy(), quantizer.sample_rate)
+
+
+def infer_text_to_audio(text, model, tokenizer, quantizer, max_seq_length=1024, top_k=20):
+    text_tokenized = tokenizer(text, return_tensors="pt")
+    text_input_tokens = text_tokenized["input_ids"].to(device)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+
+    text_tokens = torch.cat([text_input_tokens, soa], dim=1)
+    attention_mask = torch.ones(text_tokens.size(), device=device)
+
+    output_audio_tokens = model.generate(
+        text_tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        top_k=top_k,
+        do_sample=True,
+        temperature=0.8,
+        no_repeat_ngram_size=3,
+    )
+
+    audio_signal = decode_tts(output_audio_tokens[0], quantizer, 3, len(tokenizer) - codebook_size, soa, eoa)
+
+    return audio_signal
+
+
+def infer_audio_to_text(audio_path, model, tokenizer, quantizer, max_seq_length=1024, top_k=20):
+    audio_data, sample_rate = torchaudio.load(audio_path)
+
+    audio = audio_data.view(1, 1, -1).float().to(device)
+    # bandwidth_id = torch.tensor([0])
+
+    codes = quantizer.encode(audio)
+    raw_audio_tokens = codes[:, :n_codebooks_asr] + len(tokenizer) - codebook_size
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    audio_tokens = torch.cat([soa, raw_audio_tokens.view(1, -1), eoa], dim=1)
+    tokens = torch.cat([audio_tokens], dim=1)
+
+    attention_mask = torch.ones(tokens.size(), device=device)
+
+    output_text_tokens = model.generate(
+        tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        temperature=0.6,
+        top_p=0.9,
+        top_k=top_k,
+        no_repeat_ngram_size=4,
+        length_penalty=2.0,
+        repetition_penalty=1.5
+    )
+
+    output_text_tokens = output_text_tokens.cpu()[0]
+    output_text_tokens = output_text_tokens[output_text_tokens < tokenizer(start_audio_token)["input_ids"][-1]]
+    decoded_text = tokenizer.decode(output_text_tokens, skip_special_tokens=True)
+
+    return decoded_text
+
+
+device = "cuda"
+
+n_special_tokens = 3
+n_codebooks_tts = 3
+n_codebooks_asr = 1
+
+start_audio_token = "<soa>"
+end_audio_token = "<eoa>"
+end_sequence_token = "<eos>"
+
+base_model = "Vikhrmodels/salt-116k"
+
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained(base_model, cache_dir=".")
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        cache_dir=".",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+        device_map={"": 0}
+    )
+
+    quantizer_speech = SpeechTokenizer.load_from_checkpoint("speechtokenizer/config.json",
+                                                            "speechtokenizer/SpeechTokenizer.pt")
+    quantizer_speech = quantizer_speech.eval().to(device)
+    codebook_size = quantizer_speech.quantizer.bins
+
+    text = ("Say 'COUNT NUMBERS FROM ONE TO TEN' with a male speaker delivers a very monotone and "
+            "low-pitched speech with a moderate speed in a setting with almost no noise, "
+            "creating a clear and quiet recording.")
+
+    audio_signal = infer_text_to_audio(text, model, tokenizer, quantizer_speech, top_k=50)
+    audio_signal.write("output.wav")
+
+    audio_path = "./input.wav"
+    generated_text = infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech)
+    print(generated_text)