import transformers
import datasets
import torch
SAMPLE_RATE = 16000
LANGUAGE_TOKENS_WHISPER = {
"eng": 50259, "fra": 50265, "swa": 50318, "sna": 50324, "yor": 50325, "som": 50326,
"afr": 50327, "amh": 50334, "mlg": 50349, "lin": 50353, "hau": 50354,
"ach": 50357, "aka": 50356, "bam": 50355, "bem": 50352, "ber": 50351,
"cgg": 50350, "dag": 50348, "dga": 50347, "ewe": 50346, "ful": 50345,
"ibo": 50344, "kab": 50343, "kau": 50342, "kik": 50341, "kin": 50340,
"kln": 50339, "koo": 50338, "kpo": 50337, "led": 50336, "lgg": 50335,
"lth": 50333, "lug": 50332, "luo": 50331, "luy": 50330, "myx": 50329,
"nbl": 50328, "nya": 50323, "nyn": 50322, "orm": 50321, "pcm": 50320,
"ruc": 50319, "rwm": 50317, "sot": 50316, "teo": 50315, "tsn": 50314,
"ttj": 50313, "wol": 50312, "xho": 50311, "xog": 50310, "zul": 50309
}
LANGUAGE_NAMES = {
'Acholi': 'ach', 'Afrikaans': 'afr', 'Akan': 'aka', 'Amharic': 'amh', 'Ateso': 'teo',
'Bambara': 'bam', 'Bemba': 'bem', 'Berber': 'ber', 'Chichewa': 'nya', 'Dagaare': 'dga',
'Dagbani': 'dag', 'English': 'eng', 'Ewe': 'ewe', 'French': 'fra', 'Fulani': 'ful',
'Hausa': 'hau', 'Igbo': 'ibo', 'Ikposo': 'kpo', 'Kabyle': 'kab', 'Kalenjin': 'kln',
'Kanuri': 'kau', 'Kikuyu': 'kik', 'Kinyarwanda': 'kin', 'Kwamba': 'rwm', 'Lendu': 'led',
'Lingala': 'lin', 'Luganda': 'lug', 'Lugbara': 'lgg', 'Luhya': 'luy', 'Lumasaba': 'myx',
'Luo': 'luo', 'Lusoga': 'xog', 'Malagasy': 'mlg', 'Ndebele': 'nbl', 'Nigerian Pidgin': 'pcm',
'Oromo': 'orm', 'Rukiga': 'cgg', 'Rukonjo': 'koo', 'Runyankole': 'nyn', 'Ruruuli': 'ruc',
'Rutooro': 'ttj', 'Shona': 'sna', 'Somali': 'som', 'Sotho': 'sot', 'Swahili': 'swa',
'Thur': 'lth', 'Tswana': 'tsn', 'Wolof': 'wol', 'Xhosa': 'xho', 'Yoruba': 'yor', 'Zulu': 'zul'
}
model_id = "Sunbird/asr-whisper-51-african-languages"
model = transformers.WhisperForConditionalGeneration.from_pretrained(model_id)
processor = transformers.WhisperProcessor.from_pretrained(model_id)
def transcribe_by_whisper(audio_array, language):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input_features = processor(
audio_array, sampling_rate=SAMPLE_RATE, do_normalize=True, return_tensors="pt"
).input_features.to(device)
lang_tok = LANGUAGE_TOKENS_WHISPER[LANGUAGE_NAMES[language]]
transcribe_tok = processor.tokenizer.convert_tokens_to_ids("<|transcribe|>")
notimestamps_tok = processor.tokenizer.convert_tokens_to_ids("<|notimestamps|>")
forced_decoder_ids = [
(1, lang_tok),
(2, transcribe_tok),
(3, notimestamps_tok),
]
predicted_ids = model.to(device).generate(
input_features,
forced_decoder_ids=forced_decoder_ids,
num_beams=1,
do_sample=False,
)
transcription = processor.decode(
predicted_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
print(transcription[0])
import huggingface_hub
huggingface_hub.login()
ds = datasets.load_dataset('Sunbird/salt', 'multispeaker-lug', split='test')
audio_column = 'audio'
ds = ds.cast_column(audio_column, datasets.Audio(sampling_rate=SAMPLE_RATE))
audio_array = ds[0][audio_column]['array']
lang = 'Luganda'
transcribe_by_whisper(audio_array, lang)