| """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] |
|
|
|
|
| def transcribe(audio_path: str) -> str: |
| speech, sr = sf.read(audio_path, dtype="float32") |
| if speech.ndim > 1: |
| 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])) |
|
|
