diff --git a/egs/speech_llm/ASR_LLM/README.md b/egs/speech_llm/ASR_LLM/README.md
new file mode 100644
index 000000000..171240db0
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/README.md
@@ -0,0 +1,20 @@
+
+# Introduction
+
+This recipe includes scripts for training [Qwen-Audio](https://github.com/QwenLM/Qwen-Audio/tree/main) style model using multiple datasets.
+
+
+
+ 
+
+
+
+[./RESULTS.md](./RESULTS.md) contains the latest results.
+
+# ASR_LLM
+
+The following table lists the folders for different tasks.
+
+| | Speech Encoder | LLM | Comment |
+|---------------------------------------|---------------------|--------------------|---------------------------------------------------|
+| [whisper_llm_zh](./whisper_llm_zh) | Whisper | Qwen2 | [Using multiple Chinese datasets](https://github.com/k2-fsa/icefall/tree/master/egs/multi_zh-hans/ASR) |
diff --git a/egs/speech_llm/ASR_LLM/RESULTS.md b/egs/speech_llm/ASR_LLM/RESULTS.md
new file mode 100644
index 000000000..dc2479054
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/RESULTS.md
@@ -0,0 +1,62 @@
+## Results
+
+### whisper_llm_zh finetuning results
+
+| Training Dataset | Speech Encoder | LLM | Projector |Comment | CER |
+| -------------------------| ----------------|------|--------------------------------------------------|-----|--|
+| Aishell1 | whisper-large-v2-aishell1-ft, freeze| Qwen2-1.5B-Instruct, LoRA | Linear, 8x downsample| [yuekai/icefall_asr_aishell_whisper_qwen2_1.5B](https://huggingface.co/yuekai/icefall_asr_aishell_whisper_qwen2_1.5B) | Aishell1 Test 3.62% |
+
+
+Command for training is:
+```bash
+pip install -r whisper_llm_zh/requirements.txt
+
+pip install huggingface_hub['cli']
+mkdir -p models/whisper models/qwen
+
+# For aishell fine-tuned whisper model
+huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_aishell_whisper exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt
+# For multi-hans fine-tuned whisper model
+# huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_multi-hans-zh_whisper v1.1/whisper-large-v2-multi-hans-zh-epoch-3-avg-10.pt
+
+# huggingface-clie download --local-dir models/qwen Qwen/Qwen2-7B-Instruct
+huggingface-clie download --local-dir models/qwen Qwen/Qwen2-1.5B-Instruct
+
+torchrun --nproc_per_node 8 ./whisper_llm_zh/train.py \
+ --max-duration 200 \
+ --exp-dir ./whisper_llm_zh/exp_test \
+ --speech-encoder-path-or-name models/whisper/exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt \
+ --llm-path-or-name Qwen/Qwen2-1.5B-Instruct \
+ --manifest-dir data/fbank \
+ --deepspeed \
+ --deepspeed_config ./whisper_llm_zh/ds_config_zero1.json \
+ --use-flash-attn True \
+ --use-lora True --unfreeze-llm True
+```
+
+Command for decoding using fine-tuned models:
+```bash
+mkdir -p models/whisper models/qwen models/checkpoint
+huggingface-cli download --local-dir models/checkpoint yuekai/icefall_asr_aishell_whisper_qwen2_1.5B
+
+# For aishell fine-tuned whisper model
+huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_aishell_whisper exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt
+# For multi-hans fine-tuned whisper model
+# huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_multi-hans-zh_whisper v1.1/whisper-large-v2-multi-hans-zh-epoch-3-avg-10.pt
+
+huggingface-clie download --local-dir models/qwen Qwen/Qwen2-7B-Instruct
+
+mkdir -p whisper_llm_zh/exp_aishell_whisper_qwen2_1.5B
+ln -s models/checkpoint/epoch-10-avg-5.pt whisper_llm_zh/exp_aishell_whisper_qwen2_1.5B/epoch-999.pt
+
+python3 ./whisper_llm_zh/decode.py \
+ --max-duration 80 \
+ --exp-dir whisper_llm_zh/exp_aishell_whisper_qwen2_1.5B \
+ --speech-encoder-path-or-name models/whisper/exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt \
+ --llm-path-or-name models/qwen \
+ --epoch 999 --avg 1 \
+ --manifest-dir data/fbank \
+ --use-flash-attn True \
+ --use-lora True --dataset aishell
+```
diff --git a/egs/speech_llm/ASR_LLM/assets/framework.png b/egs/speech_llm/ASR_LLM/assets/framework.png
new file mode 100644
index 000000000..dc48bda78
Binary files /dev/null and b/egs/speech_llm/ASR_LLM/assets/framework.png differ
diff --git a/egs/speech_llm/ASR_LLM/prepare.sh b/egs/speech_llm/ASR_LLM/prepare.sh
new file mode 100644
index 000000000..6f5ed5448
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/prepare.sh
@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+
+# fix segmentation fault reported in https://github.com/k2-fsa/icefall/issues/674
+export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
+
+set -eou pipefail
+
+stage=0
+stop_stage=0
+# All files generated by this script are saved in "data".
+# You can safely remove "data" and rerun this script to regenerate it.
+mkdir -p data
+
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+ log "stage 0: Download whisper-large-v2 aishell 1 fbank feature from huggingface"
+
+ # pip install huggingface_hub['cli']
+ # for aishell 1
+ huggingface-cli download --local-dir data yuekai/aishell_whisper_fbank_lhotse
+
+fi
+
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+ log "stage 1: Download whisper-large-v2 multi-hans-zh fbank feature from huggingface"
+
+ # for multi-hans-zh
+ huggingface-cli download --local-dir data/fbank yuekai/wenetspeech_whisper_fbank_lhotse
+ huggingface-cli download --local-dir data/fbank yuekai/multi_hans_zh_whisper_fbank_lhotse
+ huggingface-cli download --local-dir data/fbank yuekai/alimeeting_aishell4_training_whisper_fbank_lhotse
+fi
+
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
+ log "stage 2: Download whisper-large-v2 speechio test sets fbank feature from huggingface"
+
+ # for speechio test sets
+ mkdir data_speechio
+ huggingface-cli download --local-dir data_speechio yuekai/icefall_asr_speechio
+ mv data_speechio/fbank/* data/fbank
+fi
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/asr_datamodule.py b/egs/speech_llm/ASR_LLM/whisper_llm_zh/asr_datamodule.py
new file mode 120000
index 000000000..816a4bf01
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/asr_datamodule.py
@@ -0,0 +1 @@
+../../../multi_zh-hans/ASR/zipformer/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/decode.py b/egs/speech_llm/ASR_LLM/whisper_llm_zh/decode.py
new file mode 100755
index 000000000..882ce4fbf
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/decode.py
@@ -0,0 +1,650 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corporation (Author: Liyong Guo,
+# Fangjun Kuang,
+# Wei Kang)
+# 2024 Yuekai Zhang
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Usage:
+# Command for decoding using fine-tuned models:
+
+pip install huggingface_hub['cli']
+mkdir -p models/whisper models/qwen models/checkpoint
+huggingface-cli download --local-dir models/checkpoint yuekai/icefall_asr_aishell_whisper_qwen2_1.5B
+
+# For aishell fine-tuned whisper model
+huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_aishell_whisper exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt
+# For multi-hans fine-tuned whisper model
+# huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_multi-hans-zh_whisper v1.1/whisper-large-v2-multi-hans-zh-epoch-3-avg-10.pt
+
+huggingface-clie download --local-dir models/qwen Qwen/Qwen2-7B-Instruct
+
+mkdir -p whisper_llm_zh/exp_aishell_whisper_qwen2_1.5B
+ln -s models/checkpoint/epoch-10-avg-5.pt whisper_llm_zh/exp_aishell_whisper_qwen2_1.5B/epoch-999.pt
+
+python3 ./whisper_llm_zh/decode.py \
+ --max-duration 80 \
+ --exp-dir whisper_llm_zh/exp_aishell_whisper_qwen2_1.5B \
+ --speech-encoder-path-or-name models/whisper/exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt \
+ --llm-path-or-name models/qwen \
+ --epoch 999 --avg 1 \
+ --manifest-dir data/fbank \
+ --use-flash-attn True \
+ --use-lora True --dataset aishell
+"""
+
+import argparse
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import torch
+import torch.nn as nn
+import transformers
+import whisper
+from asr_datamodule import AsrDataModule
+from lhotse.cut import Cut
+from model import SPEECH_LLM, EncoderProjector
+from multi_dataset import MultiDataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from train import DEFAULT_SPEECH_TOKEN
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from whisper_encoder_forward_monkey_patch import replace_whisper_encoder_forward
+
+from icefall.checkpoint import average_checkpoints_with_averaged_model, load_checkpoint
+from icefall.env import get_env_info
+from icefall.utils import (
+ AttributeDict,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+
+def average_checkpoints(
+ filenames: List[Path], device: torch.device = torch.device("cpu")
+) -> dict:
+ """Average a list of checkpoints.
+ The function is mainly used for deepspeed converted checkpoint averaging, which only include model state_dict.
+
+ Args:
+ filenames:
+ Filenames of the checkpoints to be averaged. We assume all
+ checkpoints are saved by :func:`save_checkpoint`.
+ device:
+ Move checkpoints to this device before averaging.
+ Returns:
+ Return a dict (i.e., state_dict) which is the average of all
+ model state dicts contained in the checkpoints.
+ """
+ n = len(filenames)
+
+ if "model" in torch.load(filenames[0], map_location=device):
+ avg = torch.load(filenames[0], map_location=device)["model"]
+ else:
+ avg = torch.load(filenames[0], map_location=device)
+
+ # Identify shared parameters. Two parameters are said to be shared
+ # if they have the same data_ptr
+ uniqued: Dict[int, str] = dict()
+
+ for k, v in avg.items():
+ v_data_ptr = v.data_ptr()
+ if v_data_ptr in uniqued:
+ continue
+ uniqued[v_data_ptr] = k
+
+ uniqued_names = list(uniqued.values())
+
+ for i in range(1, n):
+ if "model" in torch.load(filenames[i], map_location=device):
+ state_dict = torch.load(filenames[i], map_location=device)["model"]
+ else:
+ state_dict = torch.load(filenames[i], map_location=device)
+ for k in uniqued_names:
+ avg[k] += state_dict[k]
+
+ for k in uniqued_names:
+ if avg[k].is_floating_point():
+ avg[k] /= n
+ else:
+ avg[k] //= n
+
+ return avg
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--llm-path-or-name",
+ type=str,
+ default="/workspace/asr/Qwen1.5-0.5B-Chat",
+ help="Path or name of the large language model.",
+ )
+
+ parser.add_argument(
+ "--speech-encoder-path-or-name",
+ type=str,
+ default="whisper-large-v2",
+ help="Path or name of the speech encoder.",
+ )
+
+ parser.add_argument(
+ "--encoder-projector-ds-rate",
+ type=int,
+ default=8,
+ help="Downsample rate for the encoder projector.",
+ )
+
+ parser.add_argument(
+ "--use-flash-attn",
+ type=str2bool,
+ default=True,
+ help="Whether to use flash attention.",
+ )
+
+ parser.add_argument(
+ "--use-lora",
+ type=str2bool,
+ default=True,
+ help="Whether to use lora fine-tuned llm checkpoint.",
+ )
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=-1,
+ help="It specifies the checkpoint to use for decoding."
+ "Note: Epoch counts from 0.",
+ )
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=1,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch'. ",
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="beam-search",
+ help="""Decoding method.
+ Supported values are:
+ - beam-search
+ """,
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=1,
+ help="beam size for beam search decoding",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="whisper/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--remove-whisper-encoder-input-length-restriction",
+ type=str2bool,
+ default=True,
+ help="replace whisper encoder forward method to remove input length restriction",
+ )
+
+ parser.add_argument(
+ "--dataset",
+ type=str,
+ default="aishell",
+ choices=["aishell", "speechio", "wenetspeech_test_meeting", "multi_hans_zh"],
+ help="The dataset to decode",
+ )
+
+ add_model_arguments(parser)
+ return parser
+
+
+def get_params() -> AttributeDict:
+ params = AttributeDict(
+ {
+ "env_info": get_env_info(),
+ }
+ )
+ return params
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ tokenizer: AutoTokenizer,
+ batch: dict,
+) -> Dict[str, List[List[int]]]:
+ """Decode one batch and return the result in a dict. The dict has the
+ following format:
+
+ - key: "beam-search"
+ - value: A list of lists. Each sublist is a list of token IDs.
+ Args:
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ batch:
+ It is returned by :meth:`torch.utils.data.DataLoader.__iter__`.
+ Returns:
+ Return a dict, whose key may be "beam-search".
+ """
+
+ def preprocess(
+ messages,
+ tokenizer: transformers.PreTrainedTokenizer,
+ max_len: int = 128,
+ ) -> Dict:
+ """Preprocesses the data for supervised fine-tuning."""
+ texts = []
+ TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if loop.last %}{{''}}{% else %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}"
+ for i, msg in enumerate(messages):
+ texts.append(
+ tokenizer.apply_chat_template(
+ msg,
+ tokenize=True,
+ add_generation_prompt=False,
+ chat_template=TEMPLATE,
+ padding="longest",
+ max_length=max_len,
+ truncation=True,
+ )
+ )
+ max_len_texts = max([len(text) for text in texts])
+ if tokenizer.padding_side == "right":
+ texts = [
+ text + [tokenizer.pad_token_id] * (max_len_texts - len(text))
+ for text in texts
+ ]
+ else:
+ texts = [
+ [tokenizer.pad_token_id] * (max_len_texts - len(text)) + text
+ for text in texts
+ ]
+
+ input_ids = torch.tensor(texts, dtype=torch.int)
+
+ attention_mask = input_ids.ne(tokenizer.pad_token_id)
+
+ return input_ids, attention_mask
+
+ dtype = torch.float32
+ device = model.llm.device
+
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+ feature = feature.to(device, dtype=dtype).transpose(1, 2)
+ if not params.remove_whisper_encoder_input_length_restriction:
+ T = 3000
+ if feature.shape[2] < T:
+ feature = torch.cat(
+ [
+ feature,
+ torch.zeros(
+ feature.shape[0], feature.shape[1], T - feature.shape[2]
+ ).to(device, dtype=dtype),
+ ],
+ 2,
+ )
+
+ supervisions = batch["supervisions"]
+ feature_len = supervisions["num_frames"]
+ feature_len = feature_len.to(device, dtype=dtype)
+
+ messages = [
+ [
+ {"role": "user", "content": f"{DEFAULT_SPEECH_TOKEN}请转写音频为文字"},
+ {"role": "assistant", "content": ""},
+ ]
+ ] * len(feature)
+
+ input_ids, attention_mask = preprocess(messages, tokenizer, max_len=128)
+
+ generated_ids = model.decode(
+ feature, input_ids.to(device, dtype=torch.long), attention_mask.to(device)
+ )
+ hyps = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+ return {"beam-search": hyps}
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ tokenizer: AutoTokenizer,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ The dataloader.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ Returns:
+ Return a dict, whose key may be "beam-search".
+ """
+
+ def normalize_text_alimeeting(text: str, normalize: str = "m2met") -> str:
+ """
+ Text normalization similar to M2MeT challenge baseline.
+ See: https://github.com/yufan-aslp/AliMeeting/blob/main/asr/local/text_normalize.pl
+ """
+ if normalize == "none":
+ return text
+ elif normalize == "m2met":
+ import re
+
+ text = text.replace(" ", "")
+ text = text.replace("", "")
+ text = text.replace("<%>", "")
+ text = text.replace("<->", "")
+ text = text.replace("<$>", "")
+ text = text.replace("<#>", "")
+ text = text.replace("<_>", "")
+ text = text.replace("", "")
+ text = text.replace("`", "")
+ text = text.replace("&", "")
+ text = text.replace(",", "")
+ if re.search("[a-zA-Z]", text):
+ text = text.upper()
+ text = text.replace("Ａ", "A")
+ text = text.replace("ａ", "A")
+ text = text.replace("ｂ", "B")
+ text = text.replace("ｃ", "C")
+ text = text.replace("ｋ", "K")
+ text = text.replace("ｔ", "T")
+ text = text.replace("，", "")
+ text = text.replace("丶", "")
+ text = text.replace("。", "")
+ text = text.replace("、", "")
+ text = text.replace("？", "")
+ return text
+
+ results = []
+
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ batch=batch,
+ tokenizer=tokenizer,
+ )
+
+ for lm_scale, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_text = normalize_text_alimeeting(ref_text)
+ ref_words = ref_text.split()
+ print(f"ref: {ref_text}")
+ print(f"hyp: {''.join(hyp_words)}")
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[lm_scale].extend(this_batch)
+
+ num_cuts += len(batch["supervisions"]["text"])
+
+ if batch_idx % 100 == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+
+ enable_log = True
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.exp_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ store_transcripts(filename=recog_path, texts=results)
+ if enable_log:
+ logging.info(f"The transcripts are stored in {recog_path}")
+
+ # The following prints out WERs, per-word error statistics and aligned
+ # ref/hyp pairs.
+ errs_filename = (
+ params.exp_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ # we compute CER for aishell dataset.
+ results_char = []
+ for res in results:
+ results_char.append((res[0], list("".join(res[1])), list("".join(res[2]))))
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results_char, enable_log=enable_log
+ )
+ test_set_wers[key] = wer
+
+ if enable_log:
+ logging.info("Wrote detailed error stats to {}".format(errs_filename))
+
+ test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1])
+ errs_info = params.exp_dir / f"cer-summary-{test_set_name}-{params.suffix}.txt"
+ with open(errs_info, "w") as f:
+ print("settings\tCER", file=f)
+ for key, val in test_set_wers:
+ print("{}\t{}".format(key, val), file=f)
+
+ s = "\nFor {}, CER of different settings are:\n".format(test_set_name)
+ note = "\tbest for {}".format(test_set_name)
+ for key, val in test_set_wers:
+ s += "{}\t{}{}\n".format(key, val, note)
+ note = ""
+ logging.info(s)
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ AsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+ setup_logger(
+ f"{params.exp_dir}/log-{params.method}-beam{params.beam_size}/log-decode-{params.suffix}"
+ )
+
+ logging.info("Decoding started")
+ logging.info(params)
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda")
+
+ logging.info(f"device: {device}")
+
+ if params.remove_whisper_encoder_input_length_restriction:
+ replace_whisper_encoder_forward()
+
+ whisper_model = whisper.load_model(params.speech_encoder_path_or_name, "cpu")
+ speech_encoder = whisper_model.encoder
+ speech_encoder_dim = whisper_model.dims.n_audio_state
+ tokenizer = AutoTokenizer.from_pretrained(params.llm_path_or_name)
+
+ if params.use_flash_attn:
+ attn_implementation = "flash_attention_2"
+ # torch_dtype=torch.bfloat16 FIX ME
+ torch_dtype = torch.float16
+ tokenizer.padding_side = "left"
+
+ else:
+ attn_implementation = "eager"
+ torch_dtype = torch.float16
+ tokenizer.padding_side = "right"
+
+ llm = AutoModelForCausalLM.from_pretrained(
+ params.llm_path_or_name,
+ attn_implementation=attn_implementation,
+ torch_dtype=torch_dtype,
+ )
+ if params.use_lora:
+ lora_config = LoraConfig(
+ r=64,
+ lora_alpha=16,
+ target_modules=[
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ "o_proj",
+ "up_proj",
+ "gate_proj",
+ "down_proj",
+ ],
+ task_type="CAUSAL_LM",
+ )
+ llm = get_peft_model(llm, lora_config)
+ llm.print_trainable_parameters()
+
+ special_tokens_dict = {"additional_special_tokens": [DEFAULT_SPEECH_TOKEN]}
+ tokenizer.add_special_tokens(special_tokens_dict)
+ llm.config.pad_token_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
+ llm.config.bos_token_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
+ llm.config.eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+
+ llm.config.default_speech_token_id = tokenizer.convert_tokens_to_ids(
+ DEFAULT_SPEECH_TOKEN
+ )
+
+ encoder_projector = EncoderProjector(
+ speech_encoder_dim, llm.config.hidden_size, params.encoder_projector_ds_rate
+ )
+
+ model = SPEECH_LLM(
+ speech_encoder,
+ llm,
+ encoder_projector,
+ )
+
+ if params.avg > 1:
+ start = params.epoch - params.avg + 1
+ assert start >= 1, start
+ checkpoint = torch.load(
+ f"{params.exp_dir}/epoch-{params.epoch}.pt", map_location="cpu"
+ )
+ assert "model" not in checkpoint
+ # deepspeed converted checkpoint only contains model state_dict
+ filenames = [
+ f"{params.exp_dir}/epoch-{epoch}.pt"
+ for epoch in range(start, params.epoch + 1)
+ ]
+ avg_checkpoint = average_checkpoints(filenames)
+ model.load_state_dict(avg_checkpoint, strict=False)
+
+ filename = f"{params.exp_dir}/epoch-{params.epoch}-avg-{params.avg}.pt"
+ torch.save(avg_checkpoint, filename)
+ else:
+ checkpoint = torch.load(
+ f"{params.exp_dir}/epoch-{params.epoch}.pt", map_location="cpu"
+ )
+ model.load_state_dict(checkpoint, strict=False)
+
+ model.to(device)
+ model.eval()
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+
+ data_module = AsrDataModule(args)
+ multi_dataset = MultiDataset(args.manifest_dir)
+
+ def remove_long_utt(c: Cut):
+ # Keep only utterances with duration in 30 seconds
+ #
+ if c.duration > 30.0:
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ )
+ return False
+ return True
+
+ if params.dataset == "aishell":
+ test_sets_cuts = multi_dataset.aishell_test_cuts()
+ elif params.dataset == "speechio":
+ test_sets_cuts = multi_dataset.speechio_test_cuts()
+ elif params.dataset == "wenetspeech_test_meeting":
+ test_sets_cuts = multi_dataset.wenetspeech_test_meeting_cuts()
+ else:
+ test_sets_cuts = multi_dataset.test_cuts()
+
+ test_sets = test_sets_cuts.keys()
+ test_dls = [
+ data_module.test_dataloaders(test_sets_cuts[cuts_name].filter(remove_long_utt))
+ for cuts_name in test_sets
+ ]
+
+ for test_set, test_dl in zip(test_sets, test_dls):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ tokenizer=tokenizer,
+ )
+
+ save_results(params=params, test_set_name=test_set, results_dict=results_dict)
+
+ logging.info("Done!")
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/ds_config_zero1.json b/egs/speech_llm/ASR_LLM/whisper_llm_zh/ds_config_zero1.json
new file mode 100644
index 000000000..730937a21
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/ds_config_zero1.json
@@ -0,0 +1,38 @@
+{
+ "fp16": {
+ "enabled": true,
+ "loss_scale": 0,
+ "loss_scale_window": 100,
+ "initial_scale_power": 16,
+ "hysteresis": 2,
+ "min_loss_scale": 0.01
+ },
+ "zero_optimization": {
+ "stage": 1,
+ "allgather_partitions": true,
+ "allgather_bucket_size": 2e8,
+ "overlap_comm": true,
+ "reduce_scatter": true,
+ "reduce_bucket_size": 2e8,
+ "contiguous_gradients": true
+ },
+ "optimizer": {
+ "type": "Adam",
+ "params": {
+ "lr": 1e-4
+ }
+ },
+ "scheduler": {
+ "type": "WarmupLR",
+ "params": {
+ "warmup_min_lr": 0,
+ "warmup_max_lr": 1e-4,
+ "warmup_num_steps": 100
+ }
+ },
+ "gradient_accumulation_steps": 1,
+ "gradient_clipping": 5,
+ "steps_per_print": 50,
+ "train_micro_batch_size_per_gpu": 1,
+ "wall_clock_breakdown": false
+}
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/model.py b/egs/speech_llm/ASR_LLM/whisper_llm_zh/model.py
new file mode 100644
index 000000000..829ef4e2d
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/model.py
@@ -0,0 +1,285 @@
+import torch
+from torch import nn
+from transformers.trainer_pt_utils import LabelSmoother
+
+IGNORE_TOKEN_ID = LabelSmoother.ignore_index
+
+
+class EncoderProjector(nn.Module):
+ """
+ The encoder projector module. It is used to project the encoder outputs to the same dimension as the language model.
+ Modified from https://github.com/X-LANCE/SLAM-LLM/blob/main/src/slam_llm/models/projector.py.
+ Args:
+ encoder_dim (:obj:`int`): The dimension of the encoder outputs.
+ llm_dim (:obj:`int`): The dimension of the language model.
+ downsample_rate (:obj:`int`, `optional`, defaults to 5): The downsample rate to use.
+ """
+
+ def __init__(self, encoder_dim, llm_dim, downsample_rate=5):
+ super().__init__()
+ self.downsample_rate = downsample_rate
+ self.linear1 = nn.Linear(encoder_dim * self.downsample_rate, llm_dim)
+ self.relu = nn.ReLU()
+ self.linear2 = nn.Linear(llm_dim, llm_dim)
+
+ def forward(self, x):
+
+ batch_size, seq_len, feat_dim = x.size()
+ num_frames_to_discard = seq_len % self.downsample_rate
+ if num_frames_to_discard > 0:
+ x = x[:, :-num_frames_to_discard, :]
+ seq_len = x.size(1)
+
+ x = x.contiguous()
+ x = x.view(
+ batch_size, seq_len // self.downsample_rate, feat_dim * self.downsample_rate
+ )
+
+ x = self.linear1(x)
+ x = self.relu(x)
+ x = self.linear2(x)
+ return x
+
+
+class SPEECH_LLM(nn.Module):
+ """
+ The Speech-to-Text model. It consists of an encoder, a language model and an encoder projector.
+ The encoder is used to extract speech features from the input speech signal.
+ The encoder projector is used to project the encoder outputs to the same dimension as the language model.
+ The language model is used to generate the text from the speech features.
+ Args:
+ encoder (:obj:`nn.Module`): The encoder module.
+ llm (:obj:`nn.Module`): The language model module.
+ encoder_projector (:obj:`nn.Module`): The encoder projector module.
+ """
+
+ def __init__(
+ self,
+ encoder: nn.Module,
+ llm: nn.Module,
+ encoder_projector: nn.Module,
+ ):
+ super().__init__()
+ self.encoder = encoder
+ self.llm = llm
+ self.encoder_projector = encoder_projector
+
+ def _merge_input_ids_with_speech_features(
+ self, speech_features, inputs_embeds, input_ids, attention_mask, labels=None
+ ):
+ """
+ Merge the speech features with the input_ids and attention_mask. This is done by replacing the speech tokens
+ with the speech features and padding the input_ids to the maximum length of the speech features.
+ Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/modeling_llava.py#L277.
+ Args:
+ speech_features (:obj:`torch.Tensor`): The speech features to merge with the input_ids.
+ inputs_embeds (:obj:`torch.Tensor`): The embeddings of the input_ids.
+ input_ids (:obj:`torch.Tensor`): The input ids to merge.
+ attention_mask (:obj:`torch.Tensor`): The attention mask to merge.
+ labels (:obj:`torch.Tensor`, `optional`): The labels to merge.
+ Returns:
+ :obj:`Tuple(torch.Tensor)`: The merged embeddings, attention mask, labels and position ids.
+ """
+ num_speechs, speech_len, embed_dim = speech_features.shape
+ batch_size, sequence_length = input_ids.shape
+ left_padding = not torch.sum(
+ input_ids[:, -1] == torch.tensor(self.llm.config.pad_token_id)
+ )
+ # 1. Create a mask to know where special speech tokens are
+ special_speech_token_mask = input_ids == self.llm.config.default_speech_token_id
+ num_special_speech_tokens = torch.sum(special_speech_token_mask, dim=-1)
+ # Compute the maximum embed dimension
+ max_embed_dim = (
+ num_special_speech_tokens.max() * (speech_len - 1)
+ ) + sequence_length
+ batch_indices, non_speech_indices = torch.where(
+ input_ids != self.llm.config.default_speech_token_id
+ )
+
+ # 2. Compute the positions where text should be written
+ # Calculate new positions for text tokens in merged speech-text sequence.
+ # `special_speech_token_mask` identifies speech tokens. Each speech token will be replaced by `nb_text_tokens_per_speechs - 1` text tokens.
+ # `torch.cumsum` computes how each speech token shifts subsequent text token positions.
+ # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one.
+ new_token_positions = (
+ torch.cumsum((special_speech_token_mask * (speech_len - 1) + 1), -1) - 1
+ )
+ nb_speech_pad = max_embed_dim - 1 - new_token_positions[:, -1]
+ if left_padding:
+ new_token_positions += nb_speech_pad[:, None] # offset for left padding
+ text_to_overwrite = new_token_positions[batch_indices, non_speech_indices]
+
+ # 3. Create the full embedding, already padded to the maximum position
+ final_embedding = torch.zeros(
+ batch_size,
+ max_embed_dim,
+ embed_dim,
+ dtype=inputs_embeds.dtype,
+ device=inputs_embeds.device,
+ )
+ final_attention_mask = torch.zeros(
+ batch_size,
+ max_embed_dim,
+ dtype=attention_mask.dtype,
+ device=inputs_embeds.device,
+ )
+ if labels is not None:
+ final_labels = torch.full(
+ (batch_size, max_embed_dim),
+ IGNORE_TOKEN_ID,
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ # In case the Vision model or the Language model has been offloaded to CPU, we need to manually
+ # set the corresponding tensors into their correct target device.
+ target_device = inputs_embeds.device
+ batch_indices, non_speech_indices, text_to_overwrite = (
+ batch_indices.to(target_device),
+ non_speech_indices.to(target_device),
+ text_to_overwrite.to(target_device),
+ )
+ attention_mask = attention_mask.to(target_device)
+
+ # 4. Fill the embeddings based on the mask. If we have ["hey" "", "how", "are"]
+ # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the speech features
+ final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[
+ batch_indices, non_speech_indices
+ ]
+ final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[
+ batch_indices, non_speech_indices
+ ]
+ if labels is not None:
+ final_labels[batch_indices, text_to_overwrite] = labels[
+ batch_indices, non_speech_indices
+ ]
+
+ # 5. Fill the embeddings corresponding to the speechs. Anything that is not `text_positions` needs filling (#29835)
+ speech_to_overwrite = torch.full(
+ (batch_size, max_embed_dim),
+ True,
+ dtype=torch.bool,
+ device=inputs_embeds.device,
+ )
+ speech_to_overwrite[batch_indices, text_to_overwrite] = False
+ speech_to_overwrite &= speech_to_overwrite.cumsum(-1) - 1 >= nb_speech_pad[
+ :, None
+ ].to(target_device)
+
+ if speech_to_overwrite.sum() != speech_features.shape[:-1].numel():
+ raise ValueError(
+ f"The input provided to the model are wrong. The number of speech tokens is {torch.sum(special_speech_token_mask)} while"
+ f" the number of speech given to the model is {num_speechs}. This prevents correct indexing and breaks batch generation."
+ )
+
+ final_embedding[speech_to_overwrite] = (
+ speech_features.contiguous().reshape(-1, embed_dim).to(target_device)
+ )
+ final_attention_mask |= speech_to_overwrite
+ position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_(
+ (final_attention_mask == 0), 1
+ )
+
+ # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens.
+ batch_indices, pad_indices = torch.where(
+ input_ids == self.llm.config.pad_token_id
+ )
+ indices_to_mask = new_token_positions[batch_indices, pad_indices]
+
+ final_embedding[batch_indices, indices_to_mask] = 0
+
+ if labels is None:
+ final_labels = None
+
+ return final_embedding, final_attention_mask, final_labels, position_ids
+
+ def forward(
+ self,
+ fbank: torch.Tensor = None,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor = None,
+ labels: torch.LongTensor = None,
+ ):
+ encoder_outs = self.encoder(fbank)
+
+ speech_features = self.encoder_projector(encoder_outs)
+
+ inputs_embeds = self.llm.get_input_embeddings()(input_ids)
+
+ (
+ inputs_embeds,
+ attention_mask,
+ labels,
+ _,
+ ) = self._merge_input_ids_with_speech_features(
+ speech_features, inputs_embeds, input_ids, attention_mask, labels
+ )
+
+ model_outputs = self.llm(
+ inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels
+ )
+
+ with torch.no_grad():
+ preds = torch.argmax(model_outputs.logits, -1)
+ acc = compute_accuracy(
+ preds.detach()[:, :-1],
+ labels.detach()[:, 1:],
+ ignore_label=IGNORE_TOKEN_ID,
+ )
+ return model_outputs, acc
+
+ def decode(
+ self,
+ fbank: torch.Tensor = None,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor = None,
+ **kwargs,
+ ):
+
+ encoder_outs = self.encoder(fbank)
+ speech_features = self.encoder_projector(encoder_outs)
+ speech_features = speech_features.to(torch.float16)
+ inputs_embeds = self.llm.get_input_embeddings()(input_ids)
+ (
+ inputs_embeds,
+ attention_mask,
+ _,
+ position_ids,
+ ) = self._merge_input_ids_with_speech_features(
+ speech_features, inputs_embeds, input_ids, attention_mask
+ )
+ generated_ids = self.llm.generate(
+ inputs_embeds=inputs_embeds,
+ max_new_tokens=kwargs.get("max_new_tokens", 200),
+ num_beams=kwargs.get("num_beams", 1),
+ do_sample=kwargs.get("do_sample", False),
+ min_length=kwargs.get("min_length", 1),
+ top_p=kwargs.get("top_p", 1.0),
+ repetition_penalty=kwargs.get("repetition_penalty", 1.0),
+ length_penalty=kwargs.get("length_penalty", 1.0),
+ temperature=kwargs.get("temperature", 1.0),
+ bos_token_id=self.llm.config.bos_token_id,
+ eos_token_id=self.llm.config.eos_token_id,
+ pad_token_id=self.llm.config.pad_token_id,
+ )
+
+ return generated_ids
+
+
+def compute_accuracy(pad_outputs, pad_targets, ignore_label):
+ """Calculate accuracy.
+ Copied from https://github.com/X-LANCE/SLAM-LLM/blob/main/src/slam_llm/utils/metric.py
+ Args:
+ pad_outputs (LongTensor): Prediction tensors (B, Lmax).
+ pad_targets (LongTensor): Target label tensors (B, Lmax).
+ ignore_label (int): Ignore label id.
+
+ Returns:
+ float: Accuracy value (0.0 - 1.0).
+
+ """
+ mask = pad_targets != ignore_label
+ numerator = torch.sum(
+ pad_outputs.masked_select(mask) == pad_targets.masked_select(mask)
+ )
+ denominator = torch.sum(mask)
+ return numerator.float() / denominator.float()
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/multi_dataset.py b/egs/speech_llm/ASR_LLM/whisper_llm_zh/multi_dataset.py
new file mode 100644
index 000000000..eae967500
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/multi_dataset.py
@@ -0,0 +1,338 @@
+# Copyright 2023 Xiaomi Corp. (authors: Zengrui Jin)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import glob
+import logging
+import re
+from pathlib import Path
+from typing import Dict, List
+
+import lhotse
+from lhotse import CutSet, load_manifest_lazy
+
+
+class MultiDataset:
+ def __init__(self, fbank_dir: str):
+ """
+ Args:
+ manifest_dir:
+ It is expected to contain the following files:
+ - aishell_cuts_train.jsonl.gz
+ - aishell2_cuts_train.jsonl.gz
+ - aishell4_cuts_train_L.jsonl.gz
+ - aishell4_cuts_train_M.jsonl.gz
+ - aishell4_cuts_train_S.jsonl.gz
+ - alimeeting-far_cuts_train.jsonl.gz
+ - magicdata_cuts_train.jsonl.gz
+ - primewords_cuts_train.jsonl.gz
+ - stcmds_cuts_train.jsonl.gz
+ - thchs_30_cuts_train.jsonl.gz
+ - kespeech/kespeech-asr_cuts_train_phase1.jsonl.gz
+ - kespeech/kespeech-asr_cuts_train_phase2.jsonl.gz
+ - wenetspeech/cuts_L_fixed.jsonl.gz
+ """
+ self.fbank_dir = Path(fbank_dir)
+
+ def train_cuts(self) -> CutSet:
+ logging.info("About to get multidataset train cuts")
+
+ # THCHS-30
+ logging.info("Loading THCHS-30 in lazy mode")
+ thchs_30_cuts = load_manifest_lazy(
+ self.fbank_dir / "thchs_30_cuts_train.jsonl.gz"
+ )
+
+ # AISHELL-1
+ logging.info("Loading Aishell-1 in lazy mode")
+ aishell_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell_cuts_train.jsonl.gz"
+ )
+
+ # AISHELL-2
+ logging.info("Loading Aishell-2 in lazy mode")
+ aishell_2_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell2_cuts_train.jsonl.gz"
+ )
+
+ # AISHELL-4
+ logging.info("Loading Aishell-4 in lazy mode")
+ aishell_4_L_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell4_cuts_train_L.jsonl.gz"
+ )
+ aishell_4_M_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell4_cuts_train_M.jsonl.gz"
+ )
+ aishell_4_S_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell4_cuts_train_S.jsonl.gz"
+ )
+
+ # ST-CMDS
+ logging.info("Loading ST-CMDS in lazy mode")
+ stcmds_cuts = load_manifest_lazy(self.fbank_dir / "stcmds_cuts_train.jsonl.gz")
+
+ # Primewords
+ logging.info("Loading Primewords in lazy mode")
+ primewords_cuts = load_manifest_lazy(
+ self.fbank_dir / "primewords_cuts_train.jsonl.gz"
+ )
+
+ # MagicData
+ logging.info("Loading MagicData in lazy mode")
+ magicdata_cuts = load_manifest_lazy(
+ self.fbank_dir / "magicdata_cuts_train.jsonl.gz"
+ )
+
+ # Ali-Meeting
+ logging.info("Loading Ali-Meeting in lazy mode")
+ alimeeting_cuts = load_manifest_lazy(
+ self.fbank_dir / "alimeeting-far_cuts_train.jsonl.gz"
+ )
+
+ # WeNetSpeech
+ logging.info("Loading WeNetSpeech in lazy mode")
+ wenetspeech_L_cuts = load_manifest_lazy(
+ self.fbank_dir / "wenetspeech" / "cuts_L_fixed.jsonl.gz"
+ )
+
+ # KeSpeech
+ logging.info("Loading KeSpeech in lazy mode")
+ kespeech_1_cuts = load_manifest_lazy(
+ self.fbank_dir / "kespeech" / "kespeech-asr_cuts_train_phase1.jsonl.gz"
+ )
+ kespeech_2_cuts = load_manifest_lazy(
+ self.fbank_dir / "kespeech" / "kespeech-asr_cuts_train_phase2.jsonl.gz"
+ )
+
+ return CutSet.mux(
+ thchs_30_cuts,
+ aishell_cuts,
+ aishell_2_cuts,
+ aishell_4_L_cuts,
+ aishell_4_M_cuts,
+ aishell_4_S_cuts,
+ alimeeting_cuts,
+ stcmds_cuts,
+ primewords_cuts,
+ magicdata_cuts,
+ wenetspeech_L_cuts,
+ kespeech_1_cuts,
+ kespeech_2_cuts,
+ weights=[
+ len(thchs_30_cuts),
+ len(aishell_cuts),
+ len(aishell_2_cuts),
+ len(aishell_4_L_cuts),
+ len(aishell_4_M_cuts),
+ len(aishell_4_S_cuts),
+ len(alimeeting_cuts),
+ len(stcmds_cuts),
+ len(primewords_cuts),
+ len(magicdata_cuts),
+ len(wenetspeech_L_cuts),
+ len(kespeech_1_cuts),
+ len(kespeech_2_cuts),
+ ],
+ )
+
+ def dev_cuts(self) -> CutSet:
+ logging.info("About to get multidataset dev cuts")
+
+ # WeNetSpeech
+ logging.info("Loading WeNetSpeech DEV set in lazy mode")
+ wenetspeech_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "wenetspeech" / "cuts_DEV_fixed.jsonl.gz"
+ )
+
+ return wenetspeech_dev_cuts
+
+ def test_cuts(self) -> Dict[str, CutSet]:
+ logging.info("About to get multidataset test cuts")
+
+ # AISHELL
+ logging.info("Loading Aishell set in lazy mode")
+ aishell_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell_cuts_test.jsonl.gz"
+ )
+ aishell_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell_cuts_dev.jsonl.gz"
+ )
+
+ # AISHELL-2
+ logging.info("Loading Aishell-2 set in lazy mode")
+ aishell2_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell2_cuts_test.jsonl.gz"
+ )
+ aishell2_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell2_cuts_dev.jsonl.gz"
+ )
+
+ # AISHELL-4
+ logging.info("Loading Aishell-4 TEST set in lazy mode")
+ aishell4_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell4_cuts_test.jsonl.gz"
+ )
+
+ # Ali-Meeting
+ logging.info("Loading Ali-Meeting set in lazy mode")
+ alimeeting_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "alimeeting-far_cuts_test.jsonl.gz"
+ )
+ alimeeting_eval_cuts = load_manifest_lazy(
+ self.fbank_dir / "alimeeting-far_cuts_eval.jsonl.gz"
+ )
+
+ # MagicData
+ logging.info("Loading MagicData set in lazy mode")
+ magicdata_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "magicdata_cuts_test.jsonl.gz"
+ )
+ magicdata_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "magicdata_cuts_dev.jsonl.gz"
+ )
+
+ # KeSpeech
+ logging.info("Loading KeSpeech set in lazy mode")
+ kespeech_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "kespeech" / "kespeech-asr_cuts_test.jsonl.gz"
+ )
+ kespeech_dev_phase1_cuts = load_manifest_lazy(
+ self.fbank_dir / "kespeech" / "kespeech-asr_cuts_dev_phase1.jsonl.gz"
+ )
+ kespeech_dev_phase2_cuts = load_manifest_lazy(
+ self.fbank_dir / "kespeech" / "kespeech-asr_cuts_dev_phase2.jsonl.gz"
+ )
+
+ # WeNetSpeech
+ logging.info("Loading WeNetSpeech set in lazy mode")
+ wenetspeech_test_meeting_cuts = load_manifest_lazy(
+ self.fbank_dir / "wenetspeech" / "cuts_TEST_MEETING.jsonl.gz"
+ )
+ wenetspeech_test_net_cuts = load_manifest_lazy(
+ self.fbank_dir / "wenetspeech" / "cuts_TEST_NET.jsonl.gz"
+ )
+ wenetspeech_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "wenetspeech" / "cuts_DEV_fixed.jsonl.gz"
+ )
+
+ return {
+ "wenetspeech-meeting_test": wenetspeech_test_meeting_cuts,
+ "aishell_test": aishell_test_cuts,
+ "aishell_dev": aishell_dev_cuts,
+ "ali-meeting_test": alimeeting_test_cuts,
+ "ali-meeting_eval": alimeeting_eval_cuts,
+ "aishell-4_test": aishell4_test_cuts,
+ "aishell-2_test": aishell2_test_cuts,
+ "aishell-2_dev": aishell2_dev_cuts,
+ "magicdata_test": magicdata_test_cuts,
+ "magicdata_dev": magicdata_dev_cuts,
+ "kespeech-asr_test": kespeech_test_cuts,
+ "kespeech-asr_dev_phase1": kespeech_dev_phase1_cuts,
+ "kespeech-asr_dev_phase2": kespeech_dev_phase2_cuts,
+ "wenetspeech-net_test": wenetspeech_test_net_cuts,
+ "wenetspeech_dev": wenetspeech_dev_cuts,
+ }
+
+ def aishell_train_cuts(self) -> CutSet:
+ logging.info("About to get multidataset train cuts")
+ logging.info("Loading Aishell-1 in lazy mode")
+ aishell_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell_cuts_train.jsonl.gz"
+ )
+
+ return aishell_cuts
+
+ def aishell_dev_cuts(self) -> CutSet:
+ logging.info("About to get multidataset dev cuts")
+ logging.info("Loading Aishell set in lazy mode")
+ aishell_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell_cuts_dev.jsonl.gz"
+ )
+
+ return aishell_dev_cuts
+
+ def aishell_test_cuts(self) -> CutSet:
+ logging.info("About to get multidataset test cuts")
+ logging.info("Loading Aishell set in lazy mode")
+ aishell_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell_cuts_test.jsonl.gz"
+ )
+
+ return {
+ "aishell_test": aishell_test_cuts,
+ }
+
+ def aishell2_train_cuts(self) -> CutSet:
+ logging.info("About to get multidataset train cuts")
+ logging.info("Loading Aishell-2 in lazy mode")
+ aishell_2_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell2_cuts_train.jsonl.gz"
+ )
+
+ return aishell_2_cuts
+
+ def aishell2_dev_cuts(self) -> CutSet:
+ logging.info("About to get multidataset dev cuts")
+ logging.info("Loading Aishell-2 set in lazy mode")
+ aishell2_dev_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell2_cuts_dev.jsonl.gz"
+ )
+
+ return aishell2_dev_cuts
+
+ def aishell2_test_cuts(self) -> CutSet:
+ logging.info("About to get multidataset test cuts")
+ logging.info("Loading Aishell-2 set in lazy mode")
+ aishell2_test_cuts = load_manifest_lazy(
+ self.fbank_dir / "aishell2_cuts_test.jsonl.gz"
+ )
+
+ return {
+ "aishell2_test": aishell2_test_cuts,
+ }
+
+ def wenetspeech_test_meeting_cuts(self) -> CutSet:
+ logging.info("About to get multidataset test cuts")
+ logging.info("Loading WeNetSpeech set in lazy mode")
+ wenetspeech_test_meeting_cuts = load_manifest_lazy(
+ self.fbank_dir / "wenetspeech" / "cuts_TEST_MEETING.jsonl.gz"
+ )
+
+ return {
+ "wenetspeech-meeting_test": wenetspeech_test_meeting_cuts,
+ }
+
+ def speechio_test_cuts(self) -> Dict[str, CutSet]:
+ logging.info("About to get multidataset test cuts")
+ start_index = 0
+ end_index = 26
+ dataset_parts = []
+ for i in range(start_index, end_index + 1):
+ idx = f"{i}".zfill(2)
+ dataset_parts.append(f"SPEECHIO_ASR_ZH000{idx}")
+
+ prefix = "speechio"
+ suffix = "jsonl.gz"
+
+ results_dict = {}
+ for partition in dataset_parts:
+ path = f"{prefix}_cuts_{partition}.{suffix}"
+
+ logging.info(f"Loading {path} set in lazy mode")
+ test_cuts = load_manifest_lazy(self.fbank_dir / path)
+ results_dict[partition] = test_cuts
+
+ return results_dict
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/requirements.txt b/egs/speech_llm/ASR_LLM/whisper_llm_zh/requirements.txt
new file mode 100644
index 000000000..a07c7b157
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/requirements.txt
@@ -0,0 +1,11 @@
+k2
+kaldialign
+git+https://github.com/lhotse-speech/lhotse
+sentencepiece
+pypinyin
+tensorboard
+librosa
+deepspeed
+transformers>=4.37.0
+flash-attn
+peft
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/train.py b/egs/speech_llm/ASR_LLM/whisper_llm_zh/train.py
new file mode 100755
index 000000000..5f224c984
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/train.py
@@ -0,0 +1,872 @@
+#!/usr/bin/env python3
+# Copyright 2023 Xiaomi Corp. (authors: Xiaoyu Yang)
+# 2024 Yuekai Zhang
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Usage:
+# fine-tuning with whisper and Qwen2
+pip install huggingface_hub['cli']
+mkdir -p models/whisper models/qwen
+
+# For aishell fine-tuned whisper model
+huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_aishell_whisper exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt
+# For multi-hans fine-tuned whisper model
+# huggingface-cli download --local-dir models/whisper yuekai/icefall_asr_multi-hans-zh_whisper v1.1/whisper-large-v2-multi-hans-zh-epoch-3-avg-10.pt
+
+# huggingface-clie download --local-dir models/qwen Qwen/Qwen2-7B-Instruct
+huggingface-clie download --local-dir models/qwen Qwen/Qwen2-1.5B-Instruct
+
+torchrun --nproc_per_node 8 ./whisper_llm_zh/train.py \
+ --max-duration 200 \
+ --exp-dir ./whisper_llm_zh/exp_test \
+ --speech-encoder-path-or-name models/whisper/exp_large_v2/whisper-large-v2-aishell1-epoch-10-avg-6.pt \
+ --llm-path-or-name Qwen/Qwen2-1.5B-Instruct \
+ --manifest-dir data/fbank \
+ --deepspeed \
+ --deepspeed_config ./whisper_llm_zh/ds_config_zero1.json \
+ --use-flash-attn True \
+ --use-lora True --unfreeze-llm True
+"""
+
+import argparse
+import copy
+import logging
+import os
+import random
+import warnings
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import deepspeed
+import k2
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+import transformers
+import whisper
+from asr_datamodule import AsrDataModule
+from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
+from label_smoothing import LabelSmoothingLoss
+from lhotse import CutSet, load_manifest
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+from model import IGNORE_TOKEN_ID, SPEECH_LLM, EncoderProjector
+from multi_dataset import MultiDataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from torch import Tensor
+from torch.utils.tensorboard import SummaryWriter
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from whisper_encoder_forward_monkey_patch import replace_whisper_encoder_forward
+
+from icefall import diagnostics
+from icefall.dist import get_rank, get_world_size
+from icefall.env import get_env_info
+from icefall.utils import (
+ AttributeDict,
+ MetricsTracker,
+ filter_uneven_sized_batch,
+ setup_logger,
+ str2bool,
+)
+
+DEFAULT_SPEECH_TOKEN = ""
+
+
+def set_batch_count(model: nn.Module, batch_count: float) -> None:
+ for module in model.modules():
+ if hasattr(module, "batch_count"):
+ module.batch_count = batch_count
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--llm-path-or-name",
+ type=str,
+ default="/workspace/asr/Qwen1.5-0.5B-Chat",
+ help="Path or name of the large language model.",
+ )
+
+ parser.add_argument(
+ "--speech-encoder-path-or-name",
+ type=str,
+ default="whisper-large-v2",
+ help="Path or name of the speech encoder.",
+ )
+
+ parser.add_argument(
+ "--encoder-projector-ds-rate",
+ type=int,
+ default=8,
+ help="Downsample rate for the encoder projector.",
+ )
+ parser.add_argument(
+ "--use-flash-attn",
+ type=str2bool,
+ default=True,
+ help="Whether to use flash attention.",
+ )
+
+ parser.add_argument(
+ "--use-lora",
+ type=str2bool,
+ default=False,
+ help="Whether to use lora to fine-tune llm.",
+ )
+
+ parser.add_argument(
+ "--unfreeze-llm",
+ type=str2bool,
+ default=False,
+ help="Whether to unfreeze llm during training.",
+ )
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--tensorboard",
+ type=str2bool,
+ default=True,
+ help="Should various information be logged in tensorboard.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=10,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=1,
+ help="""Resume training from this epoch. It should be positive.
+ If larger than 1, it will load checkpoint from
+ exp-dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="whisper_qwen/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--pretrained-model-path",
+ type=str,
+ default=None,
+ help="""The path to the pretrained model if it is not None. Training will
+ start from this model. e.g. ./wenetspeech/ASR/whisper/exp_large_v2/epoch-4-avg-3.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--sampler-state-dict-path",
+ type=str,
+ default=None,
+ help="""The path to the sampler state dict if it is not None. Training will start from this sampler state dict.
+ """,
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=True,
+ help="Whether to use half precision training.",
+ )
+
+ parser.add_argument(
+ "--use-aishell",
+ type=str2bool,
+ default=True,
+ help="Whether to only use aishell1 dataset for training.",
+ )
+
+ parser = deepspeed.add_config_arguments(parser)
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are saved in the variable `params`.
+
+ Commandline options are merged into `params` after they are parsed, so
+ you can also access them via `params`.
+
+ Explanation of options saved in `params`:
+
+ - frame_shift_ms: The frame shift in milliseconds.
+ - allowed_excess_duration_ratio: The allowed excess duration ratio.
+ - best_train_loss: The best training loss so far.
+ - best_valid_loss: The best validation loss so far.
+ - best_train_epoch: The epoch where the best training loss is achieved.
+ - best_valid_epoch: The epoch where the best validation loss is achieved.
+ - batch_idx_train: The batch index of the current batch.
+ - log_interval: Log training stats every `log_interval` batches.
+ - reset_interval: Reset the stats every `reset_interval` batches.
+ - valid_interval: Run validation every `valid_interval` batches.
+ - env_info: The environment information.
+ """
+ params = AttributeDict(
+ {
+ "allowed_excess_duration_ratio": 0.1,
+ "subsampling_factor": 2,
+ "frame_shift_ms": 10,
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 50,
+ "reset_interval": 200,
+ "valid_interval": 5000,
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def compute_loss(
+ params: AttributeDict,
+ tokenizer: AutoTokenizer,
+ model: nn.Module,
+ batch: dict,
+ is_training: bool,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ Compute the loss for the given batch.
+ Args:
+ params:
+ It is returned by :func:`get_params`.
+ tokenizer:
+ The tokenizer used to encode the text.
+ model:
+ The model for training.
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ is_training:
+ Whether it is training.
+ Returns:
+ Return a tuple of two elements. The first element is the loss tensor.
+ """
+ # For the uneven-sized batch, the total duration after padding would possibly
+ # cause OOM. Hence, for each batch, which is sorted descendingly by length,
+ # we simply drop the last few shortest samples, so that the retained total frames
+ # (after padding) would not exceed `allowed_max_frames`:
+ # `allowed_max_frames = int(max_frames * (1.0 + allowed_excess_duration_ratio))`,
+ # where `max_frames = max_duration * 1000 // frame_shift_ms`.
+ # We set allowed_excess_duration_ratio=0.1.
+
+ def preprocess(
+ messages,
+ tokenizer: transformers.PreTrainedTokenizer,
+ max_len: int,
+ ) -> Dict:
+ """Preprocesses the data for supervised fine-tuning."""
+ texts = []
+ TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if loop.last %}{{ '<|im_end|>'}}{% else %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}"
+ for i, msg in enumerate(messages):
+ texts.append(
+ tokenizer.apply_chat_template(
+ msg,
+ tokenize=True,
+ chat_template=TEMPLATE,
+ add_generation_prompt=False,
+ padding="longest", # FIX me change padding to longest
+ max_length=max_len,
+ truncation=True,
+ )
+ )
+ # padding texts to the same length, texts is a list of list, padding with tokenzier.pad_token_id
+ max_len_texts = max([len(text) for text in texts])
+ if tokenizer.padding_side == "right":
+ texts = [
+ text + [tokenizer.pad_token_id] * (max_len_texts - len(text))
+ for text in texts
+ ]
+ else:
+ texts = [
+ [tokenizer.pad_token_id] * (max_len_texts - len(text)) + text
+ for text in texts
+ ]
+ input_ids = torch.tensor(texts, dtype=torch.int)
+ # response = tokenizer.batch_decode(input_ids, skip_special_tokens=True)[0]
+ target_ids = input_ids.clone()
+ target_ids[target_ids == tokenizer.pad_token_id] = IGNORE_TOKEN_ID
+ # mask all tokens before token_id 151646 with IGNORE_TOKEN_ID
+ # first get the indices of the tokens
+ mask_prompt = True
+ if mask_prompt:
+ mask_indices = torch.where(
+ input_ids == tokenizer.convert_tokens_to_ids("assistant")
+ )
+ for i in range(mask_indices[0].size(0)):
+ row = mask_indices[0][i]
+ col = mask_indices[1][i]
+ # + 2 to skip: 'assistant', '\n'
+ target_ids[row, : col + 2] = IGNORE_TOKEN_ID
+
+ attention_mask = input_ids.ne(tokenizer.pad_token_id)
+
+ return input_ids, attention_mask, target_ids
+
+ def normalize_text_alimeeting(text: str, normalize: str = "m2met") -> str:
+ """
+ Text normalization similar to M2MeT challenge baseline.
+ See: https://github.com/yufan-aslp/AliMeeting/blob/main/asr/local/text_normalize.pl
+ """
+ if normalize == "none":
+ return text
+ elif normalize == "m2met":
+ import re
+
+ text = text.replace(" ", "")
+ text = text.replace("", "")
+ text = text.replace("<%>", "")
+ text = text.replace("<->", "")
+ text = text.replace("<$>", "")
+ text = text.replace("<#>", "")
+ text = text.replace("<_>", "")
+ text = text.replace("", "")
+ text = text.replace("`", "")
+ text = text.replace("&", "")
+ text = text.replace(",", "")
+ if re.search("[a-zA-Z]", text):
+ text = text.upper()
+ text = text.replace("Ａ", "A")
+ text = text.replace("ａ", "A")
+ text = text.replace("ｂ", "B")
+ text = text.replace("ｃ", "C")
+ text = text.replace("ｋ", "K")
+ text = text.replace("ｔ", "T")
+ text = text.replace("，", "")
+ text = text.replace("丶", "")
+ text = text.replace("。", "")
+ text = text.replace("、", "")
+ text = text.replace("？", "")
+ return text
+
+ max_frames = params.max_duration * 1000 // params.frame_shift_ms
+ allowed_max_frames = int(max_frames * (1.0 + params.allowed_excess_duration_ratio))
+ batch = filter_uneven_sized_batch(batch, allowed_max_frames)
+
+ device = next(model.parameters()).device
+ feature = batch["inputs"]
+
+ assert feature.ndim == 3
+ feature = feature.to(device)
+ feature = feature.transpose(1, 2) # (N, C, T)
+
+ batch_idx_train = params.batch_idx_train
+ supervisions = batch["supervisions"]
+ texts = batch["supervisions"]["text"]
+ # remove spaces in texts
+ texts = [normalize_text_alimeeting(text) for text in texts]
+
+ messages = []
+ for i, text in enumerate(texts):
+ message = [
+ {"role": "user", "content": f"{DEFAULT_SPEECH_TOKEN}请转写音频为文字"},
+ {"role": "assistant", "content": text},
+ ]
+ messages.append(message)
+
+ input_ids, attention_mask, target_ids = preprocess(messages, tokenizer, max_len=128)
+
+ target_ids = target_ids.type(torch.LongTensor)
+ input_ids = input_ids.type(torch.LongTensor)
+
+ with torch.set_grad_enabled(is_training):
+ model_outputs, acc = model(
+ fbank=feature,
+ input_ids=input_ids.to(device),
+ attention_mask=attention_mask.to(device),
+ labels=target_ids.to(device),
+ )
+ loss = model_outputs.loss
+ assert loss.requires_grad == is_training
+
+ info = MetricsTracker()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ feature_lens = supervisions["num_frames"]
+ info["frames"] = (feature_lens // params.subsampling_factor).sum().item()
+
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().item()
+ info["acc"] = (
+ acc * info["frames"]
+ ) # WAR: to avoid normalization by the number of frames
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ tokenizer: whisper.tokenizer.Tokenizer,
+ model: nn.Module,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process."""
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(valid_dl):
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ tokenizer=tokenizer,
+ model=model,
+ batch=batch,
+ is_training=False,
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ tokenizer: AutoTokenizer,
+ model: nn.Module,
+ optimizer: torch.optim.Optimizer,
+ scheduler: torch.optim.lr_scheduler,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> None:
+ """Train the model for one epoch.
+
+ The training loss from the mean of all frames is saved in
+ `params.train_loss`. It runs the validation process every
+ `params.valid_interval` batches.
+
+ Args:
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The model for training.
+ optimizer:
+ The optimizer we are using.
+ scheduler:
+ The learning rate scheduler, we call step() every step.
+ train_dl:
+ Dataloader for the training dataset.
+ valid_dl:
+ Dataloader for the validation dataset.
+ scaler:
+ The scaler used for mix precision training.
+ model_avg:
+ The stored model averaged from the start of training.
+ tb_writer:
+ Writer to write log messages to tensorboard.
+ world_size:
+ Number of nodes in DDP training. If it is 1, DDP is disabled.
+ rank:
+ The rank of the node in DDP training. If no DDP is used, it should
+ be set to 0.
+ """
+ model.encoder_projector.train()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(train_dl):
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+ if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+ logging.info("Computing validation loss")
+ valid_info = compute_validation_loss(
+ params=params,
+ tokenizer=tokenizer,
+ model=model,
+ valid_dl=valid_dl,
+ world_size=world_size,
+ )
+ model.train()
+ logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+ if batch_idx != 0:
+ model.save_checkpoint(
+ save_dir=params.exp_dir,
+ tag=f"epoch-{params.cur_epoch}-checkpoint-{batch_idx}",
+ client_state={},
+ exclude_frozen_parameters=True,
+ )
+
+ if rank == 0:
+ convert_zero_checkpoint_to_fp32_state_dict(
+ params.exp_dir,
+ f"{params.exp_dir}/epoch-{params.cur_epoch}-checkpoint-{batch_idx}.pt",
+ tag=f"epoch-{params.cur_epoch}-checkpoint-{batch_idx}",
+ exclude_frozen_parameters=True,
+ )
+ # save sampler state dict into checkpoint
+ sampler_state_dict = train_dl.sampler.state_dict()
+ torch.save(
+ sampler_state_dict,
+ f"{params.exp_dir}/epoch-{params.cur_epoch}-checkpoint-{batch_idx}-sampler.pt",
+ )
+ os.system(
+ f"rm -rf {params.exp_dir}/epoch-{params.cur_epoch}-checkpoint-{batch_idx}"
+ )
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ tokenizer=tokenizer,
+ model=model,
+ batch=batch,
+ is_training=True,
+ )
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ # NOTE: We use reduction==sum and loss is computed over utterances
+ # in the batch and there is no normalization to it so far.
+
+ # deepspeed's backward() is different from torch's backward()
+ # in that it does not accept a loss tensor as input.
+ # It computes the loss internally.
+ model.backward(loss)
+ model.step()
+
+ except: # noqa
+ display_and_save_batch(batch, params=params)
+ raise
+
+ if batch_idx % params.log_interval == 0:
+ try:
+ cur_lr = scheduler.get_last_lr()[0]
+ except: # noqa
+ cur_lr = 0.0
+
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}], "
+ f"tot_loss[{tot_loss}], batch size: {batch_size}, "
+ f"lr: {cur_lr:.2e}, "
+ )
+
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ if params.train_loss < params.best_train_loss:
+ params.best_train_epoch = params.cur_epoch
+ params.best_train_loss = params.train_loss
+
+
+def run(rank, world_size, args):
+ """
+ Args:
+ rank:
+ It is a value between 0 and `world_size-1`, which is
+ passed automatically by `mp.spawn()` in :func:`main`.
+ The node with rank 0 is responsible for saving checkpoint.
+ world_size:
+ Number of GPUs for DDP training.
+ args:
+ The return value of get_parser().parse_args()
+ """
+ params = get_params()
+ params.update(vars(args))
+
+ fix_random_seed(params.seed)
+
+ setup_logger(f"{params.exp_dir}/log/log-train")
+ logging.info(params)
+
+ logging.info("About to create model")
+
+ replace_whisper_encoder_forward()
+ whisper_model = whisper.load_model(params.speech_encoder_path_or_name, "cpu")
+ speech_encoder = whisper_model.encoder
+ speech_encoder_dim = whisper_model.dims.n_audio_state
+ for name, param in speech_encoder.named_parameters():
+ param.requires_grad = False
+ speech_encoder.eval()
+
+ tokenizer = AutoTokenizer.from_pretrained(params.llm_path_or_name)
+ if params.use_flash_attn:
+ attn_implementation = "flash_attention_2"
+ # torch_dtype=torch.bfloat16 FIX ME
+ torch_dtype = torch.float16
+ tokenizer.padding_side = "left"
+
+ else:
+ attn_implementation = "eager"
+ torch_dtype = torch.float16
+ tokenizer.padding_side = "right"
+
+ llm = AutoModelForCausalLM.from_pretrained(
+ params.llm_path_or_name,
+ attn_implementation=attn_implementation,
+ torch_dtype=torch_dtype,
+ )
+
+ if not params.unfreeze_llm:
+ for name, param in llm.named_parameters():
+ param.requires_grad = False
+ llm.eval()
+ else:
+ if params.use_lora:
+ lora_config = LoraConfig(
+ r=64,
+ lora_alpha=16,
+ target_modules=[
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ "o_proj",
+ "up_proj",
+ "gate_proj",
+ "down_proj",
+ ],
+ lora_dropout=0.05,
+ task_type="CAUSAL_LM",
+ )
+ llm = get_peft_model(llm, lora_config)
+ llm.print_trainable_parameters()
+
+ special_tokens_dict = {"additional_special_tokens": [DEFAULT_SPEECH_TOKEN]}
+ tokenizer.add_special_tokens(special_tokens_dict)
+ llm.config.pad_token_id = tokenizer.pad_token_id
+ llm.config.default_speech_token_id = tokenizer.convert_tokens_to_ids(
+ DEFAULT_SPEECH_TOKEN
+ )
+
+ encoder_projector = EncoderProjector(
+ speech_encoder_dim, llm.config.hidden_size, params.encoder_projector_ds_rate
+ )
+
+ model = SPEECH_LLM(
+ speech_encoder,
+ llm,
+ encoder_projector,
+ )
+
+ if params.pretrained_model_path:
+ checkpoint = torch.load(params.pretrained_model_path, map_location="cpu")
+ missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ logging.info("Trainable parameters (excluding model.eval modules):")
+ for name, param in model.named_parameters():
+ if param.requires_grad:
+ logging.info(f"{name}: {param.shape}")
+
+ if torch.cuda.is_available():
+ device = torch.device("cuda", rank)
+ else:
+ device = torch.device("cpu")
+ logging.info(f"Device: {device}")
+ model.to(device)
+
+ assert params.deepspeed and world_size > 1
+ logging.info("Using DeepSpeed")
+ model, optimizer, _, scheduler = deepspeed.initialize(
+ args=params, model=model, model_parameters=model.parameters()
+ )
+
+ data_module = AsrDataModule(args)
+ multi_dataset = MultiDataset(args.manifest_dir)
+
+ def remove_short_and_long_utt(c: Cut):
+ # Keep only utterances with duration between 1 second and 20 seconds
+ #
+ # Caution: There is a reason to select 20.0 here. Please see
+ # ../local/display_manifest_statistics.py
+ #
+ # You should use ../local/display_manifest_statistics.py to get
+ # an utterance duration distribution for your dataset to select
+ # the threshold
+ if c.duration < 1.0 or c.duration > 20.0:
+ # logging.warning(
+ # f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ # )
+ return False
+ return True
+
+ if params.use_aishell:
+ train_cuts = multi_dataset.aishell_train_cuts()
+ else:
+ train_cuts = multi_dataset.train_cuts()
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ sampler_state_dict = None
+ if params.sampler_state_dict_path:
+ sampler_state_dict = torch.load(params.sampler_state_dict_path)
+ sampler_state_dict["max_duration"] = params.max_duration
+ # TODO: load sampler state dict
+ train_dl = data_module.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+
+ if params.use_aishell:
+ valid_cuts = multi_dataset.aishell_dev_cuts()
+ else:
+ valid_cuts = multi_dataset.dev_cuts()
+ valid_dl = data_module.valid_dataloaders(valid_cuts)
+
+ if args.tensorboard and rank == 0:
+ tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
+ else:
+ tb_writer = None
+
+ logging.info(f"start training from epoch {params.start_epoch}")
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+
+ fix_random_seed(params.seed + epoch - 1)
+ train_dl.sampler.set_epoch(epoch - 1)
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ tokenizer=tokenizer,
+ model=model,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ model.save_checkpoint(
+ save_dir=params.exp_dir,
+ tag=f"epoch-{params.cur_epoch}",
+ client_state={},
+ exclude_frozen_parameters=True,
+ )
+ if rank == 0:
+ convert_zero_checkpoint_to_fp32_state_dict(
+ params.exp_dir,
+ f"{params.exp_dir}/epoch-{params.cur_epoch}.pt",
+ tag=f"epoch-{params.cur_epoch}",
+ exclude_frozen_parameters=True,
+ )
+ # save sampler state dict into checkpoint
+ sampler_state_dict = train_dl.sampler.state_dict()
+ torch.save(
+ sampler_state_dict,
+ f"{params.exp_dir}/epoch-{params.cur_epoch}-sampler.pt",
+ )
+
+ os.system(f"rm -rf {params.exp_dir}/epoch-{params.cur_epoch}")
+
+ logging.info("Done!")
+
+
+def display_and_save_batch(
+ batch: dict,
+ params: AttributeDict,
+) -> None:
+ """Display the batch statistics and save the batch into disk.
+
+ Args:
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ params:
+ Parameters for training. See :func:`get_params`.
+ """
+ from lhotse.utils import uuid4
+
+ filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+ logging.info(f"Saving batch to {filename}")
+ torch.save(batch, filename)
+
+ supervisions = batch["supervisions"]
+ features = batch["inputs"]
+
+ logging.info(f"features shape: {features.shape}")
+
+
+def main():
+ parser = get_parser()
+ AsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ world_size = get_world_size()
+ rank = get_rank()
+
+ torch.set_num_threads(1)
+ torch.set_num_interop_threads(1)
+ run(rank=rank, world_size=world_size, args=args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/speech_llm/ASR_LLM/whisper_llm_zh/whisper_encoder_forward_monkey_patch.py b/egs/speech_llm/ASR_LLM/whisper_llm_zh/whisper_encoder_forward_monkey_patch.py
new file mode 120000
index 000000000..2a7808921
--- /dev/null
+++ b/egs/speech_llm/ASR_LLM/whisper_llm_zh/whisper_encoder_forward_monkey_patch.py
@@ -0,0 +1 @@
+../../../aishell/ASR/whisper/whisper_encoder_forward_monkey_patch.py
\ No newline at end of file