Data collection and fine-tuning Whisper for low-resource languages

import json
import csv
import re

# Read the JSON file
input_file = "/home/kiwoong/Desktop/AIhub/src/stagecoach/data/Chichewa/updated_consolidated_speech_data.json"  # Chichewa
output_file = "data/Chichewa/Chichewa.tsv"  # Output TSV file path

# Open and parse the JSON file
with open(input_file, "r", encoding="utf-8") as f:
    data = json.load(f)

# Write to a TSV file
with open(output_file, "w", encoding="utf-8", newline='') as tsvfile:
    writer = csv.DictWriter(tsvfile, fieldnames=["filepath", "text", "split"], delimiter='\t')
    writer.writeheader()  # Write header

    # Write rows with updated keys and modified filepath
    for entry in data:
        writer.writerow({
            "filepath": f"data/Chichewa/{entry['FILE_PATH']}",
            "text": re.sub('\n', ' ', entry["TRANSCRIPTION"]),
            "split": entry["Split"]
        })

print(f"Data successfully written to {output_file}")

import os
import wave

def get_wav_duration(file_path):
    with wave.open(file_path, 'rb') as wav_file:
        num_frames = wav_file.getnframes()
        frame_rate = wav_file.getframerate()
        duration = num_frames / float(frame_rate)
        return duration

def get_total_wav_duration(directory):
    total_duration = 0
    count = 0
    for filename in os.listdir(directory):
        if filename.endswith('.wav'):
            count += 1
            file_path = os.path.join(directory, filename)
            total_duration += get_wav_duration(file_path)
            
    # Convert total duration to hours
    total_duration_hours = total_duration / 3600
    avg_dur = total_duration/count
    return total_duration_hours, avg_dur

directory = '/home/kiwoong/Desktop/AIhub/src/stagecoach/data/Yoruba/train' # Calculating number of hours in Yoruba training data 
total_duration_hours, avg_dur = get_total_wav_duration(directory)
print(f"Total duration of all .wav files: {total_duration_hours} hours, avg_dur : {avg_dur}")

"filter_toolong": {
                "module": "byLength",
                "name": "ByLength",
                "args": [],
                "kwargs": {
                    "fields": ["source", "target"],
                    "maxLen": 180, #self.filter_len
                }
            }

class FilterBase:
    def __call__(self, data):
        keep = []
        for x in data:
            if self.use(x):
                keep.append(x)
        return keep


from blocks.filterBlk.filterBase import FilterBase

class ByLength(FilterBase):
    def __init__(self, fields, maxLen):
        self.fields = fields
        self.maxLen = maxLen

    def use(self, sample):
        for f in self.fields:
            if type(sample[f]) == str:
                if len(sample[f]) > self.maxLen:
                    return False
            elif type(sample[f]).__name__ == "Tensor":
                if sample[f].shape[0] > self.maxLen:
                    return False
        return True

import pandas as pd

ch = pd.read_csv('Chichewa.tsv', sep = '\t')
sw = pd.read_csv('swahili.tsv', sep = '\t')

sw = sw[:30000]
merged = pd.concat([ch, sw])
merged['split'].value_counts()
'''
split
Train    30267
Test        41
Dev         35
Name: count, dtype: int64
'''

field_limit = 131072  # Default CSV field size limit

# Filter rows without creating a new column
merged = merged[merged['text'].str.len() <= field_limit]

merged['split'].value_counts()
'''
split
Train    30264 (3 items were removed from the Swahili dataset)
Test        41
Dev         35
Name: count, dtype: int64'''

import os
import torch
import torchaudio
from datasets import load_dataset
from transformers import (
    WhisperProcessor,
    WhisperForConditionalGeneration,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
    DataCollatorSpeechSeq2SeqWithPadding
)
from evaluate import load
import numpy as np
from dotenv import load_dotenv
from torch.optim import Adafactor
from torch.optim.lr_scheduler import get_linear_schedule_with_warmup
from tqdm.auto import trange
import wandb

# Load environment variables
load_dotenv(os.path.expanduser('~/.env'))

# Configuration
class Config:
    def __init__(self):
        self.seed = 1000
        self.model_name = "openai/whisper-large-v3-turbo"
        self.data_source = "/home/kiwoong/Desktop/AIhub/src/stagecoach/data/Yoruba/Yoruba.tsv"
        self.save_model = "WhisperTestModel"
        self.save_path = f"models/forward/{self.save_model}"
        self.batch_size = 2
        self.grad_accum = 8
        self.epochs = 20
        self.learning_rate = 1e-5
        self.weight_decay = 1e-3
        self.warmup_steps = 5000
        self.max_steps = 50000
        self.filter_len = 180
        self.epd = 2  # epochs per dev evaluation
        self.eval_batch_size = 6

class LoadWhisper:
    def __init__(self, name, device):
        self.processor = WhisperProcessor.from_pretrained(name)
        self.model = WhisperForConditionalGeneration.from_pretrained(name)
        self.tokenizer = self.processor.tokenizer
        self.model.to(device)
        self.device = device
        self.default_language = "yo"  # Yoruba language code

    def genAudio2Text(self, audio_batch, withGrad=False, temp=0, language=None):
        audio_batch = [audio.squeeze().numpy() for audio in audio_batch]
        inputs = self.processor(audio_batch, return_tensors="pt", sampling_rate=16000).to(self.device)
        language = language or self.default_language

        if language:
            conditioning_tokens = [
                "<|startoftranscript|>",
                f"<|{language}|>",
                "<|transcribe|>"
            ]
            conditioning_input_ids = torch.tensor(
                [self.processor.tokenizer.convert_tokens_to_ids(conditioning_tokens)]
            ).to(self.device)
            batch_size = inputs.input_features.size(0)
            conditioning_input_ids = conditioning_input_ids.repeat(batch_size, 1)
        else:
            conditioning_input_ids = None

        self.model.eval() if not withGrad else self.model.train()
        
        with torch.set_grad_enabled(withGrad):
            outputs = self.model.generate(
                input_features=inputs.input_features,
                forced_decoder_ids=conditioning_input_ids,
                max_length=448,
                num_beams=4,
                temperature=temp or 1.0
            )
        
        return outputs

    def genHS(self, input_features):
        return self.model(input_features=input_features, return_dict=True)

    def genLossFromHS(self, hidden_states, labels):
        return self.model(inputs_embeds=hidden_states, labels=labels).loss

class LoadLocalAudio:
    def __init__(self, path, tokenizer, **kwargs):
        self.path = path
        self.tokenizer = tokenizer
        self.extract_column = kwargs.get('extract_column', {
            'audio': 'filepath',
            'source': 'text',
            'target': 'text',
            'splits': 'split',
            'src_data': 'filepath',
            'tgt_data': 'text'
        })
        self.partitions = kwargs.get('partitions', {'train': 80, "test": 10, "dev": 10})
        self.mel = kwargs.get('mel', False)
        self.sample_rate = 16000
        self.metrics = kwargs.get('metrics', {})

    def load_dataset(self):
        dataset = load_dataset('csv', data_files=self.path, delimiter='\t')
        
        # Split dataset according to partitions
        total = sum(self.partitions.values())
        train_size = self.partitions['train'] / total
        test_size = self.partitions['test'] / total
        
        dataset = dataset['train'].train_test_split(test_size=1-train_size, seed=1000)
        test_valid = dataset['test'].train_test_split(test_size=test_size/(1-train_size), seed=1000)
        
        return {
            'train': dataset['train'],
            'dev': test_valid['train'],
            'test': test_valid['test']
        }

    def resample_audio(self, audio_path, target_sr, mel=None):
        waveform, sample_rate = torchaudio.load(audio_path)
        if sample_rate != target_sr:
            resampler = torchaudio.transforms.Resample(sample_rate, target_sr)
            waveform = resampler(waveform)
        
        if mel:
            mel_spectrogram = torchaudio.transforms.MelSpectrogram(
                sample_rate=target_sr, n_mels=80
            )
            waveform = mel_spectrogram(waveform)
        
        return waveform

    def prepare_dataset(self, batch):
        audio = batch[self.extract_column['audio']]
        text = batch[self.extract_column['text']]
        
        # Process audio
        waveform = self.resample_audio(audio, self.sample_rate, self.mel)
        input_features = self.processor(
            waveform.squeeze().numpy(),
            sampling_rate=self.sample_rate,
            return_tensors="pt"
        ).input_features[0]
        
        # Process text
        labels = self.tokenizer(text, return_tensors="pt").input_ids[0]
        
        return {
            "input_features": input_features,
            "labels": labels
        }

    def evaluate(self, targets, predictions):
        wer_metric = load("wer")
        cer_metric = load("cer")
        
        wer = wer_metric.compute(predictions=predictions, references=targets)
        cer = cer_metric.compute(predictions=predictions, references=targets)
        
        return {"wer": wer, "cer": cer}

