|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
Fine-tuning the library models for sequence to sequence speech recognition |
|
with 🤗 Datasets' streaming mode. |
|
""" |
|
|
|
|
|
|
|
import json |
|
import logging |
|
import os |
|
import subprocess |
|
import sys |
|
from dataclasses import dataclass, field |
|
from typing import Any, Dict, List, Optional, Union |
|
|
|
import datasets |
|
import torch |
|
import wandb |
|
from datasets import DatasetDict, IterableDatasetDict, interleave_datasets, load_dataset |
|
from torch.utils.data import IterableDataset |
|
|
|
import evaluate |
|
import transformers |
|
from transformers import ( |
|
AutoConfig, |
|
AutoFeatureExtractor, |
|
AutoModelForSpeechSeq2Seq, |
|
AutoProcessor, |
|
AutoTokenizer, |
|
HfArgumentParser, |
|
Seq2SeqTrainer, |
|
Seq2SeqTrainingArguments, |
|
TrainerCallback, |
|
set_seed, |
|
) |
|
from transformers.models.whisper.english_normalizer import BasicTextNormalizer |
|
from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE |
|
from transformers.trainer_pt_utils import IterableDatasetShard |
|
from transformers.trainer_utils import get_last_checkpoint, is_main_process |
|
from transformers.utils import check_min_version, send_example_telemetry |
|
from transformers.utils.versions import require_version |
|
|
|
|
|
check_min_version("4.25.0.dev0") |
|
|
|
require_version("datasets>=1.18.2", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
SENDING_NOTIFICATION = "*** Sending notification to email ***" |
|
RECIPIENT_ADDRESS = "marinone94@gmail.com" |
|
|
|
wandb_token = os.environ.get("WANDB_TOKEN", "None") |
|
hf_token = os.environ.get("HF_TOKEN", None) |
|
if (hf_token is None or wandb_token == "None") and os.path.exists("./creds.txt"): |
|
with open("./creds.txt", "r") as f: |
|
lines = f.readlines() |
|
for line in lines: |
|
key, value = line.split("=") |
|
if key == "HF_TOKEN": |
|
hf_token = value.strip() |
|
if key == "WANDB_TOKEN": |
|
wandb_token = value.strip() |
|
if key == "EMAIL_ADDRESS": |
|
os.environ["EMAIL_ADDRESS"] = value.strip() |
|
if key == "EMAIL_PASSWORD": |
|
os.environ["EMAIL_PASSWORD"] = value.strip() |
|
|
|
if hf_token is not None: |
|
try: |
|
os.makedirs("/root/.huggingface", exist_ok=True) |
|
with open("/root/.huggingface/token", "w") as f: |
|
f.write(hf_token) |
|
logger.info("Huggingface API key set") |
|
except (PermissionError, OSError): |
|
logger.warning("Huggingface API key not set, relying on ~/.huggingface/token") |
|
else: |
|
logger.warning("Huggingface API key not set, relying on ~/.huggingface/token") |
|
|
|
wandb.login(key=wandb_token, relogin=True, timeout=5) |
|
wandb.init(project="whisper", entity="pn-aa") |
|
|
|
logger.info("Wandb API key set, logging to wandb") |
|
|
|
|
|
@dataclass |
|
class ModelArguments: |
|
""" |
|
Arguments pertaining to which model/config/tokenizer we are going to fine-tune from. |
|
""" |
|
|
|
model_name_or_path: str = field( |
|
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"} |
|
) |
|
config_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"} |
|
) |
|
tokenizer_name: Optional[str] = field( |
|
default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"} |
|
) |
|
feature_extractor_name: Optional[str] = field( |
|
default=None, metadata={"help": "feature extractor name or path if not the same as model_name"} |
|
) |
|
cache_dir: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"}, |
|
) |
|
use_fast_tokenizer: bool = field( |
|
default=True, |
|
metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}, |
|
) |
|
model_revision: str = field( |
|
default="main", |
|
metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}, |
|
) |
|
use_auth_token: bool = field( |
|
default=False, |
|
metadata={ |
|
"help": ( |
|
"Will use the token generated when running `huggingface-cli login` (necessary to use this script " |
|
"with private models)." |
|
) |
|
}, |
|
) |
|
freeze_feature_encoder: bool = field( |
|
default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."} |
|
) |
|
freeze_encoder: bool = field( |
|
default=False, metadata={"help": "Whether to freeze the entire encoder of the seq2seq model."} |
|
) |
|
forced_decoder_ids: List[List[int]] = field( |
|
default=None, |
|
metadata={ |
|
"help": ( |
|
"A list of pairs of integers which indicates a mapping from generation indices to token indices " |
|
"that will be forced before sampling. For example, [[0, 123]] means the first generated token " |
|
"will always be a token of index 123." |
|
) |
|
}, |
|
) |
|
suppress_tokens: List[int] = field( |
|
default=None, metadata={"help": "A list of tokens that will be suppressed at generation."} |
|
) |
|
model_index_name: str = field(default=None, metadata={"help": "Pretty name for the model card."}) |
|
|
|
|
|
@dataclass |
|
class DataTrainingArguments: |
|
""" |
|
Arguments pertaining to what data we are going to input our model for training and eval. |
|
""" |
|
|
|
dataset_train_name: str = field( |
|
default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} |
|
) |
|
dataset_train_config_name: Optional[str] = field( |
|
default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} |
|
) |
|
dataset_eval_name: str = field( |
|
default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} |
|
) |
|
dataset_eval_config_name: Optional[str] = field( |
|
default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} |
|
) |
|
text_column: Optional[str] = field( |
|
default=None, |
|
metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."}, |
|
) |
|
max_train_samples: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": ( |
|
"For debugging purposes or quicker training, truncate the number of training examples to this " |
|
"value if set." |
|
) |
|
}, |
|
) |
|
max_eval_samples: Optional[int] = field( |
|
default=None, |
|
metadata={ |
|
"help": ( |
|
"For debugging purposes or quicker training, truncate the number of evaluation examples to this " |
|
"value if set." |
|
) |
|
}, |
|
) |
|
audio_column_name: str = field( |
|
default="audio", |
|
metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"}, |
|
) |
|
text_column_name: str = field( |
|
default="text", |
|
metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"}, |
|
) |
|
max_duration_in_seconds: float = field( |
|
default=20.0, |
|
metadata={ |
|
"help": ( |
|
"Truncate audio files that are longer than `max_duration_in_seconds` seconds to" |
|
" 'max_duration_in_seconds`" |
|
) |
|
}, |
|
) |
|
min_duration_in_seconds: float = field( |
|
default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"} |
|
) |
|
train_split_name: str = field( |
|
default="train", |
|
metadata={ |
|
"help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'" |
|
}, |
|
) |
|
eval_split_name: str = field( |
|
default="test", |
|
metadata={ |
|
"help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'" |
|
}, |
|
) |
|
do_lower_case: bool = field( |
|
default=False, |
|
metadata={"help": "Whether the target text should be lower cased."}, |
|
) |
|
do_remove_punctuation: bool = field( |
|
default=False, |
|
metadata={"help": "Whether the target text should be striped of punctuation."}, |
|
) |
|
do_normalize_eval: bool = field( |
|
default=True, |
|
metadata={"help": "Whether to normalise the references and predictions in the eval WER calculation."}, |
|
) |
|
language_train: str = field( |
|
default=None, |
|
metadata={ |
|
"help": ( |
|
"Language for multilingual fine-tuning. This argument should be set for multilingual fine-tuning " |
|
"only. For English speech recognition, it should be set to `None`." |
|
) |
|
}, |
|
) |
|
language_eval: str = field( |
|
default=None, |
|
metadata={ |
|
"help": ( |
|
"Language for multilingual fine-tuning. This argument should be set for multilingual fine-tuning " |
|
"only. For English speech recognition, it should be set to `None`." |
|
) |
|
}, |
|
) |
|
task: str = field( |
|
default="transcribe", |
|
metadata={"help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."}, |
|
) |
|
shuffle_buffer_size: Optional[int] = field( |
|
default=500, |
|
metadata={ |
|
"help": ( |
|
"The number of streamed examples to download before shuffling them. The large the buffer, " |
|
"the closer it is to real offline shuffling." |
|
) |
|
}, |
|
) |
|
streaming: bool = field( |
|
default=True, |
|
metadata={"help": "Whether to use streaming mode to load and pre-process the data."}, |
|
) |
|
|
|
|
|
@dataclass |
|
class DataCollatorSpeechSeq2SeqWithPadding: |
|
""" |
|
Data collator that will dynamically pad the inputs received. |
|
Args: |
|
processor ([`WhisperProcessor`]) |
|
The processor used for processing the data. |
|
decoder_start_token_id (`int`) |
|
The begin-of-sentence of the decoder. |
|
""" |
|
|
|
processor: Any |
|
decoder_start_token_id: int |
|
task_id: int |
|
|
|
language_id: int = -100 |
|
|
|
def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]: |
|
|
|
|
|
model_input_name = self.processor.model_input_names[0] |
|
input_features = [{model_input_name: feature[model_input_name]} for feature in features] |
|
label_features = [{"input_ids": feature["labels"]} for feature in features] |
|
|
|
|
|
batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt") |
|
|
|
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:] |
|
|
|
|
|
labels = torch.cat((torch.full_like(labels[:, 0], self.task_id), labels), dim=1) |
|
labels = torch.cat((torch.full_like(labels[:, 0], self.language_id), labels), dim=1) |
|
labels = torch.cat((torch.full_like(labels[:, 0], self.decoder_start_token_id), labels), dim=1) |
|
|
|
batch["labels"] = labels |
|
|
|
return batch |
|
|
|
|
|
def notify_me(recipient, message=None): |
|
""" |
|
Send an email to the specified address with the specified message |
|
""" |
|
sender = os.environ.get("EMAIL_ADDRESS", None) |
|
password = os.environ.get("EMAIL_PASSWORD", None) |
|
if sender is None: |
|
logging.warning("No email address specified, not sending notification") |
|
if password is None: |
|
logging.warning("No email password specified, not sending notification") |
|
if message is None: |
|
message = "Training is finished!" |
|
|
|
if sender is not None: |
|
import smtplib |
|
from email.mime.text import MIMEText |
|
|
|
msg = MIMEText(message) |
|
msg["Subject"] = "Training updates..." |
|
msg["From"] = "marinone.auto@gmail.com" |
|
msg["To"] = recipient |
|
|
|
|
|
smtp_obj = smtplib.SMTP("smtp.gmail.com", 587) |
|
smtp_obj.starttls() |
|
smtp_obj.login(sender, password) |
|
smtp_obj.sendmail(sender, recipient, msg.as_string()) |
|
smtp_obj.quit() |
|
|
|
|
|
def rename_col_and_resample(dataset, dataset_name, text_column_names, text_col_name_ref, audio_column_name, sampling_rate): |
|
raw_datasets_features = list(dataset.features.keys()) |
|
logger.info(f"Dataset {dataset_name} - Features: {raw_datasets_features}") |
|
|
|
if text_col_name_ref not in raw_datasets_features: |
|
if len(text_column_names) == 1: |
|
raise ValueError("None of the text column names provided found in dataset." |
|
f"Text columns: {text_column_names}" |
|
f"Dataset columns: {raw_datasets_features}") |
|
flag = False |
|
for text_column_name in text_column_names: |
|
if text_column_name in raw_datasets_features: |
|
logger.info(f"Renaming text column {text_column_name} to {text_col_name_ref}") |
|
dataset = dataset.rename_column(text_column_name, text_col_name_ref) |
|
flag = True |
|
break |
|
if flag is False: |
|
raise ValueError("None of the text column names provided found in dataset." |
|
f"Text columns: {text_column_names}" |
|
f"Dataset columns: {raw_datasets_features}") |
|
if audio_column_name is not None and sampling_rate is not None: |
|
ds_sr = int(dataset.features[audio_column_name].sampling_rate) |
|
if ds_sr != sampling_rate: |
|
dataset = dataset.cast_column( |
|
audio_column_name, datasets.features.Audio(sampling_rate=sampling_rate) |
|
) |
|
|
|
raw_datasets_features = list(dataset.features.keys()) |
|
raw_datasets_features.remove(audio_column_name) |
|
raw_datasets_features.remove(text_col_name_ref) |
|
|
|
dataset = dataset.remove_columns(column_names=raw_datasets_features) |
|
return dataset |
|
|
|
|
|
def load_maybe_streaming_dataset( |
|
dataset_names, |
|
dataset_config_names, |
|
split="train", |
|
streaming=True, |
|
audio_column_name=None, |
|
sampling_rate=None, |
|
**kwargs |
|
): |
|
""" |
|
Utility function to load a dataset in streaming mode. For datasets with multiple splits, |
|
each split is loaded individually and then splits combined by taking alternating examples from |
|
each (interleaving). |
|
""" |
|
text_column_names = None |
|
if "text_column_name" in kwargs: |
|
text_column_names = kwargs.pop("text_column_name").split(",") |
|
text_col_name_ref = text_column_names[0] |
|
|
|
if "," in dataset_names or "+" in split: |
|
|
|
dataset_splits = [] |
|
for dataset_name, dataset_config_name, split_names in zip( |
|
dataset_names.split(","), dataset_config_names.split(","), split.split(",") |
|
): |
|
for split_name in split_names.split("+"): |
|
if dataset_config_name: |
|
dataset = load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs) |
|
else: |
|
dataset = load_dataset(dataset_name, split=split_name, streaming=streaming, **kwargs) |
|
|
|
dataset = rename_col_and_resample( |
|
dataset, |
|
dataset_name, |
|
text_column_names, |
|
text_col_name_ref, |
|
audio_column_name, |
|
sampling_rate |
|
) |
|
|
|
dataset_splits.append(dataset) |
|
|
|
|
|
interleaved_dataset = interleave_datasets(dataset_splits) |
|
return interleaved_dataset |
|
else: |
|
|
|
|
|
dataset = load_dataset(dataset_names, dataset_config_names, split=split, streaming=streaming, **kwargs) |
|
dataset = rename_col_and_resample( |
|
dataset, |
|
dataset_names, |
|
text_column_names, |
|
text_col_name_ref, |
|
audio_column_name, |
|
sampling_rate |
|
) |
|
return dataset |
|
|
|
|
|
def print_data_samples(dataset, tokenizer, max_samples=5): |
|
shown_samples = 0 |
|
for batch in dataset: |
|
print("Target: ", tokenizer.decode(batch["labels"])) |
|
shown_samples += len(batch) |
|
if shown_samples >= max_samples: |
|
break |
|
|
|
|
|
def main(): |
|
|
|
|
|
|
|
|
|
logger.info("*** Parse args ***") |
|
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) |
|
|
|
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): |
|
|
|
|
|
model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) |
|
else: |
|
model_args, data_args, training_args = parser.parse_args_into_dataclasses() |
|
|
|
|
|
|
|
send_example_telemetry("run_speech_recognition_seq2seq_streaming", model_args, data_args) |
|
|
|
|
|
logger.info("*** Setup logging ***") |
|
logging.basicConfig( |
|
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
|
datefmt="%m/%d/%Y %H:%M:%S", |
|
handlers=[logging.StreamHandler(sys.stdout)], |
|
) |
|
log_level = training_args.get_process_log_level() |
|
logger.setLevel(log_level) |
|
datasets.utils.logging.set_verbosity(log_level) |
|
transformers.utils.logging.set_verbosity(log_level) |
|
transformers.utils.logging.enable_default_handler() |
|
transformers.utils.logging.enable_explicit_format() |
|
|
|
logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN) |
|
|
|
|
|
logger.warning( |
|
f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" |
|
f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" |
|
) |
|
logger.info(f"Training/evaluation parameters {training_args}") |
|
|
|
|
|
if is_main_process(training_args.local_rank): |
|
transformers.utils.logging.set_verbosity_info() |
|
logger.info("Training/evaluation parameters %s", training_args) |
|
|
|
|
|
last_checkpoint = None |
|
if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: |
|
last_checkpoint = get_last_checkpoint(training_args.output_dir) |
|
if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0: |
|
raise ValueError( |
|
f"Output directory ({training_args.output_dir}) already exists and is not empty. " |
|
"Use --overwrite_output_dir to overcome." |
|
) |
|
elif last_checkpoint is not None and training_args.resume_from_checkpoint is None: |
|
logger.info( |
|
f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " |
|
"the `--output_dir` or add `--overwrite_output_dir` to train from scratch." |
|
) |
|
|
|
|
|
set_seed(training_args.seed) |
|
|
|
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained( |
|
model_args.feature_extractor_name if model_args.feature_extractor_name else model_args.model_name_or_path, |
|
cache_dir=model_args.cache_dir, |
|
revision=model_args.model_revision, |
|
use_auth_token=hf_token if model_args.use_auth_token else None, |
|
) |
|
|
|
|
|
logger.info("*** Load dataset ***") |
|
raw_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict() |
|
|
|
if len(data_args.language_eval.split(",")) > 1: |
|
raise ValueError("Implementation does not support multiple language evaluation.") |
|
|
|
if training_args.do_train: |
|
raw_datasets["train"] = load_maybe_streaming_dataset( |
|
data_args.dataset_train_name, |
|
data_args.dataset_train_config_name, |
|
split=data_args.train_split_name, |
|
use_auth_token=hf_token if model_args.use_auth_token else None, |
|
streaming=data_args.streaming, |
|
text_column_name=data_args.text_column_name, |
|
audio_column_name=data_args.audio_column_name, |
|
sampling_rate=int(feature_extractor.sampling_rate), |
|
|
|
) |
|
|
|
if training_args.do_eval: |
|
raw_datasets["eval"] = load_maybe_streaming_dataset( |
|
data_args.dataset_eval_name, |
|
data_args.dataset_eval_config_name, |
|
split=data_args.eval_split_name, |
|
use_auth_token=hf_token if model_args.use_auth_token else None, |
|
streaming=data_args.streaming, |
|
text_column_name=data_args.text_column_name, |
|
audio_column_name=data_args.audio_column_name, |
|
sampling_rate=int(feature_extractor.sampling_rate), |
|
|
|
) |
|
|
|
raw_datasets_features = list(next(iter(raw_datasets.values())).features.keys()) |
|
|
|
if data_args.audio_column_name not in raw_datasets_features: |
|
raise ValueError( |
|
f"--audio_column_name '{data_args.audio_column_name}' not found in dataset. " |
|
"Make sure to set `--audio_column_name` to the correct audio column - one of " |
|
f"{', '.join(raw_datasets_features)}." |
|
) |
|
|
|
data_args.text_column_name = data_args.text_column_name.split(",")[0] |
|
if data_args.text_column_name not in raw_datasets_features: |
|
raise ValueError( |
|
f"--text_column_name {data_args.text_column_name} not found in dataset. " |
|
"Make sure to set `--text_column_name` to the correct text column - one of " |
|
f"{', '.join(raw_datasets_features)}." |
|
) |
|
|
|
|
|
logger.info("*** Load pretrained model, tokenizer, and feature extractor ***") |
|
|
|
|
|
config = AutoConfig.from_pretrained( |
|
model_args.config_name if model_args.config_name else model_args.model_name_or_path, |
|
cache_dir=model_args.cache_dir, |
|
revision=model_args.model_revision, |
|
use_auth_token=hf_token if model_args.use_auth_token else None |
|
) |
|
|
|
|
|
config.update({"forced_decoder_ids": model_args.forced_decoder_ids, "suppress_tokens": model_args.suppress_tokens}) |
|
|
|
if training_args.gradient_checkpointing: |
|
config.update({"use_cache": False}) |
|
|
|
tokenizer = AutoTokenizer.from_pretrained( |
|
model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path, |
|
cache_dir=model_args.cache_dir, |
|
use_fast=model_args.use_fast_tokenizer, |
|
revision=model_args.model_revision, |
|
use_auth_token=hf_token if model_args.use_auth_token else None, |
|
) |
|
model = AutoModelForSpeechSeq2Seq.from_pretrained( |
|
model_args.model_name_or_path, |
|
config=config, |
|
cache_dir=model_args.cache_dir, |
|
revision=model_args.model_revision, |
|
use_auth_token=hf_token if model_args.use_auth_token else None, |
|
) |
|
|
|
if model.config.decoder_start_token_id is None: |
|
raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined") |
|
|
|
if model_args.freeze_feature_encoder: |
|
model.freeze_feature_encoder() |
|
|
|
if model_args.freeze_encoder: |
|
model.freeze_encoder() |
|
|
|
if data_args.language_train is not None and len(data_args.language_train.split(",")) == 1: |
|
|
|
|
|
tokenizer.set_prefix_tokens(language=data_args.language_train, task=data_args.task) |
|
elif data_args.language_train is not None and len(data_args.language_train.split(",")) > 1: |
|
|
|
model.config.forced_decoder_ids = None |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
logger.info("*** Preprocess dataset ***") |
|
max_input_length = data_args.max_duration_in_seconds * feature_extractor.sampling_rate |
|
min_input_length = data_args.min_duration_in_seconds * feature_extractor.sampling_rate |
|
audio_column_name = data_args.audio_column_name |
|
text_column_name = data_args.text_column_name |
|
model_input_name = feature_extractor.model_input_names[0] |
|
do_lower_case = data_args.do_lower_case |
|
do_remove_punctuation = data_args.do_remove_punctuation |
|
normalizer = BasicTextNormalizer() |
|
|
|
if data_args.max_train_samples is not None: |
|
raw_datasets["train"] = ( |
|
raw_datasets["train"].take(data_args.max_train_samples) |
|
if data_args.streaming |
|
else raw_datasets["train"].select(range(data_args.max_train_samples)) |
|
) |
|
|
|
if data_args.max_eval_samples is not None: |
|
raw_datasets["eval"] = ( |
|
raw_datasets["eval"].take(data_args.max_eval_samples) |
|
if data_args.streaming |
|
else raw_datasets["eval"].select(range(data_args.max_eval_samples)) |
|
) |
|
|
|
def prepare_dataset(batch): |
|
|
|
sample = batch[audio_column_name] |
|
inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"]) |
|
|
|
batch[model_input_name] = inputs.get(model_input_name)[0] |
|
batch["input_length"] = len(sample["array"]) |
|
|
|
|
|
input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name] |
|
if do_remove_punctuation: |
|
input_str = normalizer(input_str).strip() |
|
batch["labels"] = tokenizer(input_str).input_ids |
|
return batch |
|
|
|
with training_args.main_process_first(desc="dataset map pre-processing"): |
|
|
|
vectorized_datasets = raw_datasets.map( |
|
prepare_dataset, |
|
remove_columns=raw_datasets_features, |
|
).with_format("torch") |
|
|
|
if training_args.do_train and data_args.streaming: |
|
|
|
vectorized_datasets["train"] = vectorized_datasets["train"].shuffle( |
|
buffer_size=data_args.shuffle_buffer_size, |
|
seed=training_args.seed, |
|
) |
|
|
|
|
|
|
|
def is_audio_in_length_range(length): |
|
return min_input_length < length < max_input_length |
|
|
|
if training_args.do_train: |
|
vectorized_datasets["train"] = vectorized_datasets["train"].filter( |
|
is_audio_in_length_range, |
|
input_columns=["input_length"], |
|
) |
|
|
|
|
|
logger.info("*** Load metric ***") |
|
metric = evaluate.load("wer") |
|
do_normalize_eval = data_args.do_normalize_eval |
|
|
|
def compute_metrics(pred): |
|
pred_ids = pred.predictions |
|
|
|
pred.label_ids[pred.label_ids == -100] = tokenizer.pad_token_id |
|
|
|
pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True) |
|
|
|
label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True) |
|
|
|
if do_normalize_eval: |
|
pred_str = [normalizer(pred) for pred in pred_str] |
|
label_str = [normalizer(label) for label in label_str] |
|
|
|
pred_str = [pred_str[i] for i in range(len(pred_str)) if len(label_str[i]) > 0] |
|
label_str = [label_str[i] for i in range(len(label_str)) if len(label_str[i]) > 0] |
|
|
|
wer = 100 * metric.compute(predictions=pred_str, references=label_str) |
|
|
|
return {"wer": wer} |
|
|
|
|
|
logger.info("*** Init processor ***") |
|
if is_main_process(training_args.local_rank): |
|
|
|
feature_extractor.save_pretrained(training_args.output_dir) |
|
tokenizer.save_pretrained(training_args.output_dir) |
|
config.save_pretrained(training_args.output_dir) |
|
|
|
processor = AutoProcessor.from_pretrained(training_args.output_dir) |
|
|
|
|
|
task_token = data_args.task |
|
if not task_token.startswith('<|'): |
|
task_token = f'<{task_token}>' |
|
task_id = tokenizer(task_token).input_ids[0] |
|
data_collator = DataCollatorSpeechSeq2SeqWithPadding( |
|
processor=processor, |
|
decoder_start_token_id=model.config.decoder_start_token_id, |
|
task_id=task_id |
|
) |
|
|
|
|
|
|
|
|
|
logger.info("*** Set shuffle callback ***") |
|
class ShuffleCallback(TrainerCallback): |
|
def on_epoch_begin(self, args, state, control, train_dataloader, **kwargs): |
|
if isinstance(train_dataloader.dataset, IterableDatasetShard): |
|
pass |
|
elif isinstance(train_dataloader.dataset, IterableDataset): |
|
train_dataloader.dataset.set_epoch(train_dataloader.dataset._epoch + 1) |
|
|
|
|
|
|
|
logger.info("*** Init trainer ***") |
|
trainer = Seq2SeqTrainer( |
|
model=model, |
|
args=training_args, |
|
train_dataset=vectorized_datasets["train"] if training_args.do_train else None, |
|
eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None, |
|
tokenizer=feature_extractor, |
|
data_collator=data_collator, |
|
compute_metrics=compute_metrics if training_args.predict_with_generate else None, |
|
callbacks=[ShuffleCallback()] if data_args.streaming else None, |
|
) |
|
logger.info("*** Trainer initialized ***") |
|
|
|
|
|
if training_args.do_train: |
|
logger.info("*** Train ***") |
|
print_data_samples(vectorized_datasets["train"], tokenizer) |
|
checkpoint = None |
|
if training_args.resume_from_checkpoint is not None: |
|
checkpoint = training_args.resume_from_checkpoint |
|
elif last_checkpoint is not None: |
|
checkpoint = last_checkpoint |
|
train_result = trainer.train(resume_from_checkpoint=checkpoint) |
|
logger.info("*** Training completed ***") |
|
logger.info("*** Saving model ***") |
|
|
|
|
|
|
|
orig_push_to_hub = trainer.args.push_to_hub |
|
trainer.args.push_to_hub = False |
|
trainer.save_model() |
|
trainer.args.push_to_hub = orig_push_to_hub |
|
logger.info("*** Model saved ***") |
|
metrics = train_result.metrics |
|
if data_args.max_train_samples: |
|
metrics["train_samples"] = data_args.max_train_samples |
|
logger.info("*** Logging metrics ***") |
|
trainer.log_metrics("train", metrics) |
|
logger.info("*** Metrics logged ***") |
|
logger.info("*** Saving metrics ***") |
|
trainer.save_metrics("train", metrics) |
|
logger.info("*** Metrics saved ***") |
|
logger.info("*** Saving state ***") |
|
trainer.save_state() |
|
logger.info("*** State saved ***") |
|
|
|
|
|
predictions = trainer.predict( |
|
test_dataset=vectorized_datasets["eval"].shuffle(seed=training_args.seed).take(5), |
|
metric_key_prefix="test", |
|
max_length=training_args.generation_max_length, |
|
num_beams=training_args.generation_num_beams, |
|
) |
|
logger.info("*** Test prediction done ***") |
|
preds = tokenizer.batch_decode(predictions.predictions) |
|
labels = tokenizer.batch_decode(predictions.label_ids) |
|
pred_labels = [f"Prediction: {pred}\nLabel: {label}\n" for pred, label in zip(preds, labels)] |
|
logger.info("Before setting language and task") |
|
logger.info(f"{pred_labels}") |
|
trainer.model.config.forced_decoder_ids = \ |
|
tokenizer.get_decoder_prompt_ids(language=data_args.language_eval, task=data_args.task, no_timestamps=True) |
|
preds = tokenizer.batch_decode(predictions.predictions) |
|
labels = tokenizer.batch_decode(predictions.label_ids) |
|
pred_labels = [f"Prediction: {pred}\nLabel: {label}\n" for pred, label in zip(preds, labels)] |
|
logger.info("After setting language and task") |
|
logger.info(f"{pred_labels}") |
|
|
|
|
|
results = {} |
|
if training_args.do_eval: |
|
logger.info("*** Evaluate ***") |
|
print_data_samples(vectorized_datasets["eval"], tokenizer) |
|
metrics = trainer.evaluate( |
|
metric_key_prefix="eval", |
|
max_length=training_args.generation_max_length, |
|
num_beams=training_args.generation_num_beams, |
|
) |
|
logger.info("*** Evaluation done ***") |
|
if data_args.max_eval_samples: |
|
metrics["eval_samples"] = data_args.max_eval_samples |
|
logger.info("*** Logging metrics ***") |
|
trainer.log_metrics("eval", metrics) |
|
logger.info("*** Metrics logged ***") |
|
logger.info("*** Saving metrics ***") |
|
trainer.save_metrics("eval", metrics) |
|
logger.info("*** Metrics saved ***") |
|
|
|
|
|
logger.info("*** Writing training stats ***") |
|
kwargs = { |
|
"finetuned_from": model_args.model_name_or_path, |
|
"tasks": "automatic-speech-recognition", |
|
"tags": "whisper-event", |
|
} |
|
if data_args.dataset_train_name is not None: |
|
dataset_names = list(data_args.dataset_train_name.split(",")) |
|
kwargs["dataset_tags"] = dataset_names |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if data_args.language_train is not None: |
|
languages = list(set(data_args.language_train.split(","))) |
|
kwargs["language"] = languages |
|
if model_args.model_index_name is not None: |
|
kwargs["model_name"] = model_args.model_index_name |
|
|
|
logger.info("*** Training stats written ***") |
|
logger.info(json.dumps(kwargs, indent=4)) |
|
|
|
|
|
logger.info("*** Training and eval complete ***") |
|
logger.info(SENDING_NOTIFICATION) |
|
with open(os.path.join(training_args.output_dir, "train_results.json"), "r") as f: |
|
train_results = json.load(f) |
|
with open(os.path.join(training_args.output_dir, "eval_results.json"), "r") as f: |
|
eval_results = json.load(f) |
|
notify_me(recipient=RECIPIENT_ADDRESS, |
|
message=f"Training complete! {train_results = } {eval_results = }") |
|
|
|
|
|
if training_args.push_to_hub: |
|
logger.info("*** Pushing to hub ***") |
|
trainer.push_to_hub(**kwargs) |
|
logger.info("*** Pushed to hub ***") |
|
logger.info(SENDING_NOTIFICATION) |
|
else: |
|
logger.info("*** Creating model card ***") |
|
trainer.create_model_card(**kwargs) |
|
logger.info("*** Model card created ***") |
|
logger.info(SENDING_NOTIFICATION) |
|
|
|
with open(os.path.join(training_args.output_dir, "README.md"), "r") as f: |
|
readme = f.read() |
|
notify_me(recipient=RECIPIENT_ADDRESS, |
|
message=f"Model pushed to hub! {readme = }") |
|
|
|
return results |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|