File size: 7,371 Bytes
d8be796 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 |
#!/usr/bin/env python
# pip install transformers datasets torch soundfile jiwer
from datasets import load_dataset, Audio
from transformers import pipeline, WhisperProcessor
from torch.utils.data import DataLoader
import torch
from jiwer import wer as jiwer_wer
from jiwer import cer as jiwer_cer
import ipdb
import subprocess
import os
# 1. Load FLEURS Burmese test set, cast to 16 kHz audio
ws_test_net = load_dataset("wenet-e2e/wenetspeech", "TEST_NET", split="test", streaming=False)
ws_test_meeting = load_dataset("wenet-e2e/wenetspeech", "TEST_MEETING", split="test", streaming=False)
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
model_id = "pengyizhou/whisper-wenetspeech-S"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)
whisper_model = "openai/whisper-large-v3"
processor = WhisperProcessor.from_pretrained(whisper_model, language="chinese")
asr = pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=processor.tokenizer,
feature_extractor=processor.feature_extractor,
torch_dtype=torch_dtype,
chunk_length_s=30,
batch_size=64,
max_new_tokens=225,
device=device,
num_beams=1, # Use beam search for better quality
)
generate_kwargs = {
"condition_on_prev_tokens": False,
"compression_ratio_threshold": 1.35, # zlib compression ratio threshold (in token space)
"temperature": (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
"logprob_threshold": -1.0,
}
# 3. Batch‐wise transcription function
def transcribe_batch(batch):
# `batch["audio"]` is a list of {"array": np.ndarray, ...}
inputs = [ ex["array"] for ex in batch["audio"] ]
outputs = asr(inputs, generate_kwargs=generate_kwargs) # returns a list of dicts with "text"
# lower-case and strip to normalize for CER
preds = [ out["text"].lower().strip() for out in outputs ]
return {"prediction": preds}
# 4. Map over the dataset in chunks of, say, 32 examples at a time
result_net = ws_test_net.map(
transcribe_batch,
batched=True,
batch_size=64, # feed 32 audios → pipeline will sub-batch into 8s
remove_columns=ws_test_net.column_names
)
result_meet = ws_test_meeting.map(
transcribe_batch,
batched=True,
batch_size=64, # feed 32 audios → pipeline will sub-batch into 8s
remove_columns=ws_test_meeting.column_names
)
for idx, (ds, result) in enumerate([(ws_test_net, result_net), (ws_test_meeting, result_meet)]):
ids = [key for key in ds["segment_id"]]
refs = [t.lower().strip() for t in ds["text"]]
preds = [t for t in result["prediction"]]
score_cer = jiwer_cer(refs, preds)
if idx == 0:
print(f"Zero-shot CER on wenet test_net: {score_cer*100:.2f}%")
with open("./wenet_net_finetune_nolid.pred", "w") as pred_results:
for key, pred in zip(ids, preds):
pred_results.write("{} {}\n".format(key, pred))
with open("./wenet_net.ref", "w") as ref_results:
for key, ref in zip(ids, refs):
ref_results.write("{} {}\n".format(key, ref))
if idx == 1:
print(f"Zero-shot CER on wenet test_meeting: {score_cer*100:.2f}%")
with open("./wenet_meeting_finetune_nolid.pred", "w") as pred_results:
for key, pred in zip(ids, preds):
pred_results.write("{} {}\n".format(key, pred))
pred_results.write("{}\n".format(pred))
with open("./wenet_meeting.ref", "w") as ref_results:
for key, ref in zip(ids, refs):
ref_results.write("{} {}\n".format(key, ref))
# Check if compute-wer.py exists
compute_wer_script = "./compute-wer.py"
if not os.path.exists(compute_wer_script):
# Try to find it in parent directories or common locations
possible_locations = [
"./compute-wer.py",
]
for location in possible_locations:
if os.path.exists(location):
compute_wer_script = location
break
else:
print(f"Warning: compute-wer.py not found. Tried: {[compute_wer_script] + possible_locations}")
print("Skipping detailed WER analysis.")
compute_wer_script = None
if compute_wer_script:
try:
# Run compute-wer.py with character-level analysis
ref_file = "./wenet_net.ref"
hyp_file = "./wenet_net_finetune_nolid.pred"
wer_file = "./wenet_net_finetune_nolid.wer"
cmd = [
"python", compute_wer_script,
"--char=1", # Character-level analysis
"--v=1", # Verbose output
ref_file,
hyp_file
]
print(f"Running: {' '.join(cmd)} > {wer_file}")
# Run the command and redirect output to wer file
with open(wer_file, "w") as wer_output:
result = subprocess.run(
cmd,
stdout=wer_output,
stderr=subprocess.PIPE,
text=True,
check=True
)
print(f"CER analysis saved to {wer_file}")
# Optionally, print the first few lines of the WER file
if os.path.exists(wer_file):
print("\nFirst few lines of WER analysis:")
with open(wer_file, "r") as f:
lines = f.readlines()
for i, line in enumerate(lines[:10]): # Show first 10 lines
print(f" {line.rstrip()}")
if len(lines) > 10:
print(f" ... ({len(lines) - 10} more lines)")
ref_file = "./wenet_meeting.ref"
hyp_file = "./wenet_meeting_finetune_nolid.pred"
wer_file = "./wenet_meeting_finetune_nolid.wer"
cmd = [
"python", compute_wer_script,
"--char=1", # Character-level analysis
"--v=1", # Verbose output
ref_file,
hyp_file
]
print(f"Running: {' '.join(cmd)} > {wer_file}")
# Run the command and redirect output to wer file
with open(wer_file, "w") as wer_output:
result = subprocess.run(
cmd,
stdout=wer_output,
stderr=subprocess.PIPE,
text=True,
check=True
)
print(f"CER analysis saved to {wer_file}")
# Optionally, print the first few lines of the WER file
if os.path.exists(wer_file):
print("\nFirst few lines of WER analysis:")
with open(wer_file, "r") as f:
lines = f.readlines()
for i, line in enumerate(lines[:10]): # Show first 10 lines
print(f" {line.rstrip()}")
if len(lines) > 10:
print(f" ... ({len(lines) - 10} more lines)")
except subprocess.CalledProcessError as e:
print(f"Error running compute-wer.py: {e}")
if e.stderr:
print(f"Error details: {e.stderr}")
except Exception as e:
print(f"Unexpected error: {e}")
print("Inference and CER analysis completed!") |