internalhell/whisper_small_ru_model_trainer_3ep

pip install transformers evaluate soundfile
pip install -q jiwer tensorboard
pip install --upgrade datasets transformers

import re
import json
from datasets import load_dataset, DatasetDict, Audio
from transformers import WhisperForConditionalGeneration, WhisperFeatureExtractor, WhisperTokenizer, WhisperProcessor, Seq2SeqTrainingArguments, Seq2SeqTrainer
import os, numpy as np, torch, evaluate, jiwer
from huggingface_hub import login
from dataclasses import dataclass
from typing import Any, Dict, List, Union

login("***")


common_voice = DatasetDict()
common_voice["train"] = load_dataset("mozilla-foundation/common_voice_17_0", "ru", split="train")
common_voice["test"] = load_dataset("mozilla-foundation/common_voice_17_0", "ru", split="test")

common_voice = common_voice.remove_columns(["accent", "age", "client_id", "down_votes", "gender", "locale", "path", "segment", "up_votes"])
common_voice = common_voice.cast_column("audio", Audio(sampling_rate=16000))

feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")
tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-small", language="Russian", task="transcribe")
processor = WhisperProcessor.from_pretrained("openai/whisper-small", language="Russian", task="transcribe")
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
model.config.forced_decoder_ids = None
model.config.suppress_tokens = []
model.config.use_cache = False

def prepare_dataset(batch):
    audio = batch["audio"]

    batch["input_features"] = feature_extractor(
        audio["array"],
        sampling_rate=audio["sampling_rate"]
    ).input_features[0]

    batch["labels"] = tokenizer(batch["sentence"]).input_ids
    return batch

common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice.column_names["train"], num_proc=2 )

common_voice

wer_metric = evaluate.load("wer")
cer_metric = evaluate.load("cer")

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    label_ids[label_ids == -100] = tokenizer.pad_token_id

    pred_str  = tokenizer.batch_decode(pred_ids,  skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    pairs = [(ref.strip(), hyp.strip()) for ref, hyp in zip(label_str, pred_str)]
    pairs = [(ref, hyp) for ref, hyp in pairs if len(ref) > 0]

    label_str, pred_str = zip(*pairs)

    wer = 100 * wer_metric.compute(predictions=pred_str, references=label_str)
    cer = 100 * cer_metric.compute(predictions=pred_str, references=label_str)

    ser = 100 * (sum(p.strip() != r.strip() for p, r in zip(pred_str, label_str)) / len(pred_str))

    return {
        "wer":  wer,
        "cer":  cer,
        "ser":  ser
    }

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any
    decoder_start_token_id: int

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        input_features = [{"input_features": f["input_features"]} for f in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        label_features = [{"input_ids": f["labels"]} for f in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels
        return batch

data_collator = DataCollatorSpeechSeq2SeqWithPadding(
    processor=processor,
    decoder_start_token_id=model.config.decoder_start_token_id,
)

training_args = Seq2SeqTrainingArguments(
    output_dir="/content/drive/MyDrive/models/whisper_small_ru_model_trainer_3ep",
    logging_dir="/content/drive/MyDrive/models/whisper_small_ru_model_trainer_3ep",
    group_by_length=True,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=4,
    eval_strategy="steps",
    logging_strategy="steps",
    save_strategy="steps",
    num_train_epochs=3,
    generation_max_length=170,
    logging_steps=25,
    eval_steps=500,
    save_steps=500,
    fp16=True,
    optim="adamw_torch_fused",
    torch_compile=True,
    gradient_checkpointing=True,
    learning_rate=1e-5,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=False,
    predict_with_generate=True,
)

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=common_voice["train"],
    eval_dataset=common_voice["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=processor.feature_extractor,
)

trainer.train()

import os
from transformers import (WhisperProcessor, 
			WhisperForConditionalGeneration, 
			pipeline)
import torch
import torchaudio
import librosa
import numpy as np

MODEL_HUG = "internalhell/whisper_small_ru_model_trainer_3ep"

processor = None
model = None
pipe = None

def get_model_pipe():
	global model, processor, pipe
	if model is None or processor is None:
		processor = WhisperProcessor.from_pretrained(MODEL_HUG, language="russian")
		model = WhisperForConditionalGeneration.from_pretrained(MODEL_HUG)

		model.generation_config.forced_decoder_ids = None
		forced_decoder_ids = processor.get_decoder_prompt_ids(language="ru", task="transcribe")
		model.config.forced_decoder_ids = forced_decoder_ids

		pipe = pipeline(
			"automatic-speech-recognition",
			model=model,
			tokenizer=processor.tokenizer,
			feature_extractor=processor.feature_extractor,
			device=0 if torch.cuda.is_available() else -1,
		)    
		
	return model

def recognize_audio_pipe(audio_path):
    model = get_model_pipe()
    
    waveform, sr = torchaudio.load(audio_path)
    waveform = waveform.mean(dim=0, keepdim=True)  # моно

    if sr != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)
        waveform = resampler(waveform)
        sr = 16000

    waveform_np = waveform.squeeze(0).numpy()
    return pipe({"array": waveform_np, "sampling_rate": sr})["text"]

print(recognize_audio_pipe("test.wav")) # jast .wav only