Upload STT Training Script.py

STT Training Script.py

@@ -0,0 +1,216 @@
+# Import required libraries
+from datasets import load_dataset, Audio
+from transformers import (
+    WhisperProcessor,
+    WhisperForConditionalGeneration,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer
+)
+import torch
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+from functools import partial
+import evaluate
+
+# Load the dataset
+dataset = load_dataset("")          # Specify Data Repo on HF
+dataset
+
+# Split the dataset into train and test sets (80-20 split)
+split_dataset = dataset['train'].train_test_split(test_size=0.2)
+split_dataset
+
+# Select only the relevant columns for training
+split_dataset['train'] = split_dataset['train'].select_columns(["audio", "sentence"])
+split_dataset['train']
+
+# Initialize the Whisper processor for Swahili transcription
+processor = WhisperProcessor.from_pretrained(
+    "openai/whisper-small", 
+    language="swahili", 
+    task="transcribe"
+)
+
+# Print audio features before and after resampling to match Whisper's expected sampling rate
+print('BEFORE>>> ', split_dataset['train'].features['audio'])
+sampling_rate = processor.feature_extractor.sampling_rate
+split_dataset['train'] = split_dataset['train'].cast_column(
+    "audio", 
+    Audio(sampling_rate=sampling_rate)
+)
+print('AFTER>>> ', split_dataset['train'].features['audio'])
+
+# Do the same for the test set
+print('BEFORE>>> ', split_dataset['test'].features['audio'])
+split_dataset['test'] = split_dataset['test'].cast_column(
+    "audio", 
+    Audio(sampling_rate=sampling_rate)
+)
+print('AFTER>>> ', split_dataset['test'].features['audio'])
+
+def prepare_dataset(example):
+    """Preprocess audio and text data for Whisper model training"""
+    audio = example["audio"]
+    
+    # Process audio and text using Whisper processor
+    example = processor(
+        audio=audio["array"],
+        sampling_rate=audio["sampling_rate"],
+        text=example["sentence"],
+    )
+    
+    # Compute input length of audio sample in seconds
+    example["input_length"] = len(audio["array"]) / audio["sampling_rate"]
+    
+    return example
+
+# Apply preprocessing to train and test sets
+split_dataset['train'] = split_dataset['train'].map(
+    prepare_dataset, 
+    remove_columns=split_dataset['train'].column_names, 
+    num_proc=4  # Use 4 processes for faster preprocessing
+)
+
+split_dataset['test'] = split_dataset['test'].map(
+    prepare_dataset, 
+    remove_columns=split_dataset['test'].column_names, 
+    num_proc=1
+)
+
+# Filter out audio samples longer than 30 seconds
+max_input_length = 30.0
+def is_audio_in_length_range(length):
+    return length < max_input_length
+
+split_dataset['train'] = split_dataset['train'].filter(
+    is_audio_in_length_range,
+    input_columns=["input_length"],
+)
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    """Custom data collator for Whisper speech-to-sequence tasks with padding"""
+    processor: Any
+
+    def __call__(
+        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]
+    ) -> Dict[str, torch.Tensor]:
+        # Split inputs and labels since they need different padding methods
+        # First process audio inputs
+        input_features = [
+            {"input_features": feature["input_features"][0]} for feature in features
+        ]
+        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+
+        # Process label sequences
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+        labels_batch = self.processor.tokenizer.pad(label_features, 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
+        )
+
+        # Remove BOS token if it was appended previously
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+
+# Initialize data collator
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+# Load evaluation metric (Word Error Rate)
+metric = evaluate.load("wer")
+
+# Initialize text normalizer for evaluation
+from transformers.models.whisper.english_normalizer import BasicTextNormalizer
+normalizer = BasicTextNormalizer()
+
+def compute_metrics(pred):
+    """Compute WER (Word Error Rate) metrics for evaluation"""
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # Replace -100 with pad_token_id
+    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id
+
+    # Decode predictions and labels
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = processor.batch_decode(label_ids, skip_special_tokens=True)
+
+    # Compute orthographic WER
+    wer_ortho = 100 * metric.compute(predictions=pred_str, references=label_str)
+
+    # Compute normalized WER
+    pred_str_norm = [normalizer(pred) for pred in pred_str]
+    label_str_norm = [normalizer(label) for label in label_str]
+    
+    # Filter samples with non-zero references
+    pred_str_norm = [
+        pred_str_norm[i] for i in range(len(pred_str_norm)) if len(label_str_norm[i]) > 0
+    ]
+    label_str_norm = [
+        label_str_norm[i] for i in range(len(label_str_norm)) if len(label_str_norm[i]) > 0
+    ]
+
+    wer = 100 * metric.compute(predictions=pred_str_norm, references=label_str_norm)
+
+    return {"wer_ortho": wer_ortho, "wer": wer}
+
+# Load pre-trained Whisper model
+model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+
+# Disable cache during training (incompatible with gradient checkpointing)
+model.config.use_cache = False
+
+# Configure generation settings (re-enable cache for generation)
+model.generate = partial(
+    model.generate, 
+    language="swahili", 
+    task="transcribe", 
+    use_cache=True
+)
+
+# Set up training arguments
+training_args = Seq2SeqTrainingArguments(
+    output_dir="./model",                      # Output directory
+    per_device_train_batch_size=16,            # Batch size for training
+    gradient_accumulation_steps=1,             # Number of steps before gradient update
+    learning_rate=1e-6,                        # Learning rate
+    lr_scheduler_type="constant_with_warmup",  # Learning rate scheduler
+    warmup_steps=50,                           # Warmup steps
+    max_steps=10000,                           # Total training steps
+    gradient_checkpointing=True,               # Use gradient checkpointing
+    fp16=True,                                 # Use mixed precision training
+    fp16_full_eval=True,                       # Use mixed precision evaluation
+    evaluation_strategy="steps",               # Evaluation strategy
+    per_device_eval_batch_size=16,             # Batch size for evaluation
+    predict_with_generate=True,                # Use generation for evaluation
+    generation_max_length=225,                 # Maximum generation length
+    save_steps=500,                            # Save checkpoint every N steps
+    eval_steps=500,                            # Evaluate every N steps
+    logging_steps=100,                         # Log metrics every N steps
+    report_to=["tensorboard", "wandb"],        # Logging integrations
+    load_best_model_at_end=True,               # Load best model at end of training
+    metric_for_best_model="wer",               # Metric for selecting best model
+    greater_is_better=False,                   # Lower WER is better
+    push_to_hub=True,                          # Push to Hugging Face Hub
+    save_total_limit=3,                        # Maximum number of checkpoints to keep
+)
+
+# Initialize trainer
+trainer = Seq2SeqTrainer(
+    args=training_args,
+    model=model,
+    train_dataset=split_dataset['train'],
+    eval_dataset=split_dataset['test'],
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor,  # Changed from processing_class to tokenizer
+)
+
+# Start training
+trainer.train()