Fine-Tuning Whisper for Low-Resource Dravidian Languages

import os
from keras.utils import get_file

def download_archive(url):
    filename = os.path.split(url)[-1]
    try:
        filename = get_file(fname=filename, origin=url, cache_dir="./", extract=False)
    except:
        print(f"Unable to download {url}")
        return None
    return filename


import pathlib
from zipfile import ZipFile

def extract_zipfile(filename):
    filename = pathlib.Path(filename)
    dirname = filename.parent / filename.stem
    with ZipFile(filename, "r") as zipf:
        for zipinfo in zipf.infolist():
            if zipinfo.filename[-1] == "/":
                continue
            zipinfo.filename = pathlib.Path(zipinfo.filename).name
            if pathlib.Path(zipinfo.filename).suffix == ".wav":
                zipf.extract(zipinfo, dirname / "wav")
            else:
                zipf.extract(zipinfo, dirname)
    wavfiles = list((dirname / "wav").rglob("*.wav"))
    return wavfiles

from pydub import AudioSegment, effects

def get_mp3_path(wavfile):
    mp3dir = list(wavfile.parents)[1] / "mp3"
    mp3dir.mkdir(parents=True, exist_ok=True)
    mp3file = mp3dir / wavfile.with_suffix(".mp3").name
    return mp3file

def convert_wav_to_mp3(wavfile):
    mp3file = get_mp3_path(wavfile)
    try:
        wavaudio = AudioSegment.from_wav(wavfile)
        wavaudio = wavaudio.set_frame_rate(16000).set_channels(1)
        wavaudio = effects.normalize(wavaudio)
        wavaudio.export(mp3file, format="mp3", bitrate="16k")
        wavfile.unlink()
        return mp3file
    except Exception as e:
        print(f"Unable to convert {wavfile} to mp3")
        wavfile.unlink()

from concurrent.futures import as_completed, ThreadPoolExecutor
from multiprocessing import cpu_count, Pool
from tqdm import tqdm

def main():
    urls = list(map(lambda x: x.strip(), open("../data/file_list.txt").readlines()))
    print("Downloading archives")
    zipfiles = []
    with ThreadPoolExecutor(cpu_count()*4) as executor:
        results = executor.map(download_archive, urls)
        for url in tqdm(results):
            zipfiles.append(url)

    zipfiles = list(pathlib.Path("datasets").rglob("*.zip"))

    mp3files = []
    with ThreadPoolExecutor(max_workers=len(zipfiles)) as executor:
        print("Extracting archives")
        extracted = [executor.submit(extract_zipfile, file) for file in zipfiles]
        for wavfiles in as_completed(extracted):
            wavfiles = wavfiles.result()
            with Pool(cpu_count() - 1) as pool:
                results = pool.imap_unordered(convert_wav_to_mp3, wavfiles)
                print("Converting wav to mp3")
                for mp3file in tqdm(results, total=len(wavfiles)):
                    mp3files.append(mp3file)
    return mp3files

from datasets import load_dataset, Audio, DatasetDict, concatenate_datasets

def load_data_splits(is_streaming=True, stopping_strategy="all_exhausted"):
    data_dict = {}
    data_dict["openslr_dataset_train"] = load_dataset(
        "openslr", "SLR79", split="train", use_auth_token=True)
    data_dict["ucla_dataset_train"] = load_dataset(
        "audiofolder", data_dir="../data/kannada/", drop_labels=True)["train"]
    data_dict["fleurs_dataset_train"] = load_dataset(
        "google/fleurs", "kn_in", split="train",
        use_auth_token=True).rename_column("transcription", "sentence")
    data_dict["fleurs_dataset_val"] = load_dataset(
        "google/fleurs", "kn_in", split="validation",
        use_auth_token=True).rename_column("transcription", "sentence")
    data_dict["fleurs_dataset_test"] = load_dataset(
        "google/fleurs", "kn_in", split="test",
        use_auth_token=True).rename_column("transcription", "sentence")
    
    for k in data_dict:
        data_dict[k] = (
            data_dict[k].remove_columns(
                [col for col in data_dict[k].column_names if col not in ["audio", "sentence"]]))
        data_dict[k] = data_dict[k].cast_column("audio", Audio(sampling_rate=16000))
    
    dataset_dict = DatasetDict()
    train_datasets = []
    test_datasets = []
    for k in data_dict:
        if k.endswith("train") or k.endswith("val"):
            train_datasets.append(data_dict[k])
        if k.endswith("test"):
            test_datasets.append(data_dict[k])

    dataset_dict["train"] = concatenate_datasets(train_datasets)
    dataset_dict["test"] = concatenate_datasets(test_datasets)
    return dataset_dict

def audio_from_array(array):
    file = BytesIO()
    wavf.write(file, 16000, array)
    audio_segment = AudioSegment.from_file(file)
    audio_segment.set_frame_rate(16000).set_channels(1)
    audio_segment = effects.normalize(audio_segment)
    return audio_segment