class RunModelAudio:
    def __init__(self, model, data, target, batchSize=6, initial=False, saveResult=False, language=None):
        self.model = model
        self.data = data
        self.target = target
        self.batchSize = batchSize
        self.saveResult = saveResult
        self.language = language
        self.evalSet = getattr(data, target, None)

    def testAudio(self):
        res = {}
        src, tgt, trans = [], [], []

        for i in range(0, len(self.evalSet), self.batchSize):
            samps = self.evalSet[i:i+self.batchSize]
            tgt.extend([s["target"] for s in samps])
            srcToks = [s["src_data"] for s in samps]
            
            result = self.model.genAudio2Text(srcToks, language=self.language)
            
            print(f'{i} Source: Audio Data\n{" "*len(str(i))} Target: {samps[0]["target"]}\n{" "*len(str(i))} fTrans: {result[0]}')
            
            src.extend([s["source"] for s in samps])
            trans.extend(result)
        
        res[self.target] = self.data.evaluate([t.lower() for t in tgt], [tr.lower() for tr in trans])
        print("RunModel: res[%s] = %s" % (self.target, res[self.target]))
        
        if self.saveResult:
            self.save_translations(src, tgt, trans, res)
        
        return res

    def save_translations(self, src, tgt, trans, res):
        if not self.saveResult:
            return
        
        os.makedirs(os.path.dirname(self.saveResult), exist_ok=True)
        with open(self.saveResult, 'w', newline='') as f:
            writer = csv.writer(f, delimiter='\t')
            writer.writerow(['Source', 'Target', 'Translation', 'WER', 'CER'])
            for s, t, tr in zip(src, tgt, trans):
                writer.writerow([s, t, tr, res[self.target]['wer'], res[self.target]['cer']])

