{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "e8a31bae",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
"unsloth-zoo 2026.5.4 requires hf_transfer, which is not installed.\n",
"unsloth-zoo 2026.5.4 requires msgspec, which is not installed.\n",
"unsloth-zoo 2026.5.4 requires torchao>=0.13.0; sys_platform != \"darwin\" or platform_machine != \"arm64\", which is not installed.\n",
"unsloth-zoo 2026.5.4 requires tyro; sys_platform != \"darwin\" or platform_machine != \"arm64\", which is not installed.\n",
"unsloth-zoo 2026.5.4 requires wheel>=0.42.0, which is not installed.\n",
"unsloth 2026.5.6 requires diffusers, which is not installed.\n",
"unsloth 2026.5.6 requires hf_transfer, which is not installed.\n",
"unsloth 2026.5.6 requires pydantic, which is not installed.\n",
"unsloth 2026.5.6 requires tyro, which is not installed.\n",
"unsloth 2026.5.6 requires wheel>=0.42.0, which is not installed.\n",
"unsloth-zoo 2026.5.4 requires datasets!=4.0.*,!=4.1.0,<4.4.0,>=3.4.1, but you have datasets 4.8.5 which is incompatible.\n",
"unsloth-zoo 2026.5.4 requires transformers!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0,>=4.51.3, but you have transformers 5.9.0 which is incompatible.\n",
"unsloth-zoo 2026.5.4 requires trl!=0.19.0,<=0.24.0,>=0.18.2; sys_platform != \"darwin\" or platform_machine != \"arm64\", but you have trl 1.4.0 which is incompatible.\n",
"unsloth 2026.5.6 requires datasets!=4.0.*,!=4.1.0,<4.4.0,>=3.4.1, but you have datasets 4.8.5 which is incompatible.\n",
"unsloth 2026.5.6 requires transformers!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0,>=4.51.3, but you have transformers 5.9.0 which is incompatible.\n",
"unsloth 2026.5.6 requires trl!=0.19.0,<=0.24.0,>=0.18.2, but you have trl 1.4.0 which is incompatible.\u001b[0m\u001b[31m\n",
"\u001b[0mRequirement already satisfied: accelerate in /usr/local/lib/python3.12/dist-packages (1.13.0)\n",
"Requirement already satisfied: peft in /usr/local/lib/python3.12/dist-packages (0.19.1)\n",
"Requirement already satisfied: trl in /usr/local/lib/python3.12/dist-packages (1.4.0)\n",
"Requirement already satisfied: triton in /usr/local/lib/python3.12/dist-packages (3.4.0)\n",
"Requirement already satisfied: unsloth in /usr/local/lib/python3.12/dist-packages (2026.5.6)\n",
"Requirement already satisfied: unsloth-zoo in /usr/local/lib/python3.12/dist-packages (2026.5.4)\n",
"Requirement already satisfied: cut_cross_entropy in /usr/local/lib/python3.12/dist-packages (25.1.1)\n",
"Requirement already satisfied: xformers==0.0.32.post2 in /usr/local/lib/python3.12/dist-packages (0.0.32.post2)\n"
]
}
],
"source": [
"# !pip install huggingface_hub -q\n",
"!pip install huggingface_hub datasets transformers bitsandbytes soxr soundfile einops einx omegaconf torchcodec==0.7 pyworld sentencepiece protobuf phonemizer -q\n",
"!pip install accelerate peft trl triton unsloth unsloth-zoo triton cut_cross_entropy unsloth_zoo \"xformers==0.0.32.post2\" --no-deps -q"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "a82741bc",
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"\n",
"import torch\n",
"from datasets import load_dataset\n",
"from transformers import AutoTokenizer\n",
"\n",
"from model import AudioTextSpeechBridge\n",
"from speech_instruct_trainer import (\n",
" SpeechInstructTrainConfig,\n",
" train_speech_instruction,\n",
")\n",
"from sparktts.models.audio_tokenizer import BiCodecTokenizer\n",
"\n",
"torch.backends.cuda.matmul.allow_tf32 = True\n",
"torch.backends.cudnn.allow_tf32 = True"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "5f716837",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"asr_dir: models/asr\n",
"llm_dir: models/gemma-3-4b-it\n",
"tts_dir: models/tts\n"
]
}
],
"source": [
"# Model paths. These can be concrete model dirs or Hugging Face cache snapshot parents.\n",
"asr_dir = \"./models/asr\"\n",
"llm_dir = \"./models/gemma-3-4b-it\"\n",
"tts_dir = \"./models/tts\"\n",
"\n",
"def resolve_snapshot_dir(path: str) -> str:\n",
" p = Path(path)\n",
" snapshot_dir = p if p.name == \"snapshots\" else p / \"snapshots\"\n",
" if snapshot_dir.is_dir():\n",
" candidates = [x for x in snapshot_dir.iterdir() if x.is_dir()]\n",
" if candidates:\n",
" return str(sorted(candidates, key=lambda x: x.stat().st_mtime)[-1])\n",
" return str(p)\n",
"\n",
"asr_dir = resolve_snapshot_dir(asr_dir)\n",
"llm_dir = resolve_snapshot_dir(llm_dir)\n",
"tts_dir = resolve_snapshot_dir(tts_dir)\n",
"\n",
"print(\"asr_dir:\", asr_dir)\n",
"print(\"llm_dir:\", llm_dir)\n",
"print(\"tts_dir:\", tts_dir)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "4aff0650",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/usr/local/lib/python3.12/dist-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.\n",
" WeightNorm.apply(module, name, dim)\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Missing tensor: mel_transformer.spectrogram.window\n",
"Missing tensor: mel_transformer.mel_scale.fb\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "b4acad4b65784a48b3d6153cec23d080",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading weights: 0%| | 0/422 [00:00"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"model = AudioTextSpeechBridge(asr_dir=asr_dir, llm_dir=llm_dir, tts_dir=tts_dir)\n",
"state = torch.load(\"./checkpoints_audio_instruct_bn/audio_instruct_bn_best.pt\", map_location=\"cpu\", weights_only=False)\n",
"state_dict = state.get(\"model_state_dict\", state)\n",
"model.load_state_dict(state_dict, strict=True)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "b2de6246-f0f7-4cee-baaf-a7d01ea428b6",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"AudioTextSpeechBridge(\n",
" (asr): WhisperEncoder(\n",
" (conv1): Conv1d(80, 1024, kernel_size=(3,), stride=(1,), padding=(1,))\n",
" (conv2): Conv1d(1024, 1024, kernel_size=(3,), stride=(2,), padding=(1,))\n",
" (embed_positions): Embedding(1500, 1024)\n",
" (layers): ModuleList(\n",
" (0-23): 24 x WhisperEncoderLayer(\n",
" (self_attn): WhisperAttention(\n",
" (k_proj): Linear(in_features=1024, out_features=1024, bias=False)\n",
" (v_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
" (q_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
" (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
" )\n",
" (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
" (activation_fn): GELUActivation()\n",
" (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
" (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
" (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
" )\n",
" )\n",
" (layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
" )\n",
" (llm): Gemma3ForConditionalGeneration(\n",
" (model): Gemma3Model(\n",
" (vision_tower): SiglipVisionModel(\n",
" (embeddings): SiglipVisionEmbeddings(\n",
" (patch_embedding): Conv2d(3, 1152, kernel_size=(14, 14), stride=(14, 14), padding=valid)\n",
" (position_embedding): Embedding(4096, 1152)\n",
" )\n",
" (encoder): SiglipEncoder(\n",
" (layers): ModuleList(\n",
" (0-26): 27 x SiglipEncoderLayer(\n",
" (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
" (self_attn): SiglipAttention(\n",
" (k_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
" (v_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
" (q_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
" (out_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
" )\n",
" (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
" (mlp): SiglipMLP(\n",
" (activation_fn): GELUTanh()\n",
" (fc1): Linear(in_features=1152, out_features=4304, bias=True)\n",
" (fc2): Linear(in_features=4304, out_features=1152, bias=True)\n",
" )\n",
" )\n",
" )\n",
" )\n",
" (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
" )\n",
" (multi_modal_projector): Gemma3MultiModalProjector(\n",
" (mm_soft_emb_norm): Gemma3RMSNorm((1152,), eps=1e-06)\n",
" (avg_pool): AvgPool2d(kernel_size=4, stride=4, padding=0)\n",
" )\n",
" (language_model): Gemma3TextModel(\n",
" (embed_tokens): Gemma3TextScaledWordEmbedding(262208, 2560, padding_idx=0)\n",
" (layers): ModuleList(\n",
" (0-33): 34 x Gemma3DecoderLayer(\n",
" (self_attn): Gemma3Attention(\n",
" (q_proj): Linear(in_features=2560, out_features=2048, bias=False)\n",
" (k_proj): Linear(in_features=2560, out_features=1024, bias=False)\n",
" (v_proj): Linear(in_features=2560, out_features=1024, bias=False)\n",
" (o_proj): Linear(in_features=2048, out_features=2560, bias=False)\n",
" (q_norm): Gemma3RMSNorm((256,), eps=1e-06)\n",
" (k_norm): Gemma3RMSNorm((256,), eps=1e-06)\n",
" )\n",
" (mlp): Gemma3MLP(\n",
" (gate_proj): Linear(in_features=2560, out_features=10240, bias=False)\n",
" (up_proj): Linear(in_features=2560, out_features=10240, bias=False)\n",
" (down_proj): Linear(in_features=10240, out_features=2560, bias=False)\n",
" (act_fn): GELUTanh()\n",
" )\n",
" (input_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
" (post_attention_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
" (pre_feedforward_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
" (post_feedforward_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
" )\n",
" )\n",
" (norm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
" (rotary_emb): Gemma3RotaryEmbedding()\n",
" )\n",
" )\n",
" (lm_head): Linear(in_features=2560, out_features=262208, bias=False)\n",
" )\n",
" (tts): Qwen2ForCausalLM(\n",
" (model): Qwen2Model(\n",
" (embed_tokens): Embedding(166000, 896)\n",
" (layers): ModuleList(\n",
" (0-23): 24 x Qwen2DecoderLayer(\n",
" (self_attn): Qwen2Attention(\n",
" (q_proj): Linear(in_features=896, out_features=896, bias=True)\n",
" (k_proj): Linear(in_features=896, out_features=128, bias=True)\n",
" (v_proj): Linear(in_features=896, out_features=128, bias=True)\n",
" (o_proj): Linear(in_features=896, out_features=896, bias=False)\n",
" )\n",
" (mlp): Qwen2MLP(\n",
" (gate_proj): Linear(in_features=896, out_features=4864, bias=False)\n",
" (up_proj): Linear(in_features=896, out_features=4864, bias=False)\n",
" (down_proj): Linear(in_features=4864, out_features=896, bias=False)\n",
" (act_fn): SiLUActivation()\n",
" )\n",
" (input_layernorm): Qwen2RMSNorm((896,), eps=1e-06)\n",
" (post_attention_layernorm): Qwen2RMSNorm((896,), eps=1e-06)\n",
" )\n",
" )\n",
" (norm): Qwen2RMSNorm((896,), eps=1e-06)\n",
" (rotary_emb): Qwen2RotaryEmbedding()\n",
" )\n",
" (lm_head): Linear(in_features=896, out_features=166000, bias=False)\n",
" )\n",
" (downsampler): CNNDownsampler(\n",
" (conv): Conv1d(1024, 1024, kernel_size=(6,), stride=(5,), padding=(3,), bias=False)\n",
" (norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
" (mlp): Sequential(\n",
" (0): Linear(in_features=1024, out_features=2048, bias=True)\n",
" (1): GELU(approximate='none')\n",
" (2): Linear(in_features=2048, out_features=2048, bias=True)\n",
" (3): GELU(approximate='none')\n",
" (4): Linear(in_features=2048, out_features=1024, bias=True)\n",
" )\n",
" )\n",
" (proj_a): Projector(\n",
" (net): Sequential(\n",
" (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
" (1): Linear(in_features=1024, out_features=1792, bias=True)\n",
" (2): GELU(approximate='none')\n",
" (3): Linear(in_features=1792, out_features=2560, bias=True)\n",
" )\n",
" )\n",
" (proj_b): Projector(\n",
" (net): Sequential(\n",
" (0): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
" (1): Linear(in_features=2560, out_features=1728, bias=True)\n",
" (2): GELU(approximate='none')\n",
" (3): Linear(in_features=1728, out_features=896, bias=True)\n",
" )\n",
" )\n",
")"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"module_state = torch.load(\"./checkpoints_speech_instruct_bn/speech_instruct_bn_best.pt\", map_location=\"cpu\", weights_only=False)\n",
"model.load_modules(module_state, module_names=module_state.get(\"saved_modules\"), strict=True)\n",
"model.to(\"cuda\", torch.bfloat16)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "2377fbea-e792-4ef7-9d2e-13cfaf9ebc49",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[transformers] The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
"[transformers] Setting `pad_token_id` to `eos_token_id`:151645 for open-end generation.\n"
]
}
],
"source": [
"model.set_audio_tokenizer(bicodec)\n",
"\n",
"out1 = model.text_to_speech_waveform(\n",
" text=\"আমি ভাত খাই\",\n",
" llm_tokenizer=llm_tokenizer,\n",
" qwen_tokenizer=qwen_tokenizer,\n",
" voice_name=\"male\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "31d5ce30-ef65-41fb-ae0c-73b8354517b8",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
" \n",
" "
],
"text/plain": [
""
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from IPython.display import Audio\n",
"\n",
"Audio(out1[\"waveform\"], rate=16000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "de36b6a3",
"metadata": {},
"outputs": [],
"source": [
"# Example 1: conversation-style JSON dataset.\n",
"# Every adjacent pair should look like:\n",
"# source turn: {\"text\" or \"content\" or \"ipa\"}\n",
"# target turn: {\"text\" or \"content\" or \"ipa\", \"audio\" or \"target_speech\"}\n",
"#\n",
"# Example 2: direct pairs.\n",
"# Rows can contain:\n",
"# source_text, target_text, target_speech\n",
"#\n",
"# Replace this with your actual dataset path.\n",
"speech_ds = load_dataset(\"json\", data_files=\"./dataset/ban_speech/metadata.jsonl\")[\"train\"]\n",
"print(speech_ds)\n",
"speech_ds[0]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "02c52234",
"metadata": {},
"outputs": [],
"source": [
"cfg = SpeechInstructTrainConfig(\n",
" audio_root=\"./dataset/ban_speech\",\n",
" dataset_fraction=1.0,\n",
" train_split=0.98,\n",
" batch_size=2,\n",
" eval_batch_size=2,\n",
" num_epochs=1,\n",
" lr=2e-4,\n",
" weight_decay=0.01,\n",
" grad_clip=1.0,\n",
" gradient_accumulation_steps=4,\n",
" max_llm_context_tokens=8192,\n",
" max_qwen_context_tokens=2048,\n",
" max_target_tokens=512,\n",
" train_tts_projector=True,\n",
" train_llm=False,\n",
" dtype=\"bfloat16\",\n",
" init_checkpoint=None,\n",
" save_dir=\"./checkpoints_speech_instruct_bn\",\n",
" save_name=\"speech_instruct_bn_best.pt\",\n",
" load_dir=\"\",\n",
" log_every_steps=25,\n",
" eval_every_steps=100,\n",
" save_every_steps=100,\n",
")\n",
"\n",
"cfg"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2fc7be2b",
"metadata": {},
"outputs": [],
"source": [
"model_speech_instruct = train_speech_instruction(\n",
" dataset=speech_ds,\n",
" asr_dir=asr_dir,\n",
" llm_dir=llm_dir,\n",
" tts_dir=tts_dir,\n",
" llm_tokenizer=llm_tokenizer,\n",
" qwen_tokenizer=qwen_tokenizer,\n",
" bicodec_tokenizer=bicodec,\n",
" cfg=cfg,\n",
")\n",
"\n",
"model_speech_instruct"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0ee7b783",
"metadata": {},
"outputs": [],
"source": [
"ckpt_path = Path(cfg.save_dir) / cfg.save_name\n",
"ckpt = torch.load(ckpt_path, map_location=\"cpu\", weights_only=False)\n",
"print(\"Saved modules:\", ckpt.get(\"saved_modules\"))\n",
"print(\"Checkpoint path:\", ckpt_path)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e7586267",
"metadata": {},
"outputs": [],
"source": [
"# Optional: reload only the saved bridge modules into a fresh base model.\n",
"reloaded_model = AudioTextSpeechBridge(asr_dir=asr_dir, llm_dir=llm_dir, tts_dir=tts_dir)\n",
"loaded = reloaded_model.load_modules(ckpt, module_names=ckpt.get(\"saved_modules\"), strict=True)\n",
"print(\"Reloaded modules:\", loaded)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}