def filter_nans_and_short(example):
    sentence = example["sentence"]
    length = example["length"]
    if sentence is None:
        return False
    elif length < 3 or length > 30:
        return False
    else:
        return True

def export_audio_and_sentence(example, split):
    audio, sentence = example["audio"], example["sentence"]
    new_name = pathlib.Path(uuid(audio["path"])).with_suffix(".mp3").name
    new_path = pathlib.Path(f"../data/filtered_datasets/kannada/{split}/{new_name}")
    new_path.parent.mkdir(parents=True, exist_ok=True)
    length = len(audio["array"])/16000
    if new_path.is_file():
        return {"path": str(new_path),"sentence": sentence, "length": length, "split": split}
    else:
        if not split ==  "test":
            is_not_filtered = filter_nans_and_short({"length": length, "sentence": sentence})
        else:
            is_not_filtered = True
        if is_not_filtered:
            try:
                new_path.parent.mkdir(exist_ok=True, parents=True)
                audio_segment = audio_from_array(audio["array"])
                audio_segment.export(new_path, format="mp3", bitrate="16k")
                segment = AudioSegment.from_mp3(new_path)
            except:
                return None
            return {"path": str(new_path),"sentence": sentence, "length": length, "split": split}
        else:
            return None

if __name__ == "__main__":

    dataset_dict = load_data_splits(is_streaming=False)
    exports = []
    with Pool(cpu_count()-1) as pool:
        for k in dataset_dict:
            dataset = dataset_dict[k].shuffle(seed=np.random.randint(1000))
            exporter = partial(export_audio_and_sentence, split=k)
            results = pool.imap_unordered(exporter,dataset,chunksize=10)
            for result in tqdm(results, total=len(dataset)):
                if result:
                    exports.append(result)
    exports = pd.DataFrame(exports)
    exports.to_json("../data/filtered_datasets/kannada/metadata.jsonl", lines=True, orient="records")

# https://huggingface.co/datasets/parambharat/tamil_asr_corpus/blob/main/tamil_asr_corpus.py
class TamilASRCorpus(datasets.GeneratorBasedBuilder):
    """Tamil ASR Corpus contains transcribed speech corpus for training ASR systems for Tamil language."""

    VERSION = datasets.Version("1.1.0")
    def _info(self):
        features = datasets.Features(
            {
                "audio": datasets.Audio(sampling_rate=16_000),
                "path": datasets.Value("string"),
                "sentence": datasets.Value("string"),
                "length": datasets.Value("float")
            }
        )
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=features,
            supervised_keys=("sentence", "label"),
            homepage=_HOMEPAGE,
            license=_LICENSE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager):
        metadata_paths = dl_manager.download(_METADATA_URLS)
        train_archive = dl_manager.download(_URLS["train"])
        test_archive = dl_manager.download(_URLS["test"])
        local_extracted_train_archive = dl_manager.extract(train_archive) if not dl_manager.is_streaming else None
        local_extracted_test_archive = dl_manager.extract(test_archive) if not dl_manager.is_streaming else None
        test_archive = dl_manager.download(_URLS["test"])
        train_dir = "train"
        test_dir = "test"

        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={
                    "metadata_path": metadata_paths["train"],
                    "local_extracted_archive": local_extracted_train_archive,
                    "path_to_clips": train_dir,
                    "audio_files": dl_manager.iter_archive(train_archive),
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.TEST,
                gen_kwargs={
                    "metadata_path": metadata_paths["test"],
                    "local_extracted_archive": local_extracted_test_archive,
                    "path_to_clips": test_dir,
                    "audio_files": dl_manager.iter_archive(test_archive),
                },
            ),
            
        ]
        
    def _generate_examples(self, metadata_path, local_extracted_archive, path_to_clips, audio_files):
        """Yields examples as (key, example) tuples."""
        examples = {}
        with open(metadata_path, encoding="utf-8") as f:
            for key, row in enumerate(f):
                data = json.loads(row)
                examples[data["path"]] = data
        inside_clips_dir = False
        id_ = 0
        for path, f in audio_files:
            if path.startswith(path_to_clips):
                inside_clips_dir = True
                if path in examples:
                    result = examples[path]
                    path = os.path.join(local_extracted_archive, path) if local_extracted_archive else path
                    result["audio"] = {"path": path, "bytes": f.read()}
                    result["path"] = path
                    yield id_, result
                    id_ += 1
            elif inside_clips_dir:
                break

from datasets import load_dataset
dataset = load_dataset("parambharat/telugu_asr_corpus", split="train")

from datasets import load_dataset

