Update README.md

README.md

@@ -67,6 +67,183 @@ The following hyperparameters were used during training:
 - num_epochs: 5
 - mixed_precision_training: Native AMP
 
+### Training code
+
+```bash
+pip install datasets librosa scikit-learn torch torchaudio evaluate jiwer nltk
+pip install --upgrade datasets
+```
+
+```python
+from huggingface_hub import login
+from datasets import load_dataset, DatasetDict
+from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2CTCTokenizer, Wav2Vec2Processor, Wav2Vec2ForCTC, TrainingArguments, Trainer
+from datasets import load_dataset, Audio
+import torch
+import torchaudio
+import re
+import evaluate
+import numpy as np
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+
+login("***")
+
+
+common_voice = DatasetDict()
+common_voice["train"] = load_dataset("mozilla-foundation/common_voice_17_0", "ru", split="train")
+common_voice["test"] = load_dataset("mozilla-foundation/common_voice_17_0", "ru", split="test")
+
+common_voice = common_voice.remove_columns(["accent", "age", "client_id", "down_votes", "gender", "locale", "path", "segment", "up_votes"])
+common_voice = common_voice.cast_column("audio", Audio(sampling_rate=16000))
+
+tokenizer = Wav2Vec2CTCTokenizer.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-russian")
+feature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=16000, padding_value=0.0, return_attention_mask=True)
+processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+
+    # batched output is "un-batched"
+    batch["input_values"] = processor(audio["array"], sampling_rate=audio["sampling_rate"]).input_values[0]
+    batch["input_length"] = len(batch["input_values"])
+
+    with processor.as_target_processor():
+        batch["labels"] = processor(batch["sentence"]).input_ids
+    return batch
+
+
+common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice["train"].column_names, num_proc=2)
+
+
+def has_valid_labels(example):
+    return example["labels"] is not None and len(example["labels"]) > 0
+
+common_voice = common_voice.filter(has_valid_labels, num_proc=2)
+
+ya nadeusb eta poebota ne zavosnet
+
+common_voice
+
+common_voice
+
+wer_metric = evaluate.load("wer")
+cer_metric = evaluate.load("cer")
+
+def compute_metrics(pred):
+    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, skip_special_tokens=True)
+    label_str = processor.batch_decode(pred.label_ids, group_tokens=False, skip_special_tokens=True)
+
+    pairs = [(ref.strip(), hyp.strip()) for ref, hyp in zip(label_str, pred_str)]
+    pairs = [(ref, hyp) for ref, hyp in pairs if len(ref) > 0]
+
+    if len(pairs) == 0:
+        return {"wer": 1.0, "cer": 1.0, "ser": 1.0}
+
+    label_str, pred_str = zip(*pairs)
+
+    wer = 100 * wer_metric.compute(predictions=pred_str, references=label_str)
+    cer = 100 * cer_metric.compute(predictions=pred_str, references=label_str)
+
+    incorrect_sentences = sum([ref != pred for ref, pred in zip(label_str, pred_str)])
+    ser = 100 * incorrect_sentences / len(label_str)
+
+    return {
+        "wer": wer,
+        "cer": cer,
+        "ser": ser
+    }
+
+
+model = Wav2Vec2ForCTC.from_pretrained(
+    "jonatasgrosman/wav2vec2-large-xlsr-53-russian",
+    ctc_loss_reduction="mean",
+    pad_token_id=processor.tokenizer.pad_token_id,
+)
+
+@dataclass
+class DataCollatorCTCWithPadding:
+    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 lengths 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
+
+
+data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)
+
+training_args = TrainingArguments(
+  output_dir="/content/drive/MyDrive/models/wav2vec2-large-ru-5ep",
+  logging_dir="/content/drive/MyDrive/models/wav2vec2-large-ru-5ep",
+  group_by_length=True,
+  per_device_train_batch_size=8,
+  per_device_eval_batch_size=4,
+  eval_strategy="steps",
+  logging_strategy="steps",
+  save_strategy="steps",
+  num_train_epochs=5,
+  logging_steps=25,
+  eval_steps=500,
+  save_steps=500,
+  fp16=True,
+  optim="adamw_torch_fused",
+  torch_compile=True,
+  gradient_checkpointing=True,
+  learning_rate=1e-4,
+  weight_decay=0.005,
+  report_to=["tensorboard"],
+  push_to_hub=False
+)
+
+trainer = Trainer(
+    model=model,
+    data_collator=data_collator,
+    args=training_args,
+    compute_metrics=compute_metrics,
+    train_dataset=common_voice["train"],
+    eval_dataset=common_voice["test"],
+    tokenizer=processor,
+)
+
+trainer.train()
+```
 ### Training results
 
 | Training Loss | Epoch  | Step  | Validation Loss | Wer     | Cer    | Ser     |