Skip to content

GitLab

"include/ck/utility/get_id.hpp" did not exist on "7a89684f92cc39afbb13ad970c1a3282e60b9180"
Commit db566365 authored by Yoach Lacombe's avatar Yoach Lacombe
Browse files

clean artifacts

parent 64cfa64e
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 2609 deletions
audio_classification_scripts/run_dropout_sweep.yaml deleted 100644 → 0
View file @ 64cfa64e
command:
- python3
- ${program}
- --fp16
- --fp16_full_eval
- --do_train
- --do_eval
- --trust_remote_code
- --overwrite_output_dir
- --ignore_mismatched_sizes
- --gradient_checkpointing
- ${args}
method: random
metric:
goal: maximize
name: eval/accuracy
parameters:
model_name_or_path:
value: facebook/mms-lid-126
train_dataset_name:
value: stable-speech/concatenated-normalized-accent-dataset
train_dataset_config_name:
value: default
train_split_name:
value: train
train_label_column_name:
value: labels
eval_dataset_name:
value: stable-speech/concatenated-normalized-accent-dataset
eval_dataset_config_name:
value: default
eval_split_name:
value: test
eval_label_column_name:
value: labels
output_dir:
value: ./
remove_unused_columns:
value: false
learning_rate:
value: 1e-4
lr_scheduler_type:
value: constant_with_warmup
max_length_seconds:
value: 20
min_length_seconds:
value: 5
attention_mask:
value: true
warmup_steps:
value: 50
max_steps:
value: 1000
per_device_train_batch_size:
value: 32
per_device_eval_batch_size:
value: 32
preprocessing_num_workers:
value: 4
dataloader_num_workers:
value: 4
logging_strategy:
value: steps
logging_steps:
value: 10
evaluation_strategy:
value: steps
eval_steps:
value: 1000
save_strategy:
value: steps
save_steps:
value: 1000
freeze_base_model:
values:
- false
- true
push_to_hub:
value: false
filter_threshold:
value: 1
feat_proj_dropout:
values:
- 0.0
- 0.1
- 0.2
attention_dropout:
values:
- 0.0
- 0.1
- 0.2
activation_dropout:
values:
- 0.0
- 0.1
- 0.2
hidden_dropout:
values:
- 0.0
- 0.1
- 0.2
final_dropout:
values:
- 0.0
- 0.1
- 0.2
mask_time_prob:
values:
- 0.0
- 0.1
- 0.2
mask_time_length:
values:
- 10
- 15
- 20
mask_feature_prob:
values:
- 0.0
- 0.1
- 0.2
mask_feature_length:
values:
- 10
- 15
- 20
program: run_audio_classification.py
project: mms-lid-accent-classification
\ No newline at end of file
audio_classification_scripts/run_mms_lid.sh deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env bash
CUDA_VISIBLE_DEVICES=2 python run_audio_classification_one_layer.py \
--model_name_or_path "facebook/mms-lid-4017" \
--train_dataset_name "stable-speech/concatenated-normalized-accent-dataset" \
--train_dataset_config_name "default" \
--train_split_name "train" \
--train_label_column_name "labels" \
--eval_dataset_name "stable-speech/concatenated-normalized-accent-dataset" \
--eval_dataset_config_name "default" \
--eval_split_name "test" \
--eval_label_column_name "labels" \
--output_dir "./tmp/" \
--do_train \
--do_eval \
--overwrite_output_dir \
--remove_unused_columns false \
--fp16 \
--fp16_full_eval \
--learning_rate 1e-4 \
--max_length_seconds 20 \
--min_length_seconds 5 \
--attention_mask \
--warmup_steps 100 \
--max_steps 2000 \
--per_device_train_batch_size 32 \
--per_device_eval_batch_size 32 \
--preprocessing_num_workers 4 \
--dataloader_num_workers 4 \
--logging_strategy "steps" \
--logging_steps 10 \
--evaluation_strategy "steps" \
--eval_steps 300 \
--save_strategy "no" \
--save_steps 2000 \
--freeze_base_model true \
--freeze_feature_encoder true \
--push_to_hub false \
--trust_remote_code \
--use_weighted_layer_sum true \
audio_classification_scripts/run_mms_lid_with_cv.sh deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env bash
python run_audio_classification.py \
--model_name_or_path "facebook/mms-lid-126" \
--train_dataset_name "stable-speech/concatenated-normalized-accent-dataset+stable-speech/concatenated-common-voice-15-accented" \
--train_dataset_config_name "default+default" \
--train_split_name "train+train" \
--train_label_column_name "labels+labels" \
--eval_dataset_name "stable-speech/concatenated-normalized-accent-dataset" \
--eval_dataset_config_name "default" \
--eval_split_name "test" \
--eval_label_column_name "labels" \
--output_dir "./" \
--do_train \
--do_eval \
--overwrite_output_dir \
--remove_unused_columns False \
--fp16 \
--fp16_full_eval \
--learning_rate 1e-4 \
--lr_scheduler_type "constant_with_warmup" \
--max_length_seconds 20 \
--min_length_seconds 5 \
--attention_mask \
--warmup_steps 100 \
--max_steps 5000 \
--per_device_train_batch_size 32 \
--per_device_eval_batch_size 32 \
--preprocessing_num_workers 4 \
--dataloader_num_workers 4 \
--logging_strategy "steps" \
--logging_steps 10 \
--evaluation_strategy "steps" \
--eval_steps 1000 \
--save_strategy "no" \
--save_steps 5000 \
--filter_threshold 0.01 \
--freeze_base_model False \
--gradient_checkpointing \
--push_to_hub False \
--trust_remote_code
audio_classification_scripts/run_sweep.yaml deleted 100644 → 0
View file @ 64cfa64e
command:
- python3
- ${program}
- --fp16
- --fp16_full_eval
- --do_train
- --do_eval
- --trust_remote_code
- --overwrite_output_dir
- ${args}
method: random
metric:
goal: maximize
name: eval/accuracy
parameters:
model_name_or_path:
value: facebook/mms-lid-4017
train_dataset_name:
value: "stable-speech/concatenated-normalized-accent-dataset+stable-speech/concatenated-common-voice-15-accented"
train_dataset_config_name:
value: "default+default"
train_split_name:
value: "train+train"
train_label_column_name:
value: "labels+labels"
eval_dataset_name:
value: stable-speech/concatenated-normalized-accent-dataset
eval_dataset_config_name:
value: default
eval_split_name:
value: test
eval_label_column_name:
value: labels
output_dir:
value: "/raid/yoach/tmp/"
remove_unused_columns:
value: false
learning_rate:
distribution: log_uniform_values
min: 3e-6
max: 0.01
lr_scheduler_type:
value: constant
max_length_seconds:
value: 20 # give some data diversity for longer audio samples
min_length_seconds:
value: 5
attention_mask:
values:
- true
num_train_epochs:
values:
- 2
- 5
- 10
- 20
- 40
- 60
per_device_train_batch_size:
value: 32
per_device_eval_batch_size:
value: 32
preprocessing_num_workers:
value: 8
dataloader_num_workers:
value: 8
logging_strategy:
value: steps
logging_steps:
value: 10
evaluation_strategy:
value: steps
eval_steps:
value: 2000
save_strategy:
value: "no"
save_steps:
value: 2000
metric_for_best_model:
value: accuracy
push_to_hub:
value: false
use_weighted_layer_sum:
value: false
freeze_base_model:
value: true
max_samples_per_label:
value: 10000
save_to_disk:
value: "/raid/yoach/tmp_dataset_accents/"
temporary_save_to_disk:
value: "/raid/yoach/tmp_hidden_states/"
use_last_embedding_layer:
value: true
filter_threshold:
value: "0.001"
program: run_audio_classification.py
project: mms-lid-accent-classification-v2
audio_classification_scripts/run_wav2vec2_dummy.sh deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env bash
python run_audio_classification.py \
--model_name_or_path "hf-internal-testing/tiny-random-wav2vec2" \
--train_dataset_name "facebook/voxpopuli" \
--train_dataset_config_name "en_accented" \
--train_split_name "test" \
--train_label_column_name "accent" \
--eval_dataset_name "facebook/voxpopuli" \
--eval_dataset_config_name "en_accented" \
--eval_split_name "test" \
--eval_label_column_name "accent" \
--trust_remote_code \
--output_dir "./" \
--do_train \
--do_eval \
--max_train_samples 100 \
--max_eval_samples 100 \
--overwrite_output_dir \
--remove_unused_columns False \
--fp16 \
--learning_rate 1e-4 \
--min_length_seconds 5 \
--max_length_seconds 10 \
--attention_mask False \
--warmup_ratio 0.1 \
--num_train_epochs 5 \
--per_device_train_batch_size 4 \
--per_device_eval_batch_size 4 \
--dataloader_num_workers 0 \
--logging_strategy "steps" \
--logging_steps 10 \
--evaluation_strategy "epoch" \
--save_strategy "epoch" \
--load_best_model_at_end True \
--metric_for_best_model "accuracy" \
--save_total_limit 3 \
--seed 0
dataset_concatenation_scripts/run_dataset_concatenation.sh deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env bash
python run_dataset_concatenation.py \
--dataset_name "sanchit-gandhi/vctk+facebook/voxpopuli+sanchit-gandhi/edacc-normalized" \
--dataset_config_name "default+en_accented+default" \
--dataset_split_name "train+test+validation" \
--label_column_name "accent+accent+accent" \
--text_column_name "text+normalized_text+text" \
--speaker_column_name "speaker_id+speaker_id+speaker" \
--batch_size 500 \
--output_dir "./concatenated-dataset"
python run_dataset_concatenation.py \
--dataset_name "sanchit-gandhi/edacc-normalized" \
--dataset_config_name "default" \
--dataset_split_name "test" \
--label_column_name "accent" \
--text_column_name "text" \
--speaker_column_name "speaker" \
--batch_size 500 \
--output_dir "./concatenated-dataset-test"
dataset_concatenation_scripts/run_dataset_concatenation_cv.sh deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env bash
python run_dataset_concatenation.py \
--dataset_name "stable-speech/common_voice_15_0_accented" \
--dataset_config_name "en" \
--dataset_split_name "train" \
--label_column_name "accent" \
--text_column_name "sentence" \
--speaker_column_name "client_id" \
--batch_size 250 \
--preprocessing_num_workers 4 \
--output_dir "./concatenated-dataset-cv"
python run_dataset_concatenation.py \
--dataset_name "stable-speech/common_voice_15_0_accented" \
--dataset_config_name "en" \
--dataset_split_name "test" \
--label_column_name "accent" \
--text_column_name "sentence" \
--speaker_column_name "client_id" \
--batch_size 250 \
--preprocessing_num_workers 4 \
--output_dir "./concatenated-dataset-cv-test"
edacc/prepare_edacc.py deleted 100644 → 0
View file @ 64cfa64e
import csv
import os
import re
import shutil
import sys
from dataclasses import dataclass, field
import soundfile as sf
from datasets import Audio, Dataset, DatasetDict, load_dataset
from tqdm import tqdm
from transformers import HfArgumentParser
@dataclass
class DataTrainingArguments:
"""
Arguments pertaining to what data we are going to input our data for prepareation
"""
dataset_dir: str = field(
default=None,
metadata={
"help": "Path where the EdAcc tar.gz archive is extracted. Leave in it's raw format: the script will "
"assume it's unchanged from the download and use relative paths to load the relevant audio files."
},
)
output_dir: str = field(
default=None,
metadata={
"help": "Where to save the processed dataset to disk. If unspecified, uses a 'pretty' version of the "
"original dataset name. E.g. 'facebook/voxpopuli' will be saved under 'voxpopuli'."
},
)
overwrite_output_dir: bool = field(
default=True,
metadata={"help": "Overwrite the content of the output directory."},
)
push_to_hub: bool = field(
default=False,
metadata={"help": "Whether or not to push the processed dataset to the Hub."},
)
hub_dataset_id: str = field(
default=False,
metadata={"help": "Repository namespace if pushing to the Hugging Face Hub."},
)
private_repo: bool = field(
default=True,
metadata={"help": "Whether or not to push the processed dataset to a private repository on the Hub"},
)
max_samples: int = field(
default=None,
metadata={"help": "Maximum number of samples per split. Useful for debugging purposes."},
)
def main():
# 1. Parse input arguments
parser = HfArgumentParser(DataTrainingArguments)
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0]
else:
data_args = parser.parse_args_into_dataclasses()[0]
# 1. Load accents for each speaker
linguistic_background = {}
linguistic_background_csv = os.path.join(data_args.dataset_dir, "linguistic_background.csv")
with open(linguistic_background_csv, encoding="utf-8") as file:
reader = csv.DictReader(file, delimiter=",")
for line in reader:
linguistic_background[line["PARTICIPANT_ID"]] = line[
"How would you describe your accent in English? (e.g. Italian, Glaswegian)"
]
accent_dataset = load_dataset("sanchit-gandhi/edacc_accents", split="train")
def format_dataset(batch):
batch["speaker_id"] = (
batch["Final-Participant_ID"].replace("EAEC", "EDACC").replace("P1", "-A").replace("P2", "-B")
)
return batch
accent_dataset = accent_dataset.map(format_dataset, remove_columns=["Final-Participant_ID"])
# 2. Clean accents for each speaker
linguistic_background_clean = {
participant: accent.strip()
for participant, accent in zip(accent_dataset["speaker_id"], accent_dataset["English_Variety"])
}
linguistic_variety = {
participant: l1.strip() for participant, l1 in zip(accent_dataset["speaker_id"], accent_dataset["L1_Variety"])
}
# 3. Initialize dataset dict
raw_datasets = DatasetDict()
if data_args.overwrite_output_dir and os.path.exists(data_args.output_dir) and os.path.isdir(data_args.output_dir):
shutil.rmtree(data_args.output_dir)
output_dir_processed = os.path.join(data_args.output_dir, "processed")
# 4. Iterate over dev/test files
for split, split_formatted in zip(["dev", "test"], ["validation", "test"]):
data_dir = os.path.join(data_args.dataset_dir, split)
metadata = os.path.join(data_dir, "stm")
output_dir_split = os.path.join(output_dir_processed, split)
os.makedirs(output_dir_split, exist_ok=True)
all_speakers = []
all_genders = []
all_l1s = []
all_texts = []
all_audio_paths = []
all_normalized_accents = []
all_raw_accents = []
current_audio = None
current_audio_array = None
current_sampling_rate = None
current_counter = 1
gender_pat = r".*?\<(.*),.*"
l1_pat = r".*?\,(.*)>.*"
with open(metadata, "r") as file:
for idx, line in tqdm(enumerate(file), desc=split):
# example line is: 'EDACC-C06 1 EDACC-C06-A 0.00 5.27 <male,l1> C ELEVEN DASH P ONE\n
# the transcription always comes to the right of the last rangle bracket
text_idx = line.find(">") + 1
all_texts.append(line[text_idx + 1 : -1])
# the metadata immediately proceeds this
line = line[:text_idx]
file, channel, speaker, start, end, gender_l1 = line.split(" ")
# add speaker information to cumulative lists
all_raw_accents.append(linguistic_background[speaker])
all_normalized_accents.append(linguistic_background_clean[speaker])
all_speakers.append(speaker)
# add gender/l1 information
all_genders.append(re.search(gender_pat, gender_l1).group(1))
all_l1s.append(linguistic_variety[speaker])
# read audio file if different from previous
if file != current_audio:
current_audio_array, current_sampling_rate = sf.read(
os.path.join(data_args.dataset_dir, "data", file + ".wav")
)
current_audio = file
current_counter = 1
else:
current_counter += 1
# chunk audio file according to start/end times
start = int(float(start) * current_sampling_rate)
end = int(float(end) * current_sampling_rate)
end = min(end, len(current_audio_array))
chunked_audio = current_audio_array[start:end]
save_path = os.path.join(output_dir_split, f"{file}-{current_counter}.wav")
sf.write(save_path, chunked_audio, current_sampling_rate)
all_audio_paths.append(save_path)
if data_args.max_samples is not None and (data_args.max_samples - 1) == idx:
break
raw_datasets[split_formatted] = Dataset.from_dict(
{
"speaker": all_speakers,
"text": all_texts,
"accent": all_normalized_accents,
"raw_accent": all_raw_accents,
"gender": all_genders,
"l1": all_l1s,
"audio": all_audio_paths,
}
).cast_column("audio", Audio())
if data_args.push_to_hub:
raw_datasets.push_to_hub(data_args.hub_dataset_id, token=True)
raw_datasets.save_to_disk(data_args.output_dir)
if __name__ == "__main__":
main()
edacc/run_edacc.sh deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env bash
python prepare_edacc.py \
--dataset_dir "/fsx/sanchit/edacc/edacc_v1.0" \
--output_dir "/fsx/sanchit/edacc_processed" \
--hub_dataset_id "sanchit-gandhi/edacc-normalized" \
--push_to_hub
run_audio_classification.py deleted 100644 → 0
View file @ 64cfa64e
#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os
import re
import sys
from collections import Counter
from dataclasses import dataclass, field
from random import randint
from typing import List, Optional, Union, Dict
import torch
import datasets
import evaluate
import numpy as np
import transformers
from datasets import Dataset, DatasetDict, IterableDataset, concatenate_datasets, interleave_datasets, load_dataset
from tqdm import tqdm
from transformers import (
AutoConfig,
AutoFeatureExtractor,
AutoModelForAudioClassification,
HfArgumentParser,
Trainer,
TrainingArguments,
set_seed,
)
from transformers.models.whisper.tokenization_whisper import LANGUAGES
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version
from transformers.trainer_pt_utils import LengthGroupedSampler
from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Model
from transformers import Wav2Vec2BertForSequenceClassification, Wav2Vec2BertModel
from transformers.models.wav2vec2.modeling_wav2vec2 import _HIDDEN_STATES_START_POSITION
from transformers.modeling_outputs import SequenceClassifierOutput
from torch import nn
from torch.nn import CrossEntropyLoss
logger = logging.getLogger(__name__)
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.38.0.dev0")
def random_subsample(wav: np.ndarray, max_length: float, sample_rate: int = 16000) -> np.ndarray:
"""Randomly sample chunks of `max_length` seconds from the input audio"""
sample_length = int(round(sample_rate * max_length))
if len(wav) <= sample_length:
return wav
random_offset = randint(0, len(wav) - sample_length - 1)
return wav[random_offset : random_offset + sample_length]
def deterministic_subsample(wav: np.ndarray, max_length: float, sample_rate: int = 16000) -> np.ndarray:
"""Take first `max_length` seconds from the input audio"""
sample_length = int(round(sample_rate * max_length))
if len(wav) <= sample_length:
return wav
return wav[0:sample_length]
class SequenceClassificationModel(Wav2Vec2ForSequenceClassification):
def __init__(self, config):
super().__init__(config)
if hasattr(config, "add_adapter") and config.add_adapter:
raise ValueError(
"Sequence classification does not support the use of Wav2Vec2 adapters (config.add_adapter=True)"
)
self.wav2vec2 = Wav2Vec2Model(config)
num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings
if config.use_weighted_layer_sum:
self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
# To bypass w2v2
self.compute_w2v2 = True
# Initialize weights and apply final processing
self.post_init()
def forward(
self,
input_features=None,
attention_mask=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
labels=None,
hidden_states=None, # added
):
r"""
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
output_hidden_states = True
if self.compute_w2v2:
outputs = self.wav2vec2(
input_features,
attention_mask=attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
hidden_states = outputs[_HIDDEN_STATES_START_POSITION][-1] # take last embedding layer
if attention_mask is None:
pooled_output = hidden_states.mean(dim=1)
else:
padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)
hidden_states[~padding_mask] = 0.0
pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1)
else:
pooled_output = hidden_states
logits = self.classifier(pooled_output)
loss = None
if labels is not None:
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
if not return_dict:
output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
return ((loss,) + output) if loss is not None else output
return SequenceClassifierOutput(
loss=loss,
logits=logits,
)
# This list first defines the accent prefixes, which we use to strip the accent from CV
# e.g. England, southern accent, slight west-country expression -> England
# TODO(YL): update this with any CV test prefixes not present in the train set
STARTS_WITH = [
"Afrikaans",
"American",
"Australian",
"Bangladeshi",
"Canadian",
"Chinese",
"Dutch",
"Eastern European",
"European",
"England",
"English",
"German",
"Filipino",
"India",
"Irish",
"Israeli",
"Italian",
"Japanese",
"Kenyan",
"Northern Irish",
"New Zealand",
"Nigerian",
"Malaysian",
"Russian",
"Scottish",
"Singaporean",
"Slavic",
"South African",
"Southern African",
"Swedish",
"Swiss",
"United States English",
"West Indies",
"french",
"polish",
"serbian",
]
# This dictionary is used to map the un-normalised accent names to normalised ones
# TODO(YL): update this with any CV test mappings not present in the train set
ACCENT_MAPPING = {
"British": "English",
# "Canadian": "American", TODO(SG): decide whether to normalize these to closely related accents
# "New zealand": "Australian",
"Northern irish": "Irish",
"Pakistani": "Indian",
"Mainstream u s english": "American",
"Southern british english": "English",
"Indian english": "Indian",
"Scottish english": "Scottish",
"Don't know": "Unknown",
"Nigerian english": "Nigerian",
"Kenyan english": "Kenyan",
"Ghanain english": "Ghanain",
"Jamaican english": "Jamaican",
"Indonesian english": "Indonesian",
"South african english": "South african",
"Irish english": "Irish",
"Latin": "Latin american",
"European": "Unknown", # Too general
"Eastern european": "Eastern european", # TODO(SG): keep for now, but maybe remove later as too general
"Bangladeshi": "Indian",
"England": "English",
"India": "Indian",
"Afrikaans": "South african",
"California": "American",
"Nepali": "Indian",
"New york city": "American",
"New jerseyan": "American",
"Northumbrian british english": "English",
"Nottinghamshire,east midlands": "English",
"Southern african": "South african",
"United states english": "American",
"West indies": "Jamaican",
"2nd language": "Unknown", # Too vague
"A savage texas gentleman": "American",
"A variety of texan english with some german influence that has undergone the cot-caught merger": "American",
"A'lo": "Unknown", # Unclear
"Academic southern english,england english": "English",
"Argentinian english": "Latin american",
"Austrian": "German",
"Bangladesh,india and south asia (india, pakistan, sri lanka)": "Indian",
"Brazillian accent": "Brazilian",
"British accent": "English",
"Caribbean canadian": "Unknown", # Specific combination not listed
"Colombian accent": "Latin american",
"Czech accent": "Czech",
"East african khoja": "Unknown", # Specific community
"East indian": "Indian",
"East london": "English",
"England,london,academic": "English",
"Filipino": "Unknown", # Unique blend
"Fluent,e sl,european": "Unknown", # Too vague
"Generic european": "Unknown", # Too vague
"Georgian english": "Unknown", # No direct match
"Ghanaian english accent,african regular reader": "Unknown", # Specific category not listed
"Haitian creole": "Unknown", # Unique blend
"Hispanic": "Latin american",
"Hispanic/latino": "Latin american",
"Hong kong english": "Chinese",
"Hong kong english,scottish english": "Chinese",
"Hunglish": "Hungarian",
"I think mine accent is influenced by indian accent ,yes please. ,india and south asia (india, pakistan, sri lanka)": "Indian",
"I was born in england and have lived in australia, canada and france.": "English",
"International english,united states english,australian english": "American",
"Israeli": "Unknown", # No direct match
"Israeli english": "Unknown", # No direct match
"Javanese,indonesian english,malaysian english": "Indonesian",
"Kazakhstan english": "Unknown", # No direct match
"Kiwi": "New zealand", # Could be generalised to Australian
"Latin america,united states english": "Latin american",
"Latin american accent": "Latin american",
"Latin english": "Unknown", # Too vague
"Latino": "Latin american",
"Latvian": "Latvian", # Note: added new
"Little latino,united states english,second language": "Latin american",
"Liverpool english,lancashire english,england english": "English",
"Liverpudlian english": "English",
"Malaysian english": "Malaysian", # Note: added new
"Mexican accent": "Latin american",
"Mid-atlantic united states english,philadelphia, pennsylvania, united states english,united states english,philadelphia style united states english": "American",
"Mid-atlantic,england english,united states english": "American",
"Midatlantic,england english": "American",
"Midwestern states (michigan),united states english": "American",
"Mild northern england english": "English",
"Minor french accent": "French",
"Mix of american and british ,native polish": "Polish",
"Mix of american and british accent": "Unknown", # Combination not clearly mapped
"Mostly american with some british and australian inflections": "Unknown", # Combination not clearly mapped
"My accent is influenced by the phones of all letters within a sentence.,southern african (south africa, zimbabwe, namibia)": "South african",
"New zealand english": "New Zealand English",
"Nigeria english": "Nigerian", # Note: added new
"Non native speaker from france": "French",
"Non-native": "Unknown", # Too vague
"Non-native,german accent": "German",
"North european english": "Unknown", # Too broad
"Norwegian": "Norwegian", # Note: added new
"Ontario,canadian english": "Canadian", # Note: added new
"Polish english": "Polish",
"Rhode island new england accent": "American",
"Singaporean english": "Singaporean", # Note: added new
"Slavic": "Eastern european",
"Slighty southern affected by decades in the midwest, 4 years in spain and germany, speak some german, spanish, polish. have lived in nine states.": "Unknown", # Complex blend
"South african": "South african",
"South atlantic (falkland islands, saint helena)": "Unknown", # Specific regions not listed
"South australia": "Australian",
"South indian": "Indian",
"Southern drawl": "American",
"Southern texas accent,united states english": "American",
"Southern united states,united states english": "American",
"Spanish bilingual": "Spanish",
"Spanish,foreign,non-native": "Spanish",
"Strong latvian accent": "Latvian",
"Swedish accent": "Swedish", # Note: added new
"Transnational englishes blend": "Unknown", # Too vague
"U.k. english": "English",
"Very slight russian accent,standard american english,boston influence": "American",
"Welsh english": "Welsh",
"West african": "Unknown", # No specific West African category
"West indian": "Unknown", # Caribbean, but no specific match
"Western europe": "Unknown", # Too broad
"With heavy cantonese accent": "Chinese",
}
def preprocess_labels(label: str) -> str:
"""Apply pre-processing formatting to the accent labels"""
if "_" in label:
# voxpopuli stylises the accent as a language code (e.g. en_pl for "polish") - convert to full accent
language_code = label.split("_")[-1]
label = LANGUAGES[language_code]
# VCTK labels for two words are concatenated into one (NewZeleand-> New Zealand)
label = re.sub(r"(\w)([A-Z])", r"\1 \2", label).strip()
for prefix in STARTS_WITH:
if label.startswith(prefix):
label = prefix
# convert Whisper language code (polish) to capitalised (Polish)
label = label.capitalize()
if label in ACCENT_MAPPING:
label = ACCENT_MAPPING[label]
return label
@dataclass
class DataCollatorFeatureExtractorWithPadding:
"""
Data collator that will dynamically pad the inputs received to the longest sequence in the batch.
"""
feature_extractor: AutoFeatureExtractor
max_length_seconds: int
feature_extractor_input_name: Optional[str] = "input_values"
def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
# split inputs and labels since they have to be of different lengths and need
# different padding methods
audios = [torch.tensor(deterministic_subsample(feature["audio"]["array"], max_length=self.max_length_seconds, sample_rate=self.feature_extractor.sampling_rate)).squeeze().numpy() for feature in features]
batch = self.feature_extractor(audios, return_tensors="pt", padding="longest", return_attention_mask=True)
batch["labels"] = torch.tensor([feature["labels_id"] for feature in features])
return batch
@dataclass
class DataCollatorHiddenStatesPadding:
"""
Data collator that will dynamically pad the inputs received to the longest sequence in the batch.
"""
def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
hidden_states = torch.stack([torch.tensor(feature["hidden_states"]) for feature in features])
batch = {"hidden_states": hidden_states}
batch["labels"] = torch.tensor([feature["labels_id"] for feature in features])
return batch
@dataclass
class DataTrainingArguments:
"""
Arguments pertaining to what data we are going to input our model for training and eval.
Using `HfArgumentParser` we can turn this class
into argparse arguments to be able to specify them on
the command line.
"""
train_dataset_name: str = field(
default=None,
metadata={
"help": "The name of the training dataset to use (via the datasets library). Load and combine "
"multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
" librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
},
)
train_dataset_config_name: Optional[str] = field(
default=None,
metadata={
"help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
"multiple datasets by separating dataset configs by a '+' symbol."
},
)
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'")
},
)
train_dataset_samples: str = field(
default=None,
metadata={
"help": "Number of samples in the training data. Load and combine "
"multiple datasets by separating dataset samples by a '+' symbol."
},
)
eval_dataset_name: str = field(
default=None,
metadata={
"help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset name if unspecified."
},
)
eval_dataset_config_name: Optional[str] = field(
default=None,
metadata={
"help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset config name if unspecified"
},
)
eval_split_name: str = field(
default="validation",
metadata={
"help": (
"The name of the evaluation data set split to use (via the datasets"
" library). Defaults to 'validation'"
)
},
)
audio_column_name: str = field(
default="audio",
metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
)
train_label_column_name: str = field(
default="labels",
metadata={
"help": "The name of the dataset column containing the labels in the train set. Defaults to 'label'"
},
)
eval_label_column_name: str = field(
default="labels",
metadata={"help": "The name of the dataset column containing the labels in the eval set. Defaults to 'label'"},
)
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."
)
},
)
max_length_seconds: Optional[float] = field(
default=20,
metadata={"help": "Audio samples will be randomly cut to this length during training if the value is set."},
)
min_length_seconds: Optional[float] = field(
default=5,
metadata={"help": "Audio samples less than this value will be filtered during training if the value is set."},
)
preprocessing_num_workers: Optional[int] = field(
default=None,
metadata={"help": "The number of processes to use for the preprocessing."},
)
filter_threshold: Optional[float] = field(
default=1.0,
metadata={"help": "Filter labels that occur less than `filter_threshold` percent in the training/eval data."},
)
max_samples_per_label: Optional[int] = field(
default=None,
metadata={
"help": (
"If set, randomly limits the number of samples per label."
)
},
)
save_to_disk: str = field(
default=None,
metadata={
"help": "If set, will save the dataset to this path if this is an empyt folder. If not empty, will load the datasets from it."
}
)
temporary_save_to_disk: str = field(
default=None,
metadata={
"help": "Temporarily save audio labels here."
}
)
@dataclass
class ModelArguments:
"""
Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
"""
model_name_or_path: str = field(
default="facebook/wav2vec2-base",
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models. Only works with Wav2Vec2 models"},
)
config_name: Optional[str] = field(
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
)
cache_dir: Optional[str] = field(
default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from the Hub"}
)
model_revision: str = field(
default="main",
metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
)
feature_extractor_name: Optional[str] = field(
default=None, metadata={"help": "Name or path of preprocessor config."}
)
freeze_feature_encoder: bool = field(
default=False,
metadata={
"help": "Whether to freeze the feature encoder layers of the model. Only relevant for Wav2Vec2-style models."
},
)
freeze_base_model: bool = field(
default=True, metadata={"help": "Whether to freeze the base encoder of the model."}
)
attention_mask: bool = field(
default=True, metadata={"help": "Whether to generate an attention mask in the feature extractor."}
)
token: str = field(
default=None,
metadata={
"help": (
"The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
"generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
)
},
)
trust_remote_code: bool = field(
default=False,
metadata={
"help": (
"Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
"should only be set to `True` for repositories you trust and in which you have read the code, as it will "
"execute code present on the Hub on your local machine."
)
},
)
ignore_mismatched_sizes: bool = field(
default=True,
metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."},
)
attention_dropout: float = field(
default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}
)
activation_dropout: float = field(
default=0.0, metadata={"help": "The dropout ratio for activations inside the fully connected layer."}
)
feat_proj_dropout: float = field(default=0.0, metadata={"help": "The dropout ratio for the projected features."})
hidden_dropout: float = field(
default=0.0,
metadata={
"help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."
},
)
final_dropout: float = field(
default=0.0,
metadata={"help": "The dropout probability for the final projection layer."},
)
mask_time_prob: float = field(
default=0.05,
metadata={
"help": (
"Probability of each feature vector along the time axis to be chosen as the start of the vector "
"span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature "
"vectors will be masked along the time axis."
)
},
)
mask_time_length: int = field(
default=10,
metadata={"help": "Length of vector span to mask along the time axis."},
)
mask_feature_prob: float = field(
default=0.0,
metadata={
"help": (
"Probability of each feature vector along the feature axis to be chosen as the start of the vectorspan"
" to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature"
" bins will be masked along the time axis."
)
},
)
mask_feature_length: int = field(
default=10,
metadata={"help": "Length of vector span to mask along the feature axis."},
)
layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
use_weighted_layer_sum: bool = field(default=False, metadata={"help": "Whether to use a weighted average of layer outputs with learned weights."})
use_last_embedding_layer: bool = field(default=False, metadata={"help": "Whether to use the last layer hidden state. Only work with W2V model."})
def convert_dataset_str_to_list(
dataset_names,
dataset_config_names,
splits=None,
label_column_names=None,
dataset_samples=None,
default_split="train",
):
if isinstance(dataset_names, str):
dataset_names = dataset_names.split("+")
dataset_config_names = dataset_config_names.split("+")
splits = splits.split("+") if splits is not None else None
label_column_names = label_column_names.split("+") if label_column_names is not None else None
dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
# basic checks to ensure we've got the right number of datasets/configs/splits/columns/probs
if len(dataset_names) != len(dataset_config_names):
raise ValueError(
f"Ensure one config is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(dataset_config_names)} configs."
)
if splits is not None and len(splits) != len(dataset_names):
raise ValueError(
f"Ensure one split is passed for each dataset, got {len(dataset_names)} datasets and {len(splits)} splits."
)
if label_column_names is not None and len(label_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one label column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(label_column_names)} label column names."
)
if dataset_samples is not None:
if len(dataset_samples) != len(dataset_names):
raise ValueError(
f"Ensure one sample is passed for each dataset, got {len(dataset_names)} datasets and "
f"{len(dataset_samples)} samples."
)
dataset_samples = [float(ds_sample) for ds_sample in dataset_samples]
else:
dataset_samples = [None] * len(dataset_names)
label_column_names = (
label_column_names if label_column_names is not None else ["labels" for _ in range(len(dataset_names))]
)
splits = splits if splits is not None else [default_split for _ in range(len(dataset_names))]
dataset_names_dict = []
for i, ds_name in enumerate(dataset_names):
dataset_names_dict.append(
{
"name": ds_name,
"config": dataset_config_names[i],
"split": splits[i],
"label_column_name": label_column_names[i],
"samples": dataset_samples[i],
}
)
return dataset_names_dict
def load_multiple_datasets(
dataset_names: Union[List, str],
dataset_config_names: Union[List, str],
splits: Optional[Union[List, str]] = None,
label_column_names: Optional[List] = None,
sampling_rate: Optional[int] = 16000,
stopping_strategy: Optional[str] = "first_exhausted",
dataset_samples: Optional[Union[List, np.array]] = None,
streaming: Optional[bool] = False,
seed: Optional[int] = None,
audio_column_name: Optional[str] = "audio",
**kwargs,
) -> Union[Dataset, IterableDataset]:
dataset_names_dict = convert_dataset_str_to_list(
dataset_names, dataset_config_names, splits, label_column_names, dataset_samples
)
if dataset_samples is not None:
dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
else:
probabilities = None
all_datasets = []
# iterate over the datasets we want to interleave
for dataset_dict in tqdm(dataset_names_dict, desc="Combining datasets..."):
dataset = load_dataset(
dataset_dict["name"],
dataset_dict["config"],
split=dataset_dict["split"],
streaming=streaming,
**kwargs,
)
dataset_features = dataset.features.keys()
if audio_column_name not in dataset_features:
raise ValueError(
f"Audio column name '{audio_column_name}' not found in dataset"
f" '{dataset_dict['name']}'. Make sure to set `--audio_column_name` to"
f" the correct audio column - one of {', '.join(dataset_features)}."
)
# resample to specified sampling rate
dataset = dataset.cast_column("audio", datasets.features.Audio(sampling_rate))
if dataset_dict["label_column_name"] not in dataset_features:
raise ValueError(
f"Label column name {dataset_dict['label_column_name']} not found in dataset"
f" '{dataset_dict['name']}'. Make sure to set `--label_column_name` to the"
f" correct text column - one of {', '.join(dataset_features)}."
)
# blanket renaming of all label columns to label
if dataset_dict["label_column_name"] != "labels":
dataset = dataset.rename_column(dataset_dict["label_column_name"], "labels")
dataset_features = dataset.features.keys()
columns_to_keep = {"audio", "labels"}
dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
all_datasets.append(dataset)
if len(all_datasets) == 1:
# we have a single dataset so just return it as is
return all_datasets[0]
if streaming:
interleaved_dataset = interleave_datasets(
all_datasets,
stopping_strategy=stopping_strategy,
probabilities=probabilities,
seed=seed,
)
else:
interleaved_dataset = concatenate_datasets(all_datasets)
return interleaved_dataset
def main():
# See all possible arguments in src/transformers/training_args.py
# or by passing the --help flag to this script.
# We now keep distinct sets of args, for a cleaner separation of concerns.
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
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()
# 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)],
)
if training_args.should_log:
# The default of training_args.log_level is passive, so we set log level at info here to have that default.
transformers.utils.logging.set_verbosity_info()
log_level = training_args.get_process_log_level()
logger.setLevel(log_level)
transformers.utils.logging.set_verbosity(log_level)
transformers.utils.logging.enable_default_handler()
transformers.utils.logging.enable_explicit_format()
# Log on each process the small summary:
logger.warning(
f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+ f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
)
logger.info(f"Training/evaluation parameters {training_args}")
# Set seed before initializing model.
set_seed(training_args.seed)
# Detecting last checkpoint.
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 train from scratch."
)
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."
)
if data_args.save_to_disk is not None:
os.makedirs(data_args.save_to_disk, exist_ok=True)
# assume that the dataset has been saved to `save_to_disk` if the latter is not empty
dataset_was_precomputed = len(os.listdir(data_args.save_to_disk)) > 0
if dataset_was_precomputed:
raw_datasets = datasets.load_from_disk(data_args.save_to_disk)
else:
# Initialize our dataset and prepare it for the audio classification task.
raw_datasets = DatasetDict()
if training_args.do_train:
raw_datasets["train"] = load_multiple_datasets(
data_args.train_dataset_name,
data_args.train_dataset_config_name,
splits=data_args.train_split_name,
label_column_names=data_args.train_label_column_name,
dataset_samples=data_args.train_dataset_samples,
seed=training_args.seed,
cache_dir=model_args.cache_dir,
token=True if model_args.token else None,
trust_remote_code=model_args.trust_remote_code,
num_proc=data_args.preprocessing_num_workers,
# streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
)
if training_args.do_eval:
dataset_names_dict = convert_dataset_str_to_list(
data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
(
data_args.eval_dataset_config_name
if data_args.eval_dataset_config_name
else data_args.train_dataset_config_name
),
splits=data_args.eval_split_name,
label_column_names=data_args.eval_label_column_name,
)
all_eval_splits = []
# load multiple eval sets
for dataset_dict in dataset_names_dict:
pretty_name = (
f"{dataset_dict['name'].split('/')[-1]}/{dataset_dict['split'].replace('.', '-')}"
if len(dataset_names_dict) > 1
else "eval"
)
all_eval_splits.append(pretty_name)
raw_datasets[pretty_name] = load_dataset(
dataset_dict["name"],
dataset_dict["config"],
split=dataset_dict["split"],
cache_dir=model_args.cache_dir,
token=True if model_args.token else None,
trust_remote_code=model_args.trust_remote_code,
num_proc=data_args.preprocessing_num_workers,
# streaming=data_args.streaming,
)
features = raw_datasets[pretty_name].features.keys()
if dataset_dict["label_column_name"] not in features:
raise ValueError(
f"--label_column_name {data_args.eval_label_column_name} not found in dataset '{data_args.dataset_name}'. "
"Make sure to set `--label_column_name` to the correct text column - one of "
f"{', '.join(raw_datasets['train'].column_names)}."
)
elif dataset_dict["label_column_name"] != "labels":
raw_datasets[pretty_name] = raw_datasets[pretty_name].rename_column(
dataset_dict["label_column_name"], "labels"
)
raw_datasets[pretty_name] = raw_datasets[pretty_name].remove_columns(
set(raw_datasets[pretty_name].features.keys()) - {"audio", "labels"}
)
if not training_args.do_train and not training_args.do_eval:
raise ValueError(
"Cannot not train and not do evaluation. At least one of training or evaluation has to be performed."
)
# Setting `return_attention_mask=True` is the way to get a correctly masked mean-pooling over
# transformer outputs in the classifier, but it doesn't always lead to better accuracy
feature_extractor = AutoFeatureExtractor.from_pretrained(
model_args.feature_extractor_name or model_args.model_name_or_path,
return_attention_mask=model_args.attention_mask,
cache_dir=model_args.cache_dir,
revision=model_args.model_revision,
token=model_args.token,
trust_remote_code=model_args.trust_remote_code,
)
feature_extractor_input_name = feature_extractor.model_input_names[0]
if not dataset_was_precomputed:
# `datasets` takes care of automatically loading and resampling the audio,
# so we just need to set the correct target sampling rate.
raw_datasets = raw_datasets.cast_column(
data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
)
with training_args.main_process_first():
if training_args.do_train:
if data_args.max_train_samples is not None:
raw_datasets["train"] = (
raw_datasets["train"].shuffle(seed=training_args.seed).select(range(data_args.max_train_samples))
)
if training_args.do_eval:
if data_args.max_eval_samples is not None:
raw_datasets["eval"] = (
raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
)
sampling_rate = feature_extractor.sampling_rate
model_input_name = feature_extractor.model_input_names[0]
if not dataset_was_precomputed:
def prepare_dataset(audio, labels):
batch = {}
batch["length"] = len(audio["array"])
batch["labels"] = preprocess_labels(labels)
return batch
with training_args.main_process_first():
tmp_datasets = raw_datasets.map(
prepare_dataset,
num_proc=data_args.preprocessing_num_workers,
input_columns=["audio", "labels"],
remove_columns=[col for col in next(iter(raw_datasets.values())).column_names if col != "labels"], # this is a trick to avoid to rewrite the entire audio column which takes ages
desc="Computing audio length",
)
for split in raw_datasets:
raw_datasets[split] = concatenate_datasets([raw_datasets[split].remove_columns(["labels"]), tmp_datasets[split]], axis=1)
# filter training data with inputs < min_input_length
min_input_length = data_args.min_length_seconds * sampling_rate
if not dataset_was_precomputed:
def is_audio_valid(input_length):
return input_length > min_input_length
with training_args.main_process_first():
raw_datasets = raw_datasets.filter(
is_audio_valid,
input_columns=["length"],
num_proc=data_args.preprocessing_num_workers,
desc="Filtering by audio length",
)
# filter training data with non-valid labels
def is_label_valid(label):
return label != "Unknown"
with training_args.main_process_first():
raw_datasets = raw_datasets.filter(
is_label_valid,
input_columns=["labels"],
num_proc=data_args.preprocessing_num_workers,
desc="Filtering by labels",
)
if training_args.do_train and data_args.max_samples_per_label:
label_names = set(raw_datasets["train"]["labels"])
labels = np.array(raw_datasets["train"]["labels"])
indices = np.arange(len(labels))
indices_to_keep = []
set_seed(training_args.seed)
for label in label_names:
label_indices = indices[labels==label]
label_indices = np.random.choice(label_indices, size=min(data_args.max_samples_per_label, len(label_indices)), replace=False)
indices_to_keep.extend(np.random.choice(label_indices, size=min(data_args.max_samples_per_label, len(label_indices)), replace=False))
with training_args.main_process_first():
raw_datasets["train"] = raw_datasets["train"].select(indices_to_keep).shuffle(seed=training_args.seed)
# Print a summary of the labels to the stddout (helps identify low-label classes that could be filtered)
# sort by freq
count_labels_dict = Counter(raw_datasets["train"]["labels"])
count_labels_dict = sorted(count_labels_dict.items(), key=lambda item: (-item[1], item[0]))
labels, frequencies = zip(*count_labels_dict)
total_labels = sum(frequencies)
labels_to_remove = []
logger.info(f"{'Accent':<15} {'Perc.':<5}")
logger.info("-" * 20)
for lab, freq in zip(labels, frequencies):
freq = 100 * freq / total_labels
logger.info(f"{lab:<15} {freq:<5}")
if freq < data_args.filter_threshold:
labels_to_remove.append(lab)
if len(labels_to_remove):
# filter training data with label freq below threshold
def is_label_valid(label):
return label not in labels_to_remove
logger.info(f"Labels removed: {labels_to_remove}")
with training_args.main_process_first():
raw_datasets = raw_datasets.filter(
is_label_valid,
input_columns=["labels"],
num_proc=data_args.preprocessing_num_workers,
desc="Filtering low freq labels",
)
# We'll include these in the model's config to get human readable labels in the Inference API.
set_labels = set(raw_datasets["train"]["labels"])
if training_args.do_eval:
logger.info(f"The following accents are present in the eval set but not the test set: {set(raw_datasets['eval']['labels']) - set_labels}")
set_labels = set_labels.union(set(raw_datasets["eval"]["labels"]))
label2id, id2label = {}, {}
for i, label in enumerate(sorted(list(set_labels))):
label2id[label] = str(i)
id2label[str(i)] = label
with training_args.main_process_first():
raw_datasets = raw_datasets.map(
lambda label: {"labels_id": int(label2id[label])},
input_columns=["labels"],
num_proc=data_args.preprocessing_num_workers,
desc="Apply labels id",
)
else:
label2id, id2label = {}, {}
# TODO: slow - probably not best
if training_args.do_train:
for sample in raw_datasets["train"]:
candidate = label2id.get(sample["labels"])
if candidate is not None and candidate != sample["labels_id"]:
print(f"issue {candidate} should be {sample['labels_id']}")
label2id[sample["labels"]] = sample["labels_id"]
if training_args.do_eval:
for sample in raw_datasets["eval"]:
candidate = label2id.get(sample["labels"])
if candidate is not None and candidate != sample["labels_id"]:
print(f"issue {candidate} should be {sample['labels_id']}")
label2id[sample["labels"]] = sample["labels_id"]
# if training_args.do_train:
# label2id = dict(zip(raw_datasets["train"]["labels"], raw_datasets["train"]["labels_id"]))
# if training_args.do_eval:
# label2id.update(dict(zip(raw_datasets["eval"]["labels"], raw_datasets["eval"]["labels_id"])))
id2label = {str(val): key for (key, val) in label2id.items()}
# Load the accuracy metric from the datasets package
metric = evaluate.load("accuracy", cache_dir=model_args.cache_dir)
# Define our compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with
# `predictions` and `label_ids` fields) and has to return a dictionary string to float.
def compute_metrics(eval_pred):
"""Computes accuracy on a batch of predictions"""
predictions = np.argmax(eval_pred.predictions, axis=1)
return metric.compute(predictions=predictions, references=eval_pred.label_ids)
config = AutoConfig.from_pretrained(
model_args.config_name or model_args.model_name_or_path,
num_labels=len(label2id),
label2id=label2id,
id2label=id2label,
finetuning_task="audio-classification",
cache_dir=model_args.cache_dir,
revision=model_args.model_revision,
token=model_args.token,
trust_remote_code=model_args.trust_remote_code,
)
# adapt config with regularization
config.update(
{
"feat_proj_dropout": model_args.feat_proj_dropout,
"attention_dropout": model_args.attention_dropout,
"hidden_dropout": model_args.hidden_dropout,
"final_dropout": model_args.final_dropout,
"mask_time_prob": model_args.mask_time_prob,
"mask_time_length": model_args.mask_time_length,
"mask_feature_prob": model_args.mask_feature_prob,
"mask_feature_length": model_args.mask_feature_length,
"layerdrop": model_args.layerdrop,
"activation_dropout": model_args.activation_dropout,
"use_weighted_layer_sum": model_args.use_weighted_layer_sum,
}
)
if model_args.use_last_embedding_layer:
model = SequenceClassificationModel.from_pretrained(
model_args.model_name_or_path,
from_tf=bool(".ckpt" in model_args.model_name_or_path),
config=config,
cache_dir=model_args.cache_dir,
revision=model_args.model_revision,
token=model_args.token,
trust_remote_code=model_args.trust_remote_code,
ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
)
if model_args.freeze_base_model:
model.compute_w2v2 = False
else:
model = AutoModelForAudioClassification.from_pretrained(
model_args.model_name_or_path,
from_tf=bool(".ckpt" in model_args.model_name_or_path),
config=config,
cache_dir=model_args.cache_dir,
revision=model_args.model_revision,
token=model_args.token,
trust_remote_code=model_args.trust_remote_code,
ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
)
# freeze the convolutional waveform encoder for wav2vec2-style models
if model_args.freeze_feature_encoder:
if hasattr(model, "freeze_feature_encoder"):
model.freeze_feature_encoder()
else:
raise ValueError("Method for freezing the feature encoder is not defined for Whisper-style models.")
if model_args.freeze_base_model:
if hasattr(model, "freeze_base_model"):
# wav2vec2-style models
model.freeze_base_model()
if hasattr(model, "freeze_feature_encoder"):
model.freeze_feature_encoder()
elif hasattr(model, "freeze_encoder"):
# whisper-style models
model.freeze_encoder()
else:
raise ValueError("Method for freezing the base module of the audio encoder is not defined")
if model_args.freeze_base_model and model_args.use_last_embedding_layer:
if not dataset_was_precomputed:
# precomputing hidden states
from torch.utils.data import DataLoader
from accelerate import Accelerator
_HIDDEN_STATES_START_POSITION = 2
if training_args.fp16:
mixed_precision = "fp16"
elif training_args.bf16:
mixed_precision = "bf16"
else:
mixed_precision = "no"
accelerator = Accelerator(
gradient_accumulation_steps=training_args.gradient_accumulation_steps,
mixed_precision=mixed_precision,
project_dir=training_args.output_dir,
)
audio_data_collator = DataCollatorFeatureExtractorWithPadding(feature_extractor, max_length_seconds=data_args.max_length_seconds, feature_extractor_input_name=feature_extractor_input_name)
for split in raw_datasets:
data_loader = DataLoader(
raw_datasets[split],
batch_size=training_args.per_device_train_batch_size, # TODO: chose another one
collate_fn=audio_data_collator,
num_workers=training_args.dataloader_num_workers,
pin_memory=True,
)
data_loader = accelerator.prepare(data_loader)
all_encoder_outputs = []
for batch in tqdm(data_loader, disable=not accelerator.is_local_main_process):
model.wav2vec2.to(batch[feature_extractor_input_name].device)
with torch.no_grad():
encoder_outputs = model.wav2vec2(batch[feature_extractor_input_name], attention_mask=batch.get("attention_mask", None), output_hidden_states=True)
hidden_states = encoder_outputs[_HIDDEN_STATES_START_POSITION][-1]
if batch.get("attention_mask", None) is None:
hidden_states = hidden_states.mean(dim=1)
else:
padding_mask = model._get_feature_vector_attention_mask(hidden_states.shape[1], batch.get("attention_mask", None))
hidden_states[~padding_mask] = 0.0
hidden_states = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1)
encoder_outputs = accelerator.gather_for_metrics(hidden_states)
# TODO: check it works multi device
if accelerator.is_main_process:
all_encoder_outputs.extend(encoder_outputs.to("cpu").numpy())
if accelerator.is_main_process:
tmp_hidden_states = Dataset.from_dict({"hidden_states": all_encoder_outputs})
tmp_hidden_states.save_to_disk(os.path.join(data_args.temporary_save_to_disk, split))
accelerator.wait_for_everyone()
del all_encoder_outputs
tmp_hidden_states = datasets.load_from_disk(os.path.join(data_args.temporary_save_to_disk, split))
with accelerator.main_process_first():
raw_datasets[split] = concatenate_datasets([raw_datasets[split].remove_columns("audio"), tmp_hidden_states], axis=1)
accelerator.free_memory()
del data_loader, batch, accelerator
data_collator = DataCollatorHiddenStatesPadding()
else:
data_collator = DataCollatorFeatureExtractorWithPadding(feature_extractor, max_length_seconds=data_args.max_length_seconds)
if data_args.save_to_disk is not None and not dataset_was_precomputed:
raw_datasets.save_to_disk(data_args.save_to_disk)
logger.info(f"Dataset saved at {data_args.save_to_disk}. Be careful of changing data parameters, which won't change the current saved dataset if reloaded.")
# Initialize our trainer
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=raw_datasets["train"] if training_args.do_train else None,
eval_dataset=raw_datasets["eval"] if training_args.do_eval else None,
compute_metrics=compute_metrics,
tokenizer=feature_extractor,
)
ignore_keys = ["hidden_states","attentions"] if model_args.use_weighted_layer_sum else None
# Training
if training_args.do_train:
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, ignore_keys_for_eval=ignore_keys)
trainer.save_model()
trainer.log_metrics("train", train_result.metrics)
trainer.save_metrics("train", train_result.metrics)
trainer.save_state()
# Evaluation
if training_args.do_eval:
metrics = trainer.evaluate(ignore_keys=ignore_keys)
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)
# Write model card and (optionally) push to hub
kwargs = {
"finetuned_from": model_args.model_name_or_path,
"tasks": "audio-classification",
"dataset": data_args.train_dataset_name.split("+")[0],
"tags": ["audio-classification"],
}
if training_args.push_to_hub:
trainer.push_to_hub(**kwargs)
else:
trainer.create_model_card(**kwargs)
if __name__ == "__main__":
main()
run_dataset_concatenation.py deleted 100644 → 0
View file @ 64cfa64e
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
from datasets import Audio, concatenate_datasets, load_dataset
from huggingface_hub import get_full_repo_name
from transformers import HfArgumentParser, WhisperTokenizerFast
@dataclass
class DataTrainingArguments:
"""
Arguments pertaining to what data we are going to input our model for training and eval.
"""
dataset_name: str = field(
default=None,
metadata={"help": "The name of the dataset to use (via the datasets library)."},
)
dataset_config_name: str = field(
default=None,
metadata={"help": "The configuration name of the dataset to use (via the datasets library)."},
)
dataset_split_name: str = field(
default=None,
metadata={
"help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
},
)
label_column_name: str = field(
default="labels",
metadata={"help": "The name of the dataset column containing the labels in the dataset. Defaults to 'label'"},
)
text_column_name: str = field(
default="text",
metadata={
"help": "The name of the dataset column containing the text transcriptions in the dataset. Defaults to 'text'"
},
)
speaker_column_name: str = field(
default="speaker_id",
metadata={
"help": "The name of the dataset column containing the speaker ids in the dataset. Defaults to 'speaker_id'"
},
)
dataset_cache_dir: str = field(
default=None,
metadata={"help": "Path to cache directory for saving and loading datasets"},
)
preprocessing_num_workers: int = field(
default=None,
metadata={"help": "The number of processes to use for the preprocessing."},
)
batch_size: int = field(
default=500,
metadata={"help": "Number of examples per batch provided to the preprocessing function."},
)
download_only: bool = field(
default=False,
metadata={"help": "Whether to only do data download and skip pre-processing."},
)
audio_column_name: str = field(
default="audio",
metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
)
max_duration_in_seconds: float = field(
default=20.0,
metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
)
sampling_rate: int = field(
default=16_000,
metadata={
"help": "Sampling rate at which to resample the audio data. Should be set to the same sampling rate as the target model."
},
)
max_samples: int = field(
default=None,
metadata={
"help": "For debugging purposes, truncate the number of examples in the dataset to this value if set."
},
)
output_dir: str = field(
default=None,
metadata={
"help": "Where to save the processed dataset to disk. If unspecified, uses a 'pretty' version of the "
"original dataset name. E.g. 'facebook/voxpopuli' will be saved under 'voxpopuli'."
},
)
push_to_hub: bool = field(
default=False,
metadata={"help": "Whether or not to push the processed dataset to the Hub."},
)
seed: int = field(
default=0,
metadata={"help": "RNG seed for reproducibility. Used during the final shuffling of the combined dataset."},
)
def convert_dataset_str_to_list(
dataset_names,
dataset_config_names,
splits=None,
label_column_names=None,
text_column_names=None,
speaker_column_names=None,
dataset_samples=None,
default_split="train",
):
if isinstance(dataset_names, str):
dataset_names = dataset_names.split("+")
dataset_config_names = dataset_config_names.split("+")
splits = splits.split("+") if splits is not None else None
label_column_names = label_column_names.split("+") if label_column_names is not None else None
text_column_names = text_column_names.split("+") if text_column_names is not None else None
speaker_column_names = speaker_column_names.split("+") if speaker_column_names is not None else None
dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
# basic checks to ensure we've got the right number of datasets/configs/splits/columns/probs
if len(dataset_names) != len(dataset_config_names):
raise ValueError(
f"Ensure one config is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(dataset_config_names)} configs."
)
if splits is not None and len(splits) != len(dataset_names):
raise ValueError(
f"Ensure one split is passed for each dataset, got {len(dataset_names)} datasets and {len(splits)} splits."
)
if label_column_names is not None and len(label_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one label column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(label_column_names)} label column names."
)
if text_column_names is not None and len(text_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one text column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(text_column_names)} text column names."
)
if speaker_column_names is not None and len(speaker_column_names) != len(dataset_names):
raise ValueError(
f"Ensure one text column name is passed for each dataset, got {len(dataset_names)} datasets and"
f" {len(speaker_column_names)} speaker column names."
)
if dataset_samples is not None:
if len(dataset_samples) != len(dataset_names):
raise ValueError(
f"Ensure one sample is passed for each dataset, got {len(dataset_names)} datasets and "
f"{len(dataset_samples)} samples."
)
dataset_samples = [float(ds_sample) for ds_sample in dataset_samples]
else:
dataset_samples = [None] * len(dataset_names)
label_column_names = (
label_column_names if label_column_names is not None else ["labels" for _ in range(len(dataset_names))]
)
text_column_names = (
text_column_names if text_column_names is not None else ["text" for _ in range(len(dataset_names))]
)
speaker_column_names = (
speaker_column_names if speaker_column_names is not None else ["speaker_id" for _ in range(len(dataset_names))]
)
splits = splits if splits is not None else [default_split for _ in range(len(dataset_names))]
dataset_names_dict = []
for i, ds_name in enumerate(dataset_names):
dataset_names_dict.append(
{
"name": ds_name,
"config": dataset_config_names[i],
"split": splits[i],
"label_column_name": label_column_names[i],
"text_column_name": text_column_names[i],
"speaker_column_name": speaker_column_names[i],
"samples": dataset_samples[i],
}
)
return dataset_names_dict
def main():
# 1. Parse input arguments
parser = HfArgumentParser(DataTrainingArguments)
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0]
else:
data_args = parser.parse_args_into_dataclasses()[0]
dataset_names_dict = convert_dataset_str_to_list(
data_args.dataset_name,
data_args.dataset_config_name,
splits=data_args.dataset_split_name,
label_column_names=data_args.label_column_name,
text_column_names=data_args.text_column_name,
speaker_column_names=data_args.speaker_column_name,
)
# load whisper tokenizer for normalisation
sampling_rate = data_args.sampling_rate
tokenizer = WhisperTokenizerFast.from_pretrained("openai/whisper-tiny.en")
max_input_length = int(data_args.max_duration_in_seconds * sampling_rate)
batch_size = data_args.batch_size
preprocessing_num_workers = data_args.preprocessing_num_workers
all_vectorized_datasets = []
for dataset_dict in dataset_names_dict:
print(10 * "=", dataset_dict["name"], 10 * "=")
raw_datasets = load_dataset(
dataset_dict["name"],
dataset_dict["config"],
split=dataset_dict["split"],
cache_dir=data_args.dataset_cache_dir,
num_proc=data_args.preprocessing_num_workers,
)
if data_args.download_only:
continue
features = raw_datasets.column_names
if dataset_dict["label_column_name"] not in features:
raise ValueError(
f"--label_column_name {dataset_dict['label_column_name']} not found in dataset '{dataset_dict['name']}'. "
"Make sure to set `--label_column_name` to the correct text column - one of "
f"{', '.join(features)}."
)
elif dataset_dict["label_column_name"] != "labels":
raw_datasets = raw_datasets.rename_column(dataset_dict["label_column_name"], "labels")
if dataset_dict["text_column_name"] not in features:
raise ValueError(
f"--text_column_name {dataset_dict['text_column_name']} not found in dataset '{dataset_dict['name']}'. "
"Make sure to set `--text_column_name` to the correct text column - one of "
f"{', '.join(features)}."
)
elif dataset_dict["text_column_name"] != "text":
raw_datasets = raw_datasets.rename_column(dataset_dict["text_column_name"], "text")
if dataset_dict["speaker_column_name"] not in features:
raise ValueError(
f"--speaker_column_name {dataset_dict['speaker_column_name']} not found in dataset '{dataset_dict['name']}'. "
"Make sure to set `--speaker_column_name` to the correct speaker id column - one of "
f"{', '.join(features)}."
)
elif dataset_dict["speaker_column_name"] != "speaker_id":
raw_datasets = raw_datasets.rename_column(dataset_dict["speaker_column_name"], "speaker_id")
raw_datasets = raw_datasets.remove_columns(
set(raw_datasets.features.keys()) - {"audio", "labels", "text", "speaker_id"}
)
if data_args.max_samples is not None:
raw_datasets = raw_datasets.select(range(data_args.max_samples))
raw_datasets = raw_datasets.cast_column(data_args.audio_column_name, Audio(sampling_rate=sampling_rate))
raw_datasets = raw_datasets.sort("speaker_id")
def filter_transcriptions(text):
normalized_text = tokenizer.normalize(text).strip()
return bool(normalized_text) and text.lower() != "ignore_time_segment_in_scoring"
raw_datasets = raw_datasets.filter(
filter_transcriptions, input_columns=["text"], desc="Filtering non-speech transcriptions"
)
def prepare_dataset(batch):
audio = [sample["array"] for sample in batch["audio"]]
input_lengths = [len(sample) for sample in audio]
concatenated_audio = []
concatenated_text = []
concatenated_speaker = []
concatenated_labels = []
audio_sample = audio[0]
text_sample = batch["text"][0]
label_sample = batch["labels"][0]
for idx in range(1, len(audio)):
prev_speaker = batch["speaker_id"][idx - 1]
speaker = batch["speaker_id"][idx]
if len(audio_sample) + input_lengths[idx] < max_input_length:
if speaker == prev_speaker:
# we have no information about whether the segments follow on sequentially
# so we just ensure the same speaker as we concatenate across files
audio_sample = np.append(audio_sample, audio[idx])
# extra spaces in the text transcription don't matter, since we only use it for the WER computation
text_sample += " " + batch["text"][idx]
else:
# segments do not follow sequentially, save the audio and start looping again
concatenated_audio.append(audio_sample)
concatenated_text.append(text_sample)
concatenated_labels.append(label_sample)
concatenated_speaker.append(speaker)
audio_sample = audio[idx]
text_sample = batch["text"][idx]
label_sample = batch["labels"][idx]
else:
# concatenated audio exceeds max length, save the audio and start looping again
concatenated_audio.append(audio_sample)
concatenated_text.append(text_sample)
concatenated_labels.append(label_sample)
concatenated_speaker.append(speaker)
audio_sample = audio[idx]
text_sample = batch["text"][idx]
label_sample = batch["labels"][idx]
batch["audio"] = [{"array": array, "sampling_rate": sampling_rate} for array in concatenated_audio]
batch["text"] = concatenated_text
batch["labels"] = concatenated_labels
batch["speaker_id"] = concatenated_speaker
return batch
raw_datasets = raw_datasets.map(
prepare_dataset,
batched=True,
batch_size=batch_size,
num_proc=preprocessing_num_workers,
desc="Concatenating dataset...",
)
pretty_name = dataset_dict["name"].split("/")[-1]
def postprocess_ids(speaker_id, idx):
formatted_idx = f"{pretty_name}-{speaker_id}-{idx}"
return {"id": formatted_idx}
raw_datasets = raw_datasets.map(
postprocess_ids,
input_columns=["speaker_id"],
with_indices=True,
desc="Setting sample idxs...",
num_proc=preprocessing_num_workers,
)
print(f"Final length {pretty_name}: ", len(raw_datasets))
# Re-format transcriptions and condition on prev as numpy arrays
raw_datasets = raw_datasets.with_format("np")
all_vectorized_datasets.append(raw_datasets)
all_vectorized_datasets = concatenate_datasets(all_vectorized_datasets)
dataset_features = all_vectorized_datasets.features.copy()
dataset_features["audio"] = Audio(sampling_rate=sampling_rate)
all_vectorized_datasets = all_vectorized_datasets.cast(
dataset_features, batch_size=batch_size, writer_batch_size=batch_size, num_proc=preprocessing_num_workers
)
all_vectorized_datasets = all_vectorized_datasets.shuffle(seed=data_args.seed)
all_vectorized_datasets.save_to_disk(data_args.output_dir)
repo_name = get_full_repo_name(Path(data_args.output_dir).absolute().name)
if data_args.push_to_hub:
all_vectorized_datasets.push_to_hub(repo_name, config_name="train", max_shard_size="1GB")
if __name__ == "__main__":
main()
run_prompt_creation.py deleted 100644 → 0
View file @ 64cfa64e
import logging
import os
import shutil
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union
import torch
from accelerate import Accelerator
from accelerate.logging import get_logger
from datasets import DatasetDict, load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
HfArgumentParser,
)
logger = get_logger(__name__, log_level="INFO")
@dataclass
class ModelArguments:
"""
Arguments pertaining to what data we are going to input our model for training and eval.
"""
model_name_or_path: str = field(
metadata={"help": "The name of the model to use (via the transformers library) for the prompt annotation."},
)
per_device_eval_batch_size: int = field(
metadata={"help": "The per-device batch size to use for inference."},
)
model_variant: str = field(
default=None,
metadata={"help": "If specified load weights from `variant` filename, *e.g.* pytorch_model.<variant>.bin. "},
)
model_revision: str = field(
default="main",
metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
)
cache_dir: Optional[str] = field(
default=None,
metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
)
torch_dtype: Optional[str] = field(
default="float16",
metadata={
"help": (
"Floating-point format in which the model weights should be initialized"
" and the computations run. Choose one of `[float32, float16, bfloat16]`."
)
},
)
attn_implementation: Optional[str] = field(
default="sdpa",
metadata={"help": "Which attn type to use: ['eager', 'sdpa', 'flash_attention_2']"},
)
load_in_8bit: Optional[bool] = field(
default=False, metadata={"help": "Whether to use 8-bit precision for inference."}
)
load_in_4bit: Optional[bool] = field(
default=False, metadata={"help": "Whether to use 4-bit precision for inference."}
)
bnb_4bit_quant_type: Optional[str] = field(
default="nf4", metadata={"help": "precise the quantization type (fp4 or nf4)"}
)
use_bnb_nested_quant: Optional[bool] = field(default=False, metadata={"help": "use nested quantization"})
trust_remote_code: Optional[bool] = field(
default=False,
metadata={
"help": (
"Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
"should only be set to `True` for repositories you trust and in which you have read the code, as it will "
"execute code present on the Hub on your local machine."
)
},
)
use_fast_tokenizer: Optional[bool] = field(
default=True, metadata={"help": "Use fast tokenizer for encoding/decoding input ids"}
)
token: Optional[bool] = field(
default=True,
metadata={
"help": "Whether or not to use an authentication token when loading/uploading from the Hugging Face Hub"
},
)
do_sample: Optional[bool] = field(default=True, metadata={"help": "Whether to use sampling mode for generation"})
temperature: Optional[float] = field(default=0.6, metadata={"help": "Temperature for sampling-based generation"})
max_new_tokens: Optional[int] = field(
default=256, metadata={"help": "Maximum number of new tokens during generation"}
)
compile_generate: Optional[bool] = field(
default=False, metadata={"help": "Whether to compile the forward pass (not sampling) in generate."}
)
@dataclass
class DataArguments:
"""
Arguments pertaining to what data we are going to input our model for training and eval.
"""
output_dir: str = field(
metadata={
"help": "Where to save the processed dataset to disk. If unspecified, uses a 'pretty' version of the "
"original dataset name. E.g. 'facebook/voxpopuli' will be saved under 'voxpopuli'."
},
)
dataset_name: str = field(
default=None,
metadata={"help": "The name of the dataset to use (via the datasets library)"},
)
dataset_config_name: Optional[str] = field(
default=None,
metadata={"help": "The configuration name of the dataset to use (via the datasets library)."},
)
dataset_split_name: Optional[str] = field(
default=None,
metadata={"help": "The split name of the dataset to use (via the datasets library)."},
)
dataset_cache_dir: Optional[str] = field(
default=None,
metadata={"help": "Path to cache directory for saving and loading datasets"},
)
max_eval_samples: Optional[int] = field(
default=None,
metadata={"help": "Maximum number of samples for generation - use for debugging purposes."},
)
overwrite_cache: bool = field(
default=False,
metadata={"help": "Overwrite the cached training and evaluation sets"},
)
preprocessing_num_workers: Optional[int] = field(
default=None,
metadata={"help": "The number of processes to use for the preprocessing."},
)
dataloader_num_workers: Optional[int] = field(
default=0,
metadata={"help": "The number of processes to use for the dataloader."},
)
push_to_hub: Optional[bool] = field(
default=False,
metadata={"help": "Whether or not to push the processed dataset to the Hub."},
)
hub_dataset_id: Optional[str] = field(
default=None,
metadata={"help": "Repository namespace if pushing to the Hugging Face Hub."},
)
overwrite_output_dir: Optional[bool] = field(
default=False,
metadata={"help": "Overwrite the content of the output directory each time the script is run."},
)
def __post_init__(self):
if self.push_to_hub and self.hub_dataset_id is None:
raise ValueError("You must specify the `hub_dataset_id` when setting `--push_to_hub=True`")
def get_quantization_config(model_args: ModelArguments) -> Union[BitsAndBytesConfig, None]:
if model_args.load_in_4bit:
compute_dtype = torch.float16
if model_args.torch_dtype not in {"auto", None}:
compute_dtype = getattr(torch, model_args.torch_dtype)
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=compute_dtype,
bnb_4bit_quant_type=model_args.bnb_4bit_quant_type,
bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant,
)
elif model_args.load_in_8bit:
quantization_config = BitsAndBytesConfig(
load_in_8bit=True,
)
else:
quantization_config = None
return quantization_config
def get_current_device() -> int:
"""Get the current device. For GPU we return the local process index to enable multiple GPU training."""
return Accelerator().local_process_index if torch.cuda.is_available() else "cpu"
def get_kbit_device_map() -> Union[Dict[str, int], None]:
"""Useful for running inference with quantized models by setting `device_map=get_peft_device_map()`"""
return {"": get_current_device()} if torch.cuda.is_available() else None
@dataclass
class DataCollatorWithPadding:
"""
Data collator that will dynamically pad the inputs received to the longest sequence in the batch.
"""
tokenizer: Any
def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
# split inputs and labels since they have to be of different lengths and need
# different padding methods
input_ids = {"input_ids": [feature["input_ids"] for feature in features]}
batch = self.tokenizer.pad(input_ids, return_tensors="pt", padding="longest", return_attention_mask=True)
return batch
# TODO(SG): add accent keyword
PROMPT = """You will be given six descriptive keywords related to an audio sample of a person's speech. These keywords include:
1. The gender (e.g., male, female)
2. The level of reverberation (e.g., very roomy sounding, quite roomy sounding, slightly roomy sounding, moderate reverberation, slightly confined sounding, quite confined sounding, very confined sounding)
3. The amount of noise the sample (e.g., very noisy, quite noisy, slightly noisy, moderate ambient sound, slightly clear, quite clear, very clear)
4. The tone of the speaker's voice (e.g., very monotone, quite monotone, slightly monotone, moderate intonation, slightly expressive, quite expressive, very expressive)
5. The pace of the speaker's delivery (e.g., very slowly, quite slowly, slightly slowly, moderate speed, slightly fast, quite fast, very fast)
6. The pitch of the speaker's voice (e.g., very low pitch, quite low pitch, slightly low pitch, moderate pitch, slightly high pitch, quite high pitch, very high pitch)
Your task is to create a text description using these keywords that accurately describes the speech sample while ensuring the description remains grammatically correct and easy to understand. You should rearrange the keyword order as necessary, and substitute synonymous terms where appropriate. If the amount of noise is 'very noisy' and the level of reverberation is 'very roomy sounding', include terms like 'very bad recording' in the description. Likewise, if the amount of noise is 'very clear' and the level of reverberation is 'very confined sounding', include terms like 'very good recording' in the description. Otherwise, do not add extra details beyond what has been provided, and only return the generated description.
For example, given the following keywords: 'female', 'slightly roomy sounding', 'slightly noisy', 'very expressive', 'slightly low pitch', 'very slowly', a valid description would be: 'a woman with a deep voice speaks slowly but has an animated delivery in an echoey room with some background noise'.
For the keywords: '[gender]', '[reverberation]', '[noise]', '[speech_monotony]', '[pitch]', '[speaking_rate]', the corresponding description is:"
"""
def main():
# 1. Parse input arguments
parser = HfArgumentParser((ModelArguments, DataArguments))
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
model_args, data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
else:
model_args, data_args = parser.parse_args_into_dataclasses()
# 2. Setup logging
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
handlers=[logging.StreamHandler(sys.stdout)],
)
accelerator = Accelerator()
if data_args.overwrite_output_dir and os.path.exists(data_args.output_dir) and os.path.isdir(data_args.output_dir):
logger.info("Cleaning output dir from previous run...")
shutil.rmtree(data_args.output_dir)
# 3. Load annotated dataset
logger.info("*** Load annotated dataset ***")
if data_args.dataset_split_name is not None:
raw_datasets = DatasetDict()
data_splits = data_args.dataset_split_name.split("+")
# load on a split-wise basis
for split in data_splits:
raw_datasets[split] = load_dataset(
data_args.dataset_name,
data_args.dataset_config_name,
split=split,
cache_dir=model_args.cache_dir,
token=model_args.token,
num_proc=data_args.preprocessing_num_workers,
)
else:
# load all splits for annotation
raw_datasets = load_dataset(
data_args.dataset_name,
data_args.dataset_config_name,
cache_dir=model_args.cache_dir,
token=model_args.token,
num_proc=data_args.preprocessing_num_workers,
)
raw_datasets_features = set(raw_datasets[next(iter(raw_datasets))].features.keys())
if data_args.max_eval_samples is not None:
for split in raw_datasets:
raw_datasets[split] = raw_datasets[split].select(range(data_args.max_eval_samples))
# TODO(SG): add accent
EXPECTED_COLUMNS = {"gender", "pitch", "noise", "reverberation", "speech_monotony", "speaking_rate"}
if not EXPECTED_COLUMNS.issubset(raw_datasets_features):
missing_columns = EXPECTED_COLUMNS - raw_datasets_features
raise ValueError(
f"Missing columns {missing_columns} from the dataset features. Got dataset features {raw_datasets_features}"
)
# 4. Load pre-trained model
logger.info("*** Load pretrained model ***")
torch_dtype = (
model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
)
quantization_config = get_quantization_config(model_args)
model = AutoModelForCausalLM.from_pretrained(
model_args.model_name_or_path,
revision=model_args.model_revision,
variant=model_args.model_variant,
trust_remote_code=model_args.trust_remote_code,
attn_implementation=model_args.attn_implementation,
torch_dtype=torch_dtype,
device_map=get_kbit_device_map() if quantization_config is not None else None,
quantization_config=quantization_config,
low_cpu_mem_usage=True,
token=model_args.token,
).eval()
if model_args.compile_generate:
if not callable(getattr(model, "_setup_cache", None)):
raise ValueError(
f"Static k/v cache is not compatible with the model {model.__class__.__name__}. Set `--compile_generate=False"
"for dynamic k/v cache"
)
model.generation_config.cache_implementation = "static"
model._forward = model.forward
compiled_forward = torch.compile(model.forward)
def compiled(func, input_ids, **kwargs):
return func(input_ids, **kwargs)
def call(input_ids, **kwargs):
if input_ids.shape[-1] == 1:
return compiled(compiled_forward, input_ids, **kwargs)
return model._forward(input_ids, **kwargs)
model.forward = call
tokenizer = AutoTokenizer.from_pretrained(
model_args.model_name_or_path,
revision=model_args.model_revision,
trust_remote_code=model_args.trust_remote_code,
use_fast=model_args.use_fast_tokenizer,
padding_side="left",
)
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.bos_token_id
model.generation_config.pad_token_id = model.generation_config.eos_token_id
def prepare_dataset(sample):
sample_prompt = PROMPT
for key in EXPECTED_COLUMNS:
sample_prompt = sample_prompt.replace(f"[{key}]", sample[key])
sample_prompt = [{"role": "user", "content": sample_prompt}]
token_ids = tokenizer.apply_chat_template(sample_prompt)
sample["input_ids"] = token_ids
return sample
with accelerator.main_process_first():
vectorized_datasets = raw_datasets.map(
prepare_dataset, num_proc=data_args.preprocessing_num_workers, desc="Preparing prompts"
)
# Prepare everything with our `accelerator`
model = accelerator.prepare(model)
data_collator = DataCollatorWithPadding(tokenizer)
def generate_step(batch):
output_ids = accelerator.unwrap_model(model).generate(
batch["input_ids"],
attention_mask=batch["attention_mask"],
do_sample=model_args.do_sample,
temperature=model_args.temperature,
max_new_tokens=model_args.max_new_tokens,
)
output_ids = accelerator.pad_across_processes(output_ids, dim=1, pad_index=tokenizer.pad_token_id)
return output_ids
def postprocess_dataset(sample):
prompt_text = tokenizer.decode(sample["input_ids"], skip_special_tokens=True)
generated_text = tokenizer.decode(sample["generated_ids"], skip_special_tokens=True)
sample["text_description"] = generated_text[len(prompt_text) :]
return sample
for split in vectorized_datasets:
data_loader = DataLoader(
vectorized_datasets[split],
batch_size=model_args.per_device_eval_batch_size,
collate_fn=data_collator,
num_workers=data_args.dataloader_num_workers,
pin_memory=True,
)
data_loader = accelerator.prepare(data_loader)
all_generated_ids = []
for batch in tqdm(data_loader, disable=not accelerator.is_local_main_process):
generated_ids = generate_step(batch)
generated_ids = accelerator.gather_for_metrics(generated_ids)
all_generated_ids.extend(generated_ids.cpu().numpy())
vectorized_datasets[split] = vectorized_datasets[split].add_column("generated_ids", all_generated_ids)
if accelerator.is_main_process:
vectorized_datasets[split] = vectorized_datasets[split].map(
postprocess_dataset,
num_proc=data_args.preprocessing_num_workers,
desc="Postprocessing dataset",
remove_columns=["input_ids", "generated_ids"],
)
if accelerator.is_main_process:
vectorized_datasets.save_to_disk(data_args.output_dir)
if data_args.push_to_hub:
vectorized_datasets.push_to_hub(
data_args.hub_dataset_id,
config_name=data_args.dataset_config_name if data_args.dataset_config_name is not None else "default",
token=model_args.token,
)
accelerator.end_training()
if __name__ == "__main__":
main()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment