local/decode.py

fine_tuning_dir = "fine_tuned/SSD/model/Negel_79_AVA_script_conv_train_conv_dev/checkpoint-50"

from typing import Any, Dict, List, Union
from dataclasses import dataclass
from transformers import Seq2SeqTrainer
from transformers import WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor, Seq2SeqTrainingArguments, Seq2SeqTrainer, WhisperModel
import evaluate
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
from random import sample
from sys import flags
import gradio as gr
import torchaudio
import torch.nn as nn
import jiwer
import numpy as np
from rich import print as rprint
from rich.progress import track
from transformers import pipeline
import argparse
import yaml
import torch
from pathlib import Path
from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
from datasets import load_dataset, concatenate_datasets
from datasets import Dataset, Audio
import pdb
import string
import librosa
# local import
import sys

sys.path.append("src")
import lightning_module

torch.cuda.set_device("cuda:0")

audio_dir = "./data/Patient_sil_trim_16k_normed_5_snr_40"
healthy_dir = "./data/Healthy"
Fary_PAL_30 = "./data/Fary_PAL_p326_20230110_30"
John_p326 = "./data/John_p326/output"
John_video = "./data/20230103_video"
negel_79 = "./data/4_negel_79"

patient_T = "data/Patient_T/Patient_T"
patient_L = "data/Patient_L/Patient_L"
# Get Transcription, WER and PPM
"""
TODO:
    [DONE]: Automatic generating Config
"""


sys.path.append("./src")


wer = evaluate.load("wer")

# root_path = Path(__file__).parents[1]


class ChangeSampleRate(nn.Module):
    def __init__(self, input_rate: int, output_rate: int):
        super().__init__()
        self.output_rate = output_rate
        self.input_rate = input_rate

    def forward(self, wav: torch.tensor) -> torch.tensor:
        # Only accepts 1-channel waveform input
        wav = wav.view(wav.size(0), -1)
        new_length = wav.size(-1) * self.output_rate // self.input_rate
        indices = torch.arange(new_length) * (
            self.input_rate / self.output_rate
        )
        round_down = wav[:, indices.long()]
        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
            0
        ) + round_up * indices.fmod(1.0).unsqueeze(0)
        return output

# resample and clean text data


def dataclean(example):
    # pdb.set_trace()
    if example['audio']['sampling_rate'] != 16000:
        resampled_audio = librosa.resample(y=example['audio']['array'],
                                           orig_sr=example['audio']['sampling_rate'],
                                           target_sr=16000)
        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
        # resampled_audio = resampler(example['audio']['array'])

        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
    else:
        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}

processor = AutoFeatureExtractor.from_pretrained(
    "facebook/wav2vec2-base-960h"
)

def prepare_dataset(batch):
    audio = batch["audio"]
    batch = processor(
        audio["array"], sampling_rate=audio["sampling_rate"], text=batch['transcription'])
    batch["input_length"] = len(batch["input_values"][0])
    return batch


negel_79_dataset = load_dataset("audiofolder", data_dir=negel_79, split="train")
negel_79_dataset = negel_79_dataset.map(dataclean)

def train_dev_test_split(dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1):
    """
    input: dataset
    dev_rate,
    test_rate
    seed
    -------
    Output:
    dataset_dict{"train", "dev", "test"}
    """
    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
    test = train_dev_test["test"]
    train_dev = train_dev_test['train']
    
    # pdb.set_trace()
    if len(train_dev) <= int(len(dataset)*dev_rate):
        train = Dataset.from_dict({"audio": [], "transcription": []})
        dev = train_dev
    else:
        train_dev = train_dev.train_test_split(test_size=int(len(dataset)*dev_rate), seed=seed)
        train = train_dev['train']
        dev = train_dev['test']
    return train, dev, test

# pdb.set_trace()
# P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1)
# P1tony_train_ = concatenate_datasets([P1tony_train,P1tony_scripted]) 
# pdb.set_trace()

Negel_79_train, Negel_79_dev, Negel_79_test = train_dev_test_split(negel_79_dataset, dev_rate=0.1, test_rate=0.1, seed=1)
 
# src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
# src_dataset = src_dataset.map(dataclean)

# healthy_test_dataset = load_dataset(
#     "audiofolder", data_dir=healthy_dir, split='train')
# healthy_test_dataset = healthy_test_dataset.map(dataclean)

# Fary_PAL_test_dataset = load_dataset(
#     "audiofolder", data_dir=Fary_PAL_30, split='train')
# Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)