class SupervisedAudio:
    def __init__(self, model, batchSz, epochs, optimizer, train, grad_accum=1, dev=None, epochsPerDev=1, evalBatchSz=6, test=None, dtype="float32", savePath=None):
        self.model = model
        self.batchSz = batchSz
        self.epochs = epochs
        self.optimizer, self.scheduler = optimizer(self.model.model)
        self.train = train
        self.gradAccum = grad_accum
        self.dev = dev
        self.test = test
        self.epochsPerDev = epochsPerDev
        self.dtype = getattr(torch, dtype, torch.float32)
        self.scaler = torch.cuda.amp.GradScaler() if dtype != "float32" else None
        self.savePath = savePath

    def supervisedBatch(self, x, y, doBackward=True):
        try:
            input_features = x.to(self.model.device, dtype=self.dtype)
            labels = y['input_ids'].to(self.model.device)
            attention_mask = y['attention_mask'].to(self.model.device)

            self.model.train()
            loss = self.model(input_features=input_features, labels=labels).loss

            if doBackward:
                if self.scaler:
                    self.scaler.scale(loss).backward()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    loss.backward()
                    self.optimizer.step()

                self.optimizer.zero_grad(set_to_none=True)
                self.scheduler.step()

            return loss.item()

        except RuntimeError as e:
            self.optimizer.zero_grad(set_to_none=True)
            print("RuntimeError:", e)
            return 0

    def train(self):
        for epoch in range(self.epochs):
            total_loss = 0
            num_batches = 0
            
            for i in range(0, len(self.train), self.batchSz):
                batch = self.train[i:i+self.batchSz]
                x = [b["input_features"] for b in batch]
                y = {"input_ids": torch.stack([b["labels"] for b in batch])}
                
                loss = self.supervisedBatch(x, y)
                total_loss += loss
                num_batches += 1
                
                if num_batches % self.gradAccum == 0:
                    print(f"Epoch {epoch+1}/{self.epochs}, Batch {num_batches}, Loss: {loss:.4f}")
            
            avg_loss = total_loss / num_batches
            print(f"Epoch {epoch+1}/{self.epochs}, Average Loss: {avg_loss:.4f}")
            
            if (epoch + 1) % self.epochsPerDev == 0:
                self.evaluate()
            
            if self.savePath:
                self.save_model()

    def evaluate(self):
        if self.dev:
            dev_results = self.dev.testAudio()
            print(f"Dev results: {dev_results}")
        
        if self.test:
            test_results = self.test.testAudio()
            print(f"Test results: {test_results}")

    def save_model(self):
        if not self.savePath:
            return
        
        os.makedirs(os.path.dirname(self.savePath), exist_ok=True)
        self.model.model.save_pretrained(self.savePath)
        self.model.processor.save_pretrained(self.savePath)

def main():
    # Set random seed
    torch.manual_seed(1000)
    
    # Initialize configuration
    config = Config()
    
    # Initialize model and processor
    model_handler = LoadWhisper(config.model_name, "cuda")
    processor = model_handler.processor
    model = model_handler.model
    
    # Load and prepare dataset
    data_handler = LoadLocalAudio(
        config.data_source,
        processor.tokenizer,
        extract_column={
            'audio': 'filepath',
            'source': 'text',
            'target': 'text',
            'splits': 'split',
            'src_data': 'filepath',
            'tgt_data': 'text'
        },
        partitions={'train': 80, "test": 10, "dev": 10}
    )
    
    datasets = data_handler.load_dataset()
    
    # Prepare datasets
    for split in datasets:
        datasets[split] = datasets[split].map(
            lambda x: data_handler.prepare_dataset(x),
            remove_columns=datasets[split].column_names
        )
    
    # Initialize optimizer
    optimizer = lambda model: (
        Adafactor(
            model.parameters(),
            scale_parameter=False,
            relative_step=False,
            lr=config.learning_rate,
            clip_threshold=1.0,
            weight_decay=config.weight_decay
        ),
        get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=config.warmup_steps,
            num_training_steps=config.max_steps
        )
    )
    
    # Initialize evaluation
    eval_dev = RunModelAudio(model_handler, data_handler, "dev", batchSize=config.eval_batch_size)
    eval_test = RunModelAudio(model_handler, data_handler, "test", batchSize=config.eval_batch_size)
    
    # Initialize trainer
    trainer = SupervisedAudio(
        model_handler,
        config.batch_size,
        config.epochs,
        optimizer,
        datasets["train"],
        grad_accum=config.grad_accum,
        dev=eval_dev,
        epochsPerDev=config.epd,
        evalBatchSz=config.eval_batch_size,
        test=eval_test,
        savePath=config.save_path
    )
    
    # Start training
    trainer.train()

if __name__ == "__main__":
    main()