dataset = load_dataset("parambharat/telugu_asr_corpus", split="train", streaming=True)
next(iter(dataset))
>>{'audio': {'path': 'train/9zKDqLgBykz9FSyL6bFPRQ.mp3',
  'array': array([ 1.2912806e-03,  2.4933945e-03,  3.6223403e-03, ...,
          5.2775860e-05, -1.2118297e-05,  8.1833590e-05], dtype=float32),
  'sampling_rate': 16000},
 'path': 'train/9zKDqLgBykz9FSyL6bFPRQ.mp3',
 'sentence': 'అక్కడ మీడియాతో మాట్లాడిన అనంతరం తిరిగి హైదరాబాద్ బయలుదేరుతారు',
 'length': 3.87}

def load_data_splits(is_streaming=True,):
    data_dict = load_dataset("parambharat/tamil_asr_corpus", streaming=is_streaming)
    return data_dict

augment_waveform = Compose([
    AddGaussianNoise(min_amplitude=0.005, max_amplitude=0.015, p=0.2),
    TimeStretch(min_rate=0.8, max_rate=1.25, p=0.2, leave_length_unchanged=False),
    PitchShift(min_semitones=-4, max_semitones=4, p=0.2)
    ,])

def augment_dataset(batch):
    audio = batch["audio"]["array"]
    # apply augmentation
    augmented_audio = augment_waveform(samples=audio, sample_rate=16000)
    batch["audio"]["array"] = augmented_audio
    return batch

# call augment dataset on the training set
dataset_dict["train"] = dataset_dict["train"].map(augment_dataset)

def fix_sentence(sentence):
    transcription = sentence
  
    if transcription.startswith('"') and transcription.endswith('"'):
        # we can remove trailing quotation marks as they do not affect the transcription
        transcription = transcription[1:-1]
  
    if transcription[-1] not in [".", "?", "!"]:
        # append a full-stop to sentences that do not end in punctuation
        transcription = transcription + "."
    transcription = transcription[:-1].translate(str.maketrans('', '', string.punctuation)) + transcription[-1]
    return transcription
    
def filter_empty_strings(sentence):
    if len(sentence) < 2:
        return False
    else: return True

for k in dataset_dict:
    dataset_dict[k] = dataset_dict[k].filter(filter_empty_strings, input_columns=["sentence"])

feature_extractor = WhisperFeatureExtractor.from_pretrained(
    "openai/whisper-tiny"
)
tokenizer = WhisperTokenizer.from_pretrained(
    "openai/whisper-tiny", 
     language="Telugu",
     task="transcribe",
     model_max_length=225
)

def prepare_dataset(examples):
    # compute log-Mel input features from input audio array 
    audio = examples["audio"]
    
    examples["input_features"] = feature_extractor(
        audio["array"], sampling_rate=16000).input_features[0]
    
    sentences = fix_sentence(examples["sentence"])
    
    # encode target text to label ids 
    examples["labels"] = tokenizer(sentences, max_length=225, truncation=True).input_ids
    return examples

for k in dataset_dict:
    dataset_dict[k] = dataset_dict[k].map(
        prepare_dataset,).with_format("torch")

features = feature_extractor(sample["audio"]["array"], sampling_rate=16000)["input_features"][0]
print(features)
print(features.shape)

[Output:]

[[ 0.03961742 -0.4417838  -0.423414   ... -0.6383085  -0.6383085
  -0.6383085 ]
 [ 0.04983616 -0.28957796 -0.087901   ... -0.6383085  -0.6383085
  -0.6383085 ]
 [ 0.13291407 -0.16156828 -0.17585051 ... -0.6383085  -0.6383085
  -0.6383085 ]
 ...
 [-0.6383085  -0.6383085  -0.6383085  ... -0.6383085  -0.6383085
  -0.6383085 ]
 [-0.6383085  -0.6383085  -0.6383085  ... -0.6383085  -0.6383085
  -0.6383085 ]
 [-0.6383085  -0.6383085  -0.6383085  ... -0.6383085  -0.6383085
  -0.6383085 ]]
(80, 3000)

input_str = sample["sentence"]
labels = tokenizer(input_str).input_ids
decoded_with_special = tokenizer.decode(labels, skip_special_tokens=False)
decoded_str = tokenizer.decode(labels, skip_special_tokens=True)

print(f"Input:                 {input_str}")
print(f"Decoded w/ special:    {decoded_with_special}")
print(f"Decoded w/out special: {decoded_str}")
print(f"Are equal:             {input_str == decoded_str}")

[Output:]
Input:                 అక్కడ మీడియాతో మాట్లాడిన అనంతరం తిరిగి హైదరాబాద్ బయలుదేరుతారు
Decoded w/ special:    <|startoftranscript|><|te|><|transcribe|><|notimestamps|>అక్కడ మీడియాతో మాట్లాడిన అనంతరం తిరిగి హైదరాబాద్ బయలుదేరుతారు<|endoftext|>
Decoded w/out special: అక్కడ మీడియాతో మాట్లాడిన అనంతరం తిరిగి హైదరాబాద్ బయలుదేరుతారు
Are equal:             True