wav2vec2-large-xls-r-300m-Urdu: Model Card

import json
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from pyctcdecode import build_ctcdecoder
from transformers import AutoModelForCTC, AutoProcessor

MODEL_ID = "kingabzpro/wav2vec2-large-xls-r-300m-Urdu"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
INFERENCE_DTYPE = torch.float16 if DEVICE.type == "cuda" else torch.float32

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForCTC.from_pretrained(MODEL_ID).eval().to(
    device=DEVICE, dtype=INFERENCE_DTYPE
)

kenlm_path = hf_hub_download(MODEL_ID, "language_model/5gram.bin")
unigrams_path = hf_hub_download(MODEL_ID, "language_model/unigrams.txt")
attrs_path = hf_hub_download(MODEL_ID, "language_model/attrs.json")

with open(unigrams_path, encoding="utf-8") as file:
    unigrams = [line.strip() for line in file if line.strip()]
with open(attrs_path, encoding="utf-8") as file:
    attrs = json.load(file)

# Keep only acoustic tokens. The matching ID list is then applied to logits.
vocab_items = sorted(processor.tokenizer.get_vocab().items(), key=lambda item: item[1])
blank_id = processor.tokenizer.pad_token_id
delimiter = processor.tokenizer.word_delimiter_token
decoder_pairs = [
    (token, token_id)
    for token, token_id in vocab_items
    if token_id == blank_id or token == delimiter or len(token) == 1
]
kept_token_ids = [token_id for _, token_id in decoder_pairs]
labels = [
    "" if token_id == blank_id else " " if token == delimiter else token
    for token, token_id in decoder_pairs
]
decoder = build_ctcdecoder(
    labels,
    kenlm_model_path=kenlm_path,
    unigrams=unigrams,
    alpha=attrs.get("alpha", 0.5),
    beta=attrs.get("beta", 1.0),
)

waveform, sample_rate = torchaudio.load("audio.wav")
waveform = waveform.mean(dim=0)
if sample_rate != 16_000:
    waveform = torchaudio.functional.resample(waveform, sample_rate, 16_000)

inputs = processor(
    waveform.numpy(), sampling_rate=16_000, return_tensors="pt"
).input_values.to(device=DEVICE, dtype=INFERENCE_DTYPE)
with torch.inference_mode():
    logits = model(inputs).logits[0].float().cpu().numpy()

transcription = decoder.decode(logits[:, kept_token_ids])
print(transcription)

from datasets import Audio, load_dataset

stream = load_dataset(
    "fixie-ai/common_voice_17_0", "ur", split="test", streaming=True
).cast_column("audio", Audio(sampling_rate=16_000))

example = next(iter(stream))
audio = example["audio"]
samples = audio.get_all_samples().data if hasattr(audio, "get_all_samples") else audio["array"]
waveform = samples.detach().cpu().numpy() if torch.is_tensor(samples) else samples
if waveform.ndim == 2:
    waveform = waveform.mean(axis=0 if waveform.shape[0] <= waveform.shape[-1] else 1)

inputs = processor(waveform, sampling_rate=16_000, return_tensors="pt")
with torch.inference_mode():
    logits = model(inputs.input_values.to(device=DEVICE, dtype=INFERENCE_DTYPE)).logits[0]

prediction = decoder.decode(logits.float().cpu().numpy()[:, kept_token_ids])
print("Reference: ", example["sentence"])
print("Prediction:", prediction)