# John_p326_test_dataset = load_dataset(
#     "audiofolder", data_dir=John_p326, split='train')
# John_p326_test_dataset = John_p326_test_dataset.map(dataclean)

# John_video_test_dataset = load_dataset(
#     "audiofolder", data_dir=John_video, split='train')
# John_video_test_dataset = John_video_test_dataset.map(dataclean)

# patient_T_test_dataset = load_dataset("audiofolder", data_dir=patient_T, split='train')
# patient_T_test_dataset = patient_T_test_dataset.map(dataclean)

# patient_L_test_dataset = load_dataset("audiofolder", data_dir=patient_L, split='train')
# patient_L_test_dataset = patient_L_test_dataset.map(dataclean)
# pdb.set_trace()

# train_dev / test
# ds = src_dataset.train_test_split(test_size=0.1, seed=1)

# dataset_libri = load_dataset(
#     "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")

# train_dev = ds['train']
# # train / dev
# train_dev = train_dev.train_test_split(
#     test_size=int(len(src_dataset)*0.1), seed=1)
# # train/dev/test
# train = train_dev['train']
# test = ds['test']
# dev = train_dev['test']

# # pdb.set_trace()
# encoded_train = train.map(prepare_dataset, num_proc=4)
# encoded_dev = dev.map(prepare_dataset, num_proc=4)
# encoded_test = test.map(prepare_dataset, num_proc=4)

# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
# pdb.set_trace()

WER = evaluate.load("wer")

# Whisper decoding

processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
model = WhisperForConditionalGeneration.from_pretrained(
    "openai/whisper-medium").to("cuda:0")
tokenizer = WhisperTokenizer.from_pretrained(
    "openai/whisper-medium", language="English", task="transcribe")

# Need to push tokenizer to hugginface/model to activate online API

# tokenizer.push_to_hub("KevinGeng/whipser_medium_en_PAL300_step25")
# import pdb
# pdb.set_trace()

feature_extractor = WhisperFeatureExtractor.from_pretrained(
    "openai/whisper-medium")


def whisper_prepare_dataset(batch):
    # load and resample audio data from 48 to 16kHz
    audio = batch["audio"]

    # compute log-Mel input features from input audio array
    batch["input_features"] = feature_extractor(
        audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]

    # encode target text to label ids
    batch["labels"] = tokenizer(batch["transcription"]).input_ids
    return batch


torch.cuda.empty_cache()

training_args = Seq2SeqTrainingArguments(
    # change to a repo name of your choice
    output_dir="./whisper-medium-PAL128-25step",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
    learning_rate=1e-5,
    warmup_steps=100,
    max_steps=1000,
    gradient_checkpointing=True,
    fp16=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=8,
    predict_with_generate=True,
    generation_max_length=512,
    save_steps=100,
    eval_steps=25,
    logging_steps=100,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=True,
)


def my_map_to_pred(batch):
    # pdb.set_trace()
    audio = batch["audio"]
    input_features = processor(
        audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])

    with torch.no_grad():
        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
        predicted_ids = model.generate(input_features.to("cuda"))[0]
    transcription = model.decode(predicted_ids)
    batch["prediction"] = model.tokenizer._normalize(transcription)
    return batch


@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # split inputs and labels since they have to be of different lengths and need different padding methods
        # first treat the audio inputs by simply returning torch tensors
        input_features = [{"input_features": feature["input_features"]}
                          for feature in features]
        batch = self.processor.feature_extractor.pad(
            input_features, return_tensors="pt")

        # get the tokenized label sequences
        label_features = [{"input_ids": feature["labels"]}
                          for feature in features]
        # pad the labels to max length
        labels_batch = self.processor.tokenizer.pad(
            label_features, return_tensors="pt")

        # replace padding with -100 to ignore loss correctly
        labels = labels_batch["input_ids"].masked_fill(
            labels_batch.attention_mask.ne(1), -100)

        # if bos token is appended in previous tokenization step,
        # cut bos token here as it's append later anyways
        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch


data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)


def compute_metrics(pred):
    pdb.set_trace()
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    # replace -100 with the pad_token_id
    label_ids[label_ids == -100] = tokenizer.pad_token_id

    # we do not want to group tokens when computing the metrics
    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    wer = 100 * WER.compute(predictions=pred_str, references=label_str)

    return {"wer": wer}

encode_negel_79_train = Negel_79_train.map(whisper_prepare_dataset, num_proc=4)
encode_negel_79_dev = Negel_79_dev.map(whisper_prepare_dataset, num_proc=4)
encode_negel_79_test = Negel_79_test.map(whisper_prepare_dataset, num_proc=4)
pdb.set_trace()
torch.cuda.empty_cache()

torch.cuda.empty_cache()

fine_tuned_model = WhisperForConditionalGeneration.from_pretrained(
    fine_tuning_dir
).to("cuda")
    # "fine_tuned/SSD/model/whipser_medium_TEP_patient_T"
    # "./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400"
    #"./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-200"


def fine_tuned_map_to_pred(batch):
    # pdb.set_trace()
    audio = batch["audio"]
    input_features = processor(
        audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])

    with torch.no_grad():
        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
        predicted_ids = fine_tuned_model.generate(input_features.to("cuda"))[0]
    transcription = tokenizer.decode(predicted_ids)
    batch["prediction"] = tokenizer._normalize(transcription)
    return batch


    # output_dir="./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400",
testing_args = Seq2SeqTrainingArguments(
    # change to a repo name of your choice
    output_dir="fine_tuned/SSD/model/whipser_medium_TEP_patient_TL_TL",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
    learning_rate=1e-5,
    warmup_steps=100,
    max_steps=1000,
    gradient_checkpointing=True,
    fp16=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=8,
    predict_with_generate=True,
    generation_max_length=512,
    save_steps=100,
    eval_steps=25,
    logging_steps=100,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=False,
)

predict_trainer = Seq2SeqTrainer(
    args=testing_args,
    model=fine_tuned_model,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=processor.feature_extractor,
)

# trainer.train()
# fine tuned
# z_result = encoded_test.map(fine_tuned_map_to_pred)
pdb.set_trace()
z_result= encode_negel_79_test.map(fine_tuned_map_to_pred)
# 0.4692737430167598
z = WER.compute(references=z_result['reference'], predictions=z_result['prediction'])
# pdb.set_trace()
# z_hel_result = encoded_healthy.map(fine_tuned_map_to_pred)
# z_hel = WER.compute(references=z_hel_result['reference'], predictions=z_hel_result['prediction'])
# # 0.1591610117211598

# # pdb.set_trace()
# # z_fary_result = encoded_Fary.map(fine_tuned_map_to_pred)
# # z_far = WER.compute(references=z_fary_result['reference'], predictions=z_fary_result['prediction'])
# # 0.1791044776119403
# z_patient_LT = encoded_patient_TL_test.map(fine_tuned_map_to_pred)
# z_patient_LT_result = WER.compute(references=z_patient_LT['reference'], predictions=z_patient_LT['prediction'])
# z_patient_L = encoded_patient_L_test.map(fine_tuned_map_to_pred)
# z_patient_L_result = WER.compute(references=z_patient_L['reference'], predictions=z_patient_L['prediction'])
# z_patient_T = encoded_patient_T_test.map(fine_tuned_map_to_pred)
# z_patient_T_result = WER.compute(references=z_patient_T['reference'], predictions=z_patient_T['prediction'])

# # z_john_p326_result = encoded_John_p326.map(fine_tuned_map_to_pred)
# # pdb.set_trace() 

# # z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
# # 0.4648241206030151
pdb.set_trace() 

# # y_John_video= fine_tuned_trainer.predict(encoded_John_video)
# # metrics={'test_loss': 2.665189743041992, 'test_wer': 0.7222222222222222, 'test_runtime': 0.1633, 'test_samples_per_second': 48.979, 'test_steps_per_second': 6.122})
# pdb.set_trace()

# p326 training
# metrics={'test_loss': 0.4804028868675232, 'test_wer': 0.21787709497206703, 'test_runtime': 0.3594, 'test_samples_per_second': 44.517, 'test_steps_per_second': 5.565})
# hel metrics={'test_loss': 1.6363693475723267, 'test_wer': 0.17951881554595928, 'test_runtime': 3.8451, 'test_samples_per_second': 41.611, 'test_steps_per_second': 5.201})
# Fary: metrics={'t est_loss': 1.4633615016937256, 'test_wer': 0.5572139303482587, 'test_runtime': 0.6627, 'test_samples_per_second': 45.27, 'test_steps_per_second': 6.036})
# p326 large: metrics={'test_loss': 0.6568527817726135, 'test_wer': 0.2889447236180904, 'test_runtime': 0.7169, 'test_samples_per_second': 51.613, 'test_steps_per_second': 6.975})