File size: 1,992 Bytes
f2a5b60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Greedy-CTC inference for anshulsc/w2v-bert-ethio-v5.
Pipeline matches the evaluation harness exactly:
1. 16 kHz mono audio -> SeamlessM4TFeatureExtractor (from facebook/w2v-bert-2.0)
2. Wav2Vec2BertForCTC forward pass -> argmax over logits
3. Mask the six [LID:*] token ids (311-316) to the CTC blank id
4. Wav2Vec2CTCTokenizer.decode (collapses CTC repeats, maps | to spaces)
Needs: torch transformers soundfile librosa
"""
import sys
import soundfile as sf
import torch
from transformers import SeamlessM4TFeatureExtractor, Wav2Vec2BertForCTC, Wav2Vec2CTCTokenizer
MODEL_ID = "anshulsc/w2v-bert-ethio-v5"
LID_TOKEN_IDS = [311, 312, 313, 314, 315, 316] # [LID:*] tokens must not survive decoding
def transcribe(audio_path: str) -> str:
speech, sr = sf.read(audio_path, dtype="float32")
if speech.ndim > 1: # downmix stereo to mono
speech = speech.mean(axis=1)
feature_extractor = SeamlessM4TFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
if sr != feature_extractor.sampling_rate:
import librosa
speech = librosa.resample(speech, orig_sr=sr, target_sr=feature_extractor.sampling_rate)
inputs = feature_extractor(
speech,
sampling_rate=feature_extractor.sampling_rate,
return_tensors="pt",
return_attention_mask=True,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Wav2Vec2BertForCTC.from_pretrained(MODEL_ID, torch_dtype=torch.float16).to(device).eval()
inputs = {k: v.to(device) for k, v in inputs.items()}
tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(MODEL_ID)
with torch.no_grad():
logits = model(**inputs).logits
pred_ids = torch.argmax(logits, dim=-1)
lid = torch.tensor(LID_TOKEN_IDS, device=pred_ids.device)
pred_ids = pred_ids.masked_fill(torch.isin(pred_ids, lid), tokenizer.pad_token_id)
return tokenizer.decode(pred_ids[0]).strip()
if __name__ == "__main__":
print(transcribe(sys.argv[1]))
|