In your current code, the Seq2SeqTrainingArguments class is responsible for saving checkpoints. The save_steps attribute dictates how often a checkpoint is saved. If a checkpoint is saved, you can resume training from that checkpoint at any time. Your current save_steps is set to 1000, meaning a checkpoint will be saved every 1000 steps.
If you want to save checkpoints more frequently, you can decrease the value of save_steps.
Here's the complete refactored version of your code:
import torch
from dataclasses import dataclass
from typing import Any, Dict, List, Union
# import the relevant libraries for logging in
from huggingface_hub import HfApi, HfFolder
from datasets import load_dataset, DatasetDict, Audio
from transformers import (
WhisperFeatureExtractor,
WhisperTokenizer,
WhisperProcessor,
WhisperForConditionalGeneration,
Seq2SeqTrainingArguments,
Seq2SeqTrainer,
)
import evaluate
def load_datasets():
common_voice = DatasetDict()
common_voice["train"] = load_dataset("mozilla-foundation/common_voice_11_0", "bn", split="train+validation", use_auth_token=False)
common_voice["test"] = load_dataset("mozilla-foundation/common_voice_11_0", "bn", split="test", use_auth_token=False)
common_voice = common_voice.remove_columns(
["accent",
"age",
"client_id",
"down_votes",
"gender",
"locale",
"path",
"segment",
"up_votes"]
)
return common_voice
def get_processor():
feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")
tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-small", language="bengali", task="transcribe")
processor = WhisperProcessor.from_pretrained("openai/whisper-small", language="bengali", task="transcribe")
return processor
def prepare_dataset(batch, feature_extractor, tokenizer):
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
def prepare_datasets(common_voice, processor):
common_voice = common_voice.cast_column("audio", Audio(sampling_rate=16000))
common_voice = common_voice.map(
lambda batch: prepare_dataset(batch, processor.feature_extractor, processor.tokenizer),
remove_columns=common_voice.column_names["train"],
num_proc=2,
)
return common_voice
def get_data_collator(processor):
@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
processor: Any
def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
input_features = [{"input_features": feature["input_features"]} for feature in features]
batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
label_features = [{"input_ids": feature["labels"]} for feature 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.processor.tokenizer.bos_token_id).all().cpu().item():
labels = labels[:, 1:]
batch["labels"] = labels
return batch
return DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
def get
_metric():
return evaluate.load("wer")
def get_model():
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
model.config.forced_decoder_ids = None
model.config.suppress_tokens = []
return model
def get_training_args():
return Seq2SeqTrainingArguments(
output_dir="output/whisper-small-bn",
per_device_train_batch_size=8,
gradient_accumulation_steps=2,
learning_rate=1e-5,
warmup_steps=500,
max_steps=4000,
gradient_checkpointing=True,
fp16=True,
evaluation_strategy="steps",
per_device_eval_batch_size=4,
predict_with_generate=True,
generation_max_length=225,
save_steps=500, # decreased this to save more frequently
eval_steps=1000,
logging_steps=25,
report_to=["tensorboard"],
load_best_model_at_end=True,
metric_for_best_model="wer",
greater_is_better=False,
push_to_hub=False,
)
def compute_metrics(pred, tokenizer, metric):
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)
wer = 100 * metric.compute(predictions=pred_str, references=label_str)
return {"wer": wer}
def get_trainer(training_args, model, common_voice, data_collator, processor, metric):
return Seq2SeqTrainer(
args=training_args,
model=model,
train_dataset=common_voice["train"],
eval_dataset=common_voice["test"],
data_collator=data_collator,
compute_metrics=lambda pred: compute_metrics(pred, processor.tokenizer, metric),
tokenizer=processor.feature_extractor,
)
def train_model(trainer):
print('Training is started.')
trainer.train()
print('Training is finished.')
def upload_to_hub(trainer):
kwargs = {
"dataset_tags": "mozilla-foundation/common_voice_11_0",
"dataset": "Common Voice 11.0",
"dataset_args": "config: lt, split: test",
"language": "lt",
"model_name": "Whisper Large LT - Vytautas Bielinskas",
"finetuned_from": "openai/whisper-large",
"tasks": "automatic-speech-recognition",
"tags": "hf-asr-leaderboard",
}
trainer.push_to_hub(**kwargs)
print('Trained model uploaded to the Hugging Face Hub')
def main():
common_voice = load_datasets()
processor = get_processor()
common_voice = prepare_datasets(common_voice, processor)
data_collator = get_data_collator(processor)
metric = get_metric()
model = get_model()
training_args = get_training_args()
trainer = get_trainer(training_args, model, common_voice, data_collator, processor, metric)
# Save processor object before starting training
processor.save_pretrained(training_args.output_dir)
train_model(trainer)
upload_to_hub(trainer)
if __name__ == "__main__":
main()This refactored version encapsulates each step into its own function, making the code easier to understand and modify. The main() function calls each of these in the correct order.
