MMaDA/inference_s2t_WER.py

# coding=utf-8
# Copyright 2025 AIDAS Lab
#
# 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 os
import random
import editdistance
from functools import partial
from normalizer import data_utils

os.environ["TOKENIZERS_PARALLELISM"] = "true"

from tqdm import tqdm
import torch
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP

import wandb
from datasets import load_dataset
from models import MMadaModelLM
from models.modeling_emova_speech_tokenizer import EMOVASpeechTokenizer
from training.data import S2T_INSTRUCTION
from training.prompting_utils import UniversalPrompting
from training.utils import get_config, flatten_omega_conf
from transformers import AutoTokenizer

import argparse
import logging 
import re

os.environ["TOKENIZERS_PARALLELISM"] = "true"

from tqdm import tqdm
import torch
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP

import wandb
from datasets import load_dataset
from models import MMadaModelLM
from models.modeling_emova_speech_tokenizer import EMOVASpeechTokenizer

from training.data import S2T_INSTRUCTION
from training.prompting_utils import UniversalPrompting
from training.utils import get_config, flatten_omega_conf
from transformers import AutoTokenizer

def setup_logger(rank):

    logger = logging.getLogger(__name__)
    
    # 핸들러 중복 추가 방지
    if logger.hasHandlers():
        logger.handlers.clear()

    formatter = logging.Formatter(f'%(asctime)s - [RANK {rank}] - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    
    ch = logging.StreamHandler()
    ch.setFormatter(formatter)
    
    logger.addHandler(ch)

    if rank == 0:
        logger.setLevel(logging.INFO)
    else:
        logger.setLevel(logging.WARNING)
        
    return logger

def calculate_WER(recognized_text_list, groundtruth_text_list):
    """Calculates the Word Error Rate (WER) between predicted and ground truth texts."""
    word_num = 0.0
    scores = 0.0
    for recognized_text, groundtruth_text in zip(recognized_text_list, groundtruth_text_list):
        
        recognized_text = recognized_text.lower()
        groundtruth_text = groundtruth_text.lower()

        recognized_text = re.sub(r"[^\w\s']", "", recognized_text)
        groundtruth_text = re.sub(r"[^\w\s']", "", groundtruth_text)
        
        recognized_word_list = recognized_text.split()
        groundtruth_word_list = groundtruth_text.split()
        
        current_word_num = len(groundtruth_word_list)
        word_num += current_word_num
        
        current_score = editdistance.eval(recognized_word_list, groundtruth_word_list)
        scores += current_score
        
    WER = scores / word_num if word_num > 0 else 0.0
    return WER, scores, word_num

def get_vq_model_class(model_type):
    """Returns the speech tokenizer model class based on the model type."""
    if model_type == "magvitv2":
        raise NotImplementedError("MAGVITv2 is not implemented in this script.")
    elif model_type == "emova":
        return EMOVASpeechTokenizer
    else:
        raise ValueError(f"model_type {model_type} not supported.")

def get_librispeech_dataset(logger):
    """Loads the Librispeech ASR dataset (test-clean split) from Hugging Face."""
    logger.info("Loading EMOVA dataset (clean/test)...")
    dataset = load_dataset("Emova-ollm/emova-asr-tts-eval/", "librispeech-asr-tts")['test']
    logger.info("Dataset loaded successfully.")
    return dataset

def form_ann_rst_list(ann, results, key):

    ann_dict = {}
    for item in ann:
        if key in item['id']:
            ann_dict[item['id']] = item['conversations'][-1]['value']
    
    rst_dict = {}
    for item in results:
        if key in item['id']:
            rst_dict[item['id']] = item['text']
    
    return ann_dict, rst_dict
    
# --- DDP Setup and Cleanup Functions ---

def setup_distributed(rank, world_size):
    """Initializes the distributed process group."""
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup_distributed():
    """Cleans up the distributed process group."""
    dist.destroy_process_group()

# --- Custom Dataset and Collate Function ---

class LibrispeechEvalDataset(Dataset):
    def __init__(self, hf_dataset, root_path, vq_model, text_vocab_size, image_vocab_size):
        self.hf_dataset = hf_dataset
        self.root_path = root_path
        self.vq_model = vq_model 
        self.text_vocab_size = text_vocab_size
        self.image_vocab_size = image_vocab_size

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        example = self.hf_dataset[idx]
        gt_text = example['text']
        sample_id = example['id']

        speaker_id, chapter_id, _ = sample_id.split('-')
        audio_path = os.path.join(self.root_path, speaker_id, chapter_id, f"{sample_id}.flac")

        if not os.path.exists(audio_path):
            return None

        speech_token_ids = self.vq_model.encode(audio_path)
        speech_token_ids += self.text_vocab_size + self.image_vocab_size

        return {
            "speech_token_ids": speech_token_ids,
            "gt_text": gt_text,
            "sample_id": sample_id
        }

def evaluation_collate_fn(batch, text_tokenizer, uni_prompting, config):
    batch = [b for b in batch if b is not None]
    if not batch:
        return None

    device = batch[0]["speech_token_ids"].device

    max_text_len = config.dataset.preprocessing.max_seq_length
    max_audio_len = config.dataset.preprocessing.max_aud_length + 1
    audio_pad_id = 126093
    sptids_dict = uni_prompting.sptids_dict
    
    batched_input_ids = [] 
    gt_texts = [item["gt_text"] for item in batch]
    sample_ids = [item["sample_id"] for item in batch]

    for item in batch:
        current_audio_tokens = item["speech_token_ids"].to(device)
        task_tensor = sptids_dict['<|s2t|>'].to(device).unsqueeze(0)
        soa_tensor = sptids_dict['<|soa|>'].to(device).unsqueeze(0)
        eoa_tensor = sptids_dict['<|eoa|>'].to(device).unsqueeze(0)

        effective_max_audio = max_audio_len - 3
        if current_audio_tokens.shape[1] > effective_max_audio:
            current_audio_tokens = current_audio_tokens[:, :effective_max_audio]
        
        audio_block = torch.cat([task_tensor, soa_tensor, current_audio_tokens, eoa_tensor], dim=1)
        
        num_padding = max_audio_len - audio_block.shape[1]
        if num_padding > 0:
            padding_tensor = torch.full((1, num_padding), audio_pad_id, dtype=torch.long, device=device)
            padded_audio_block = torch.cat([padding_tensor, audio_block], dim=1)
        else:
            padded_audio_block = audio_block

        chosen_prompt = random.choice(S2T_INSTRUCTION)
        prompt_text = f'<|start_header_id|>user<|end_header_id|>\n{chosen_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n'
        
        prompt_encoding = text_tokenizer(
            prompt_text,
            max_length=max_text_len,
            truncation=True,
            return_tensors="pt"
        )
        prompt_tensor = prompt_encoding.input_ids.to(device)

        final_sequence = torch.cat([padded_audio_block, prompt_tensor], dim=1)
        batched_input_ids.append(final_sequence.squeeze(0))

    pad_token_id = 126093
    
    max_len = max(seq.size(0) for seq in batched_input_ids)

    
    final_batch = torch.full((len(batched_input_ids), max_len), 
                             pad_token_id, 
                             dtype=torch.long, 
                             device=device)

    
    for i, seq in enumerate(batched_input_ids):
        final_batch[i, -len(seq):] = seq
    
    return {
        "input_ids": final_batch,
        "gt_texts": gt_texts,
        "sample_ids": sample_ids
    }

def main():
    """Main function to run the distributed evaluation."""
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    setup_distributed(rank, world_size)
    device = torch.device(f"cuda:{rank}")

    logger = setup_logger(rank)
    parser = argparse.ArgumentParser(description="Run DDP evaluation for MMadaModelLM.")
    parser.add_argument('--train_step', type=int, required=True, help='The training step of the checkpoint to evaluate.')
    parser.add_argument('--remasking', type=str, default='random', help='Remasking Strategy.')
    parser.add_argument('--generation_step', type=int, default=512, help='The training step of the checkpoint to evaluate.')
    parser.add_argument('--new_tok', type=int, default=256, help='The training step of the checkpoint to evaluate.')
    args, unknown = parser.parse_known_args()
    config = get_config()

    if rank == 0:
        run_id = config.wandb.get("run_id", None) or wandb.util.generate_id()
        config.wandb.run_id = run_id
        wandb_config = {k: v for k, v in flatten_omega_conf(config, resolve=True)}
        wandb.init(
            project="librispeech_test-clean",
            name=config.experiment.name + f'_STEP-{args.train_step}_Remasking-{args.remasking}_GS-{args.generation_step}_NT-{args.new_tok}',
            config=wandb_config,
        )

    text_tokenizer = AutoTokenizer.from_pretrained(config.model.omada.pretrained_model_path, padding_side="left")
    
    uni_prompting = UniversalPrompting(text_tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
                                       special_tokens=("<|s2t|>", "<|soa|>", "<|eoa|>", "<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),
                                       ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob, use_reserved_token=True)
    
    vq_model_class = get_vq_model_class(config.model.vq_model_audio.type)
    vq_model = vq_model_class.from_pretrained(config.model.vq_model_audio.vq_model_name).to(device)
    vq_model.requires_grad_(False)
    vq_model.eval()

    train_step = args.train_step
    trained_checkpoint_path = f"/home/work/AIDAS/ckpts/omada/omada-training-stage1/checkpoint-{train_step}/unwrapped_model/"
    
    if rank == 0:
        logger.info(f"Loading trained model from: {trained_checkpoint_path}")

    model = MMadaModelLM.from_pretrained(
        trained_checkpoint_path, 
        trust_remote_code=True, 
        torch_dtype=torch.bfloat16,
        config="/home/work/AIDAS/ckpts/omada/omada-training-stage1/config.json"
    ).to(device)

    model = DDP(model, device_ids=[rank])
    if rank == 0:
        logger.info("✅ Trained model loaded and wrapped with DDP successfully!")

    text_vocab_size = len(uni_prompting.text_tokenizer)
    image_vocab_size = config.model.omada.codebook_size

    # --- Setup DataLoader ---
    hf_dataset = get_librispeech_dataset(logger)
    root_path = "/home/work/AIDAS/data/audio/LibriSpeech/test-clean"
    
    eval_dataset = LibrispeechEvalDataset(hf_dataset, root_path, vq_model, text_vocab_size, image_vocab_size)
    sampler = DistributedSampler(eval_dataset, num_replicas=world_size, rank=rank, shuffle=False)
    
    collate_for_eval = partial(
        evaluation_collate_fn,
        text_tokenizer=text_tokenizer,
        uni_prompting=uni_prompting,
        config=config
    )
    
    dataloader = DataLoader(
        eval_dataset,
        batch_size=16,
        sampler=sampler,
        num_workers=0, 
        collate_fn=collate_for_eval,
        pin_memory=True
    )

    # --- Evaluation Loop ---
    local_results = []
    model.eval()
    
    progress_bar = tqdm(dataloader, desc="Evaluating on Librispeech", disable=(rank != 0))
    for batch_idx, batch in enumerate(progress_bar):
        # if batch_idx > 1:
        #     break
        if batch is None:
            continue
            
        input_ids = batch["input_ids"].to(device)
        gt_texts = batch["gt_texts"]
        sample_ids = batch["sample_ids"]

        with torch.no_grad():
            output_ids = model.module.mmu_generate(input_ids, max_new_tokens=args.new_tok, steps=args.generation_step, block_length=args.new_tok, remasking=args.remasking)
            decoded_texts = text_tokenizer.batch_decode(output_ids[:, input_ids.shape[1]:], skip_special_tokens=True)

            # print(decoded_texts)
            for i in range(len(decoded_texts)):
                local_results.append({
                    "sample_id": sample_ids[i],
                    "gt_text": gt_texts[i],
                    "decoded_text": decoded_texts[i]
                })

                if rank == 0 and i == 0:
                    logger.info(f"  ID: {sample_ids[i]}")
                    logger.info(f"  GT: {data_utils.normalizer(gt_texts[i])}") 
                    logger.info(f"  PD: {data_utils.normalizer(decoded_texts[i])}") 

    # --- Gather Results from All GPUs ---
    all_results = [None] * world_size
    dist.all_gather_object(all_results, local_results)

    # --- Final Processing and Logging (only on rank 0) ---
    if rank == 0:
        logger.info("Gathering and processing results from all GPUs...")
        final_results = [item for sublist in all_results for item in sublist]
        
        groundtruth_text_list = [data_utils.normalizer(res["gt_text"]) for res in final_results]
        recognized_text_list = [data_utils.normalizer(res["decoded_text"]) for res in final_results]
        
        results_table = wandb.Table(columns=["ID", "Ground Truth", "Response"])
        for res in final_results:
            results_table.add_data(res["sample_id"], res["gt_text"], res["decoded_text"])

        wandb.log({"Speech-to-Text Response Examples": results_table})
        
        wer, errors, words = calculate_WER(groundtruth_text_list, recognized_text_list)
        logger.info(f"Final WER (Librispeech test-clean): {wer:.4f}  | Word Errors: {errors}  | Total Words: {words}")

        wandb.log({
            "WER": wer,
            "Total Word Errors": errors,
            "Total Words": words
        })

    # --- Cleanup ---
    if rank == 0:
        wandb.finish()
    cleanup_distributed()


if __name__ == '__main__':
    main()