初めてのTTSモデルをずんだもんのデータで作ってみた with LLaSA

"""
preprocess_llasa.py
------------------------
・XCodec2 で音声→トークン
・chat template を経由して Llasa 推論時と同じ input_ids を生成
・labels は  以降のみ学習
・可変長 JSONL を Axolotl 用に保存
"""

import argparse, json, os, sys, torch, torchaudio, numpy as np
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer
from xcodec2.modeling_xcodec2 import XCodec2Model
from tqdm import tqdm


p = argparse.ArgumentParser()
p.add_argument("--data", required=True, help="HF dataset name / local path / JSONL …")
p.add_argument("--split", default="train")
p.add_argument("--output", required=True)
p.add_argument("--base_model", default="HKUSTAudio/Llasa-1B-Multilingual")
p.add_argument("--codec", default="HKUSTAudio/xcodec2")
p.add_argument("--sample_rate", type=int, default=16000)
p.add_argument("--max_len", type=int, default=2048)
p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = p.parse_args()


tok = AutoTokenizer.from_pretrained(args.base_model, use_fast=True)
special = [
    "", "",
    "",  ""
]
if not all(t in tok.get_vocab() for t in special):
    tok.add_tokens([t for t in special if t not in tok.get_vocab()])


if "" not in tok.get_vocab():
    tok.add_tokens([f"{i}|>" for i in range(65536)])

codec = XCodec2Model.from_pretrained(args.codec).to(args.device).eval()

BOS, EOS = tok.bos_token_id, tok.eos_token_id
SP_S     = tok.convert_tokens_to_ids("")


ds: Dataset = load_dataset(args.data, split=args.split)


os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
written, skipped = 0, 0

def replace_token(seq, target_id, new_ids):
    """seq(list[int]) の target_id を new_ids に置換（1 回だけ）"""
    try:
        idx = seq.index(target_id)
        return seq[:idx] + new_ids + seq[idx+1:]
    except ValueError:
        return seq

with open(args.output, "w", encoding="utf-8") as fout, torch.inference_mode():
    for ex in tqdm(ds, desc="build jsonl"):
        try:
            
            if not ex.get("text"):
                raise ValueError("no text")
            text_tagged = f"{ex['text']}"
            text_ids = tok.encode(text_tagged, add_special_tokens=False)

            
            if not (ex.get("audio") and "array" in ex["audio"]):
                raise ValueError("no audio")
            wav = torch.tensor(ex["audio"]["array"]).float()
            sr  = ex["audio"].get("sampling_rate", args.sample_rate)
            if sr != args.sample_rate:
                wav = torchaudio.functional.resample(wav, sr, args.sample_rate)
            wav = wav.squeeze()
            if wav.ndim == 0:
                raise ValueError("empty audio")
            wav = wav.unsqueeze(0).to(args.device)  
            codes_np = codec.encode_code(wav)[0][0].cpu().numpy()
            speech_ids = [SP_S] + tok.convert_tokens_to_ids(
                [f"{c}|>" for c in codes_np]
            ) + [tok.convert_tokens_to_ids("")]

            
            chat = [
                {"role": "user",
                 "content": "Convert the text to speech:"},
                {"role": "assistant",
                 "content": ""}
            ]
            
            from inspect import signature
            kw = {}
            if "add_generation_prompt" in signature(tok.apply_chat_template).parameters:
                kw["add_generation_prompt"] = False
            templ_ids = tok.apply_chat_template(chat, tokenize=True, **kw)

            
            seq = replace_token(templ_ids, tok.convert_tokens_to_ids(""), text_ids)
            seq = replace_token(seq, tok.convert_tokens_to_ids(""), speech_ids)

            
            if seq[0] != BOS:
                seq = [BOS] + seq
            if seq[-1] != EOS:
                seq.append(EOS)

            if len(seq) > args.max_len:
                skipped += 1
                continue

            
            labels = seq.copy()
            sp_idx = labels.index(SP_S)
            labels[:sp_idx] = [-100] * sp_idx
            attn = [1] * len(seq)

            
            fout.write(json.dumps({
                "input_ids": seq,
                "attention_mask": attn,
                "labels": labels
            }, ensure_ascii=False) + "\n")
            written += 1

        except Exception as e:
            skipped += 1
            print(e)
            continue

print(f"✅ written : {written}")
print(f"⚠️ skipped : {skipped}")
print(f"➡️ saved   : {args.output}")

base_model: HKUSTAudio/Llasa-1B-Multilingual
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer

hub_model_id: Aratako/Llasa-1B-Zundamon
hub_strategy: "end"
push_dataset_to_hub:
hf_use_auth_token: true


plugins:
  - axolotl.integrations.liger.LigerPlugin

liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_fused_linear_cross_entropy: true
liger_layer_norm: true

load_in_8bit: false
load_in_4bit: false
strict: false

chat_template: tokenizer_default

datasets:
  - path: /workspace/data/preprocessed.jsonl
    ds_type: json
    type:

dataset_processes: 128

shuffle_merged_datasets: true
dataset_prepared_path: /workspace/data/llasa-data
val_set_size: 0.05
output_dir: /workspace/models/Llasa-1B-Zundamon

sequence_len: 2048
sample_packing: true
eval_sample_packing: false
pad_to_sequence_len: true

adapter:
lora_model_dir:
lora_r:
lora_alpha:
lora_dropout:
lora_target_linear:
lora_fan_in_fan_out:

wandb_project: Llasa-1B-Multilingual
wandb_entity: aratako-lm
wandb_watch:
wandb_name: Zundamon-LR1e-5-BS16
wandb_log_model:

gradient_accumulation_steps: 1
micro_batch_size: 2
num_epochs: 30
optimizer: adamw_torch
lr_scheduler: cosine
cosine_min_lr_ratio: 0.1
learning_rate: 1e-5
max_grad_norm: 1.0

train_on_inputs: false
group_by_length: false
bf16: true
fp16: false
tf32: false
bfloat16: true

gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false
early_stopping_patience:
auto_resume_from_checkpoints: true
local_rank:
logging_steps: 1
xformers_attention:
flash_attention: true

save_strategy: epoch
save_total_limit: 2
eval_strategy: epoch
eval_batch_size: 4

save_only_model: true

warmup_steps: 50
debug:
deepspeed: /workspace/axolotl/deepspeed_configs/zero1.json
weight_decay: 0.01
fsdp:
fsdp_config:
metric_for_best_model: eval_loss
load_best_model_at_end: true

import soundfile as sf
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from xcodec2.modeling_xcodec2 import XCodec2Model


model_path = "Aratako/Llasa-1B-Zundamon"





tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16)
model.eval().cuda()

codec_path = "HKUSTAudio/xcodec2"

codec_model = XCodec2Model.from_pretrained(codec_path)
codec_model.eval().cuda()

input_text = "ボク、ずんだもんなのだ！よろしくなのだ！"


def ids_to_speech_tokens(speech_ids):

    speech_tokens_str = []
    for speech_id in speech_ids:
        speech_tokens_str.append(f"{speech_id}|>")
    return speech_tokens_str


def extract_speech_ids(speech_tokens_str):

    speech_ids = []
    for token_str in speech_tokens_str:
        if token_str.startswith(") and token_str.endswith("|>"):
            num_str = token_str[4:-2]

            num = int(num_str)
            speech_ids.append(num)
        else:
            print(f"Unexpected token: {token_str}")
    return speech_ids


with torch.no_grad():

    formatted_text = (
        f"{input_text}"
    )

    chat = [
        {"role": "user", "content": "Convert the text to speech:" + formatted_text},
        {"role": "assistant", "content": ""},
    ]

    input_ids = tokenizer.apply_chat_template(
        chat, tokenize=True, return_tensors="pt", continue_final_message=True
    )
    input_ids = input_ids.to("cuda")
    speech_end_id = tokenizer.convert_tokens_to_ids("")

    
    outputs = model.generate(
        input_ids,
        max_length=2048,
        eos_token_id=speech_end_id,
        do_sample=True,
        top_p=1,
        temperature=0.8,
    )

    generated_ids = outputs[0][input_ids.shape[1] : -1]

    speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

    speech_tokens = extract_speech_ids(speech_tokens)

    speech_tokens = torch.tensor(speech_tokens).cuda().unsqueeze(0).unsqueeze(0)

    gen_wav = codec_model.decode_code(speech_tokens)


sf.write("gen.wav", gen_wav[0, 0, :].cpu().numpy(), 16000)

import soundfile as sf
import torch
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
from xcodec2.modeling_xcodec2 import XCodec2Model






model_path = "Aratako/Llasa-8B-Zundamon"

codec_path = "HKUSTAudio/xcodec2"


def ids_to_speech_tokens(speech_ids):

    speech_tokens_str = []
    for speech_id in speech_ids:
        speech_tokens_str.append(f"{speech_id}|>")
    return speech_tokens_str


def extract_speech_ids(speech_tokens_str):

    speech_ids = []
    for token_str in speech_tokens_str:
        if token_str.startswith(") and token_str.endswith("|>"):
            num_str = token_str[4:-2]

            num = int(num_str)
            speech_ids.append(num)
        else:
            print(f"Unexpected token: {token_str}")
    return speech_ids


def main():
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = LLM(model=model_path, max_model_len=2048, gpu_memory_utilization=0.8)
    codec_model = XCodec2Model.from_pretrained(codec_path)
    codec_model.eval().cuda()
    input_text = "ボク、ずんだもんなのだ！よろしくなのだ！"

    formatted_text = (
        f"{input_text}"
    )

    chat = [
        {"role": "user", "content": "Convert the text to speech:" + formatted_text},
        {"role": "assistant", "content": ""},
    ]

    prompt = tokenizer.apply_chat_template(
        chat, tokenize=False, continue_final_message=True
    )
    print(prompt)
    speech_end_token = ""

    sampling_params = SamplingParams(
        temperature=0.8, top_p=1.0, max_tokens=1024, stop=[speech_end_token]
    )

    
    outputs = model.generate([prompt], sampling_params)

    generated_ids = outputs[0].outputs[0].token_ids

    speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

    speech_tokens = extract_speech_ids(speech_tokens)

    speech_tokens = torch.tensor(speech_tokens).cuda().unsqueeze(0).unsqueeze(0)

    gen_wav = codec_model.decode_code(speech_tokens)

    sf.write("gen.wav", gen_wav[0, 0, :].cpu().numpy(), 16000)


if __name__ == "__main__":
    main()