Upload main.py

main.py

@@ -0,0 +1,229 @@
+import os
+import re
+import json
+import torch
+import argparse
+from functools import partial
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import pandas as pd
+
+from datasets import set_caching_enabled
+set_caching_enabled(False)
+
+from datasets import (
+    load_dataset,
+    load_from_disk,
+    load_metric,)
+
+from transformers import (
+    Wav2Vec2CTCTokenizer,
+    Wav2Vec2FeatureExtractor,
+    Wav2Vec2Processor,
+    Wav2Vec2ForCTC,
+    TrainingArguments,
+    Trainer,
+)
+
+import torchaudio
+
+
+def preprocess_data(example, tok_func = word_tokenize):
+    example['sentence'] = ' '.join(tok_func(example['sentence']))
+    return example
+
+
+def speech_file_to_array_fn(batch,
+                            text_col="sentence",
+                            fname_col="path",
+                            resampling_to=16000):
+    speech_array, sampling_rate = torchaudio.load(batch[fname_col])
+    resampler=torchaudio.transforms.Resample(sampling_rate, resampling_to)
+    batch["speech"] = resampler(speech_array)[0].numpy()
+    batch["sampling_rate"] = resampling_to
+    batch["target_text"] = batch[text_col]
+    return
+
+@dataclass
+class DataCollatorCTCWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        processor (:class:`~transformers.Wav2Vec2Processor`)
+            The processor used for proccessing the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        max_length (:obj:`int`, `optional`):
+            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).
+        max_length_labels (:obj:`int`, `optional`):
+            Maximum length of the ``labels`` returned list and optionally padding length (see above).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+    """
+
+    processor: Wav2Vec2Processor
+    padding: Union[bool, str] = True
+    max_length: Optional[int] = None
+    max_length_labels: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    pad_to_multiple_of_labels: Optional[int] = None
+
+    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 lenghts and need
+        # different padding methods
+        input_features = [{"input_values": feature["input_values"]} for feature in features]
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+
+        batch = self.processor.pad(
+            input_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        with self.processor.as_target_processor():
+            labels_batch = self.processor.pad(
+                label_features,
+                padding=self.padding,
+                max_length=self.max_length_labels,
+                pad_to_multiple_of=self.pad_to_multiple_of_labels,
+                return_tensors="pt",
+            )
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+
+        batch["labels"] = labels
+
+        return batch
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--pre_trained_model", default='', type=str, help='Local path to pre-trained wav2vec2 model')
+    parser.add_argument("--train_file_path", default='', type=str, help='Local path to train file')
+    parser.add_argument("--valid_file_path", default='', type=str, help='Local path to valid file')
+
+    parser.add_argument("--warmup_steps", default=20000, type=int, help='')
+    parser.add_argument("--learning_rate", default=3e-5, type=float, help='')
+    args = parser.parse_args()
+
+    def prepare_dataset(batch):
+        # check that all files have the correct sampling rate
+        # assert (
+                # len(set(batch["sampling_rate"])) == 1
+        # ), f"Make sure all inputs have the same sampling rate of {processor.feature_extractor.sampling_rate}."
+
+        batch["input_values"] = processor(batch["speech"], sampling_rate=batch["sampling_rate"][0]).input_values
+
+        with processor.as_target_processor():
+            batch["labels"] = processor(batch["target_text"]).input_ids
+        return
+
+    def compute_metrics(pred, processor, metric):
+        pred_logits = pred.predictions
+        pred_ids = np.argmax(pred_logits, axis=-1)
+
+        pred.label_ids[pred.label_ids == -100] = processor.tokenizer.pad_token_id
+
+        pred_str = processor.batch_decode(pred_ids)
+        # we do not want to group tokens when computing the metrics
+        label_str = processor.batch_decode(pred.label_ids, group_tokens=False)
+
+        wer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+        return {"cer": cer}
+
+    # load dataset
+    print('Loading dataset....')
+    datasets = load_dataset('csv', name='cn', data_files={'train': args.train_file_path, 'valid': args.valid_file_path},
+                            cache_dir='/path/to/csv')
+    datasets = datasets.map(preprocess_data)
+
+    dataset_train = datasets['train']
+    dataset_valid = datasets['valid']
+
+    dataset_train = dataset_train.map(speech_file_to_array_fn,
+                                      remove_columns=dataset_train.column_names,
+                                      cache_file_name='/path/to/cache/of/train/speech/file')
+
+    dataset_valid = dataset_valid.map(speech_file_to_array_fn,
+                                      remove_columns=dataset_valid.column_names,
+                                      cache_file_name='/path/to/cache/of/valid/speech/file')
+
+    print('Tokenization')
+    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(args.pre_trained_model)
+
+    print('Feature extracting....')
+    feature_extractor = Wav2Vec2FeatureExtractor(args.pre_trained_model)
+    processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+    dataset_train = dataset_train.map(prepare_dataset,
+                                      remove_columns=dataset_train.column_names,
+                                      batched=True,
+                                      load_from_cache_file=True,
+                                      cache_file_name='/path/to/train')
+
+    dataset_valid = dataset_valid.map(prepare_dataset,
+                                      remove_columns=dataset_valid.column_names,
+                                      batched=True,
+                                      load_from_cache_file=True,
+                                      cache_file_name='/path/to/valid')
+
+
+    data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)
+    wer_metric = load_metric("cer")
+
+    # create model
+    model = Wav2Vec2ForCTC.from_pretrained(
+        args.pre_trained_model,
+        vocab_size=len(processor.tokenizer)
+    )
+    model.freeze_feature_extractor()
+
+    training_args = TrainingArguments(
+        output_dir="/path/to/output",
+        group_by_length=True,
+        per_device_train_batch_size=3,
+        gradient_accumulation_steps=1,
+        per_device_eval_batch_size=1,
+        metric_for_best_model='cer',
+        evaluation_strategy="steps",
+        eval_steps=15000,
+        logging_strategy="steps",
+        logging_steps=15000,
+        save_strategy="steps",
+        save_steps=15000,
+        num_train_epochs=100,
+        fp16=True,
+        learning_rate=args.learning_rate,
+        warmup_steps=args.warmup_steps,
+        save_total_limit=3,
+        report_to="tensorboard"
+    )
+
+    print('Training model....')
+    # Train
+    trainer = Trainer(
+        model=model,
+        data_collator=data_collator,
+        args=training_args,
+        compute_metrics=partial(compute_metrics, metric=cer_metric, processor=processor),
+        train_dataset=dataset_train,
+        eval_dataset=dataset_valid,
+        tokenizer=processor.feature_extractor,
+    )
+
+    trainer.train()