File size: 5,777 Bytes
e112632 |
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 |
# scripts/finetune_whisper_small.py
import argparse, os, sys, pandas as pd, numpy as np, soundfile as sf, librosa, torch
from datasets import Dataset
from transformers import (
WhisperProcessor, WhisperForConditionalGeneration,
Seq2SeqTrainingArguments, Seq2SeqTrainer
)
SAMPLE_RATE = 16000
DUR_SEC = 2.0
N_SAMPLES = int(SAMPLE_RATE * DUR_SEC)
def load_2s(path: str) -> np.ndarray:
try:
audio, sr = sf.read(path, dtype="float32", always_2d=False)
if audio.ndim > 1:
audio = np.mean(audio, axis=1)
except Exception:
# fallback if soundfile cannot read
audio, sr = librosa.load(path, sr=None, mono=True)
audio = audio.astype(np.float32)
if sr != SAMPLE_RATE:
audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE).astype(np.float32)
if len(audio) >= N_SAMPLES:
audio = audio[:N_SAMPLES]
else:
pad = np.zeros(N_SAMPLES, dtype=np.float32)
pad[: len(audio)] = audio
audio = pad
return np.ascontiguousarray(audio, dtype=np.float32)
def make_transform():
# set_transform can be called on single example or batch; handle both
def _tx(batch):
paths = batch["path"] if isinstance(batch["path"], list) else [batch["path"]]
texts = batch["text"] if isinstance(batch["text"], list) else [batch["text"]]
# we only attach raw 2s audio; features are made in data_collator
audios = [load_2s(p) for p in paths]
out = {
"audio_2s": audios if len(audios) > 1 else audios[0],
"text": texts if len(texts) > 1 else texts[0],
"path": paths if len(paths) > 1 else paths[0],
}
return out
return _tx
class DataCollatorWhisper2s:
def __init__(self, processor: WhisperProcessor):
self.processor = processor
self.tok = processor.tokenizer
self.feat = processor.feature_extractor
def __call__(self, features):
# features contain: {"audio_2s": np.array, "text": str, "path": str}
audios = [f["audio_2s"] for f in features]
texts = [f["text"] for f in features]
# Whisper input features (80x3000), computed on-the-fly
# processor.feature_extractor already pads to fixed shape
inputs = self.feat(audios, sampling_rate=SAMPLE_RATE, return_tensors="pt")
input_features = inputs.input_features # [B, 80, 3000]
# Tokenize targets; pad, then replace pad_token_id with -100
labels = self.tok(texts, return_tensors="pt", padding=True).input_ids
labels[labels == self.tok.pad_token_id] = -100
return {"input_features": input_features, "labels": labels}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--csv", required=True, help="CSV with columns: path,transcription OR path,text")
ap.add_argument("--out_dir", required=True)
ap.add_argument("--base_model", default="openai/whisper-small")
ap.add_argument("--lang", default="en")
ap.add_argument("--epochs", type=int, default=3)
ap.add_argument("--batch", type=int, default=8)
ap.add_argument("--num_workers", type=int, default=4)
args = ap.parse_args()
out_dir_abs = os.path.abspath(args.out_dir)
os.makedirs(out_dir_abs, exist_ok=True)
print(f"[INFO] Will save to: {out_dir_abs}")
head = pd.read_csv(args.csv, nrows=1)
df = pd.read_csv(args.csv)
if "transcription" in head.columns:
df = df.rename(columns={"transcription": "text"})
assert {"path","text"}.issubset(df.columns), "CSV must have columns path and transcription/text"
df = df.dropna(subset=["path","text"])
df = df[(df["path"].astype(str).str.len()>0) & (df["text"].astype(str).str.len()>0)]
print(f"[INFO] Rows: {len(df)}")
# Tiny, path-only dataset; we transform lazily
ds = Dataset.from_pandas(df[["path","text"]], preserve_index=False)
ds = ds.train_test_split(test_size=0.1, seed=42)
train_ds, eval_ds = ds["train"], ds["test"]
processor = WhisperProcessor.from_pretrained(args.base_model, language=args.lang, task="transcribe")
model = WhisperForConditionalGeneration.from_pretrained(args.base_model)
model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language=args.lang, task="transcribe")
model.config.suppress_tokens = [] # better for short clips
model.gradient_checkpointing_enable() # less VRAM
# lazy transforms (no big arrays stored)
tx = make_transform()
train_ds.set_transform(tx)
eval_ds.set_transform(tx)
collator = DataCollatorWhisper2s(processor)
train_args = Seq2SeqTrainingArguments(
output_dir=out_dir_abs,
per_device_train_batch_size=args.batch,
per_device_eval_batch_size=args.batch,
dataloader_num_workers=args.num_workers,
gradient_accumulation_steps=2,
learning_rate=1e-5,
num_train_epochs=args.epochs,
evaluation_strategy="epoch",
save_strategy="epoch",
logging_strategy="steps",
logging_steps=50,
predict_with_generate=False,
fp16=True,
remove_unused_columns=False, # keep our custom fields
report_to=[]
)
trainer = Seq2SeqTrainer(
model=model,
args=train_args,
train_dataset=train_ds,
eval_dataset=eval_ds,
tokenizer=processor.tokenizer,
data_collator=collator,
)
try:
trainer.train()
finally:
print("[INFO] Forcing save...")
trainer.save_model(out_dir_abs)
model.save_pretrained(out_dir_abs)
processor.save_pretrained(out_dir_abs)
open(os.path.join(out_dir_abs, "SAVE_OK"), "w").write("ok\n")
print(f"[DONE] Saved to: {out_dir_abs}")
if __name__ == "__main__":
main()
|