Create train_pipeline.py

train_pipeline.py

@@ -0,0 +1,104 @@
+# train_pipeline.py
+import re
+import numpy as np
+from jiwer import wer
+from transformers import (
+    Wav2Vec2ForCTC,
+    Wav2Vec2CTCTokenizer,
+    Wav2Vec2FeatureExtractor,
+    Wav2Vec2Processor,
+    TrainingArguments,
+    Trainer
+)
+from datasets import Audio
+import torch
+from dataclasses import dataclass
+from preprocess import load_telugu_dataset, normalize_text
+from vocab import build_vocab
+
+def prepare_dataset(batch, processor):
+    speech = batch["audio"]["array"]
+    batch["input_values"] = processor(speech, sampling_rate=16000).input_values[0]
+    batch["labels"] = processor.tokenizer(normalize_text(batch["text"])).input_ids
+    return batch
+
+@dataclass
+class DataCollatorCTC:
+    processor: Wav2Vec2Processor
+    padding: bool = True
+
+    def __call__(self, features):
+        inputs = [{"input_values": f["input_values"]} for f in features]
+        labels = [{"input_ids": f["labels"]} for f in features]
+
+        batch = self.processor.pad(inputs, return_tensors="pt")
+        with self.processor.as_target_processor():
+            labels_batch = self.processor.pad(labels, return_tensors="pt")
+
+        labels = labels_batch["input_ids"]
+        labels[labels == self.processor.tokenizer.pad_token_id] = -100
+        batch["labels"] = labels
+        return batch
+
+def train_model():
+    # 1. Load dataset
+    ds = load_telugu_dataset()
+    ds = ds.train_test_split(test_size=0.1)
+    train = ds["train"]
+    test = ds["test"]
+
+    # 2. Build vocab
+    build_vocab(train, text_col="text")
+
+    # 3. Processor
+    tokenizer = Wav2Vec2CTCTokenizer("vocab.json", pad_token="[PAD]", unk_token="[UNK]")
+    extractor = Wav2Vec2FeatureExtractor(sampling_rate=16000, do_normalize=True)
+    processor = Wav2Vec2Processor(extractor, tokenizer)
+
+    # 4. Prepare
+    train = train.map(lambda x: prepare_dataset(x, processor))
+    test = test.map(lambda x: prepare_dataset(x, processor))
+
+    # 5. Load XLS-R model
+    model = Wav2Vec2ForCTC.from_pretrained(
+        "facebook/wav2vec2-xls-r-300m",
+        vocab_size=len(tokenizer),
+        pad_token_id=tokenizer.pad_token_id,
+        ctc_loss_reduction="mean"
+    )
+    model.freeze_feature_extractor()
+
+    data_collator = DataCollatorCTC(processor)
+
+    def compute_metrics(pred):
+        pred_ids = np.argmax(pred.predictions, axis=-1)
+        pred.label_ids[pred.label_ids == -100] = tokenizer.pad_token_id
+
+        preds = processor.batch_decode(pred_ids)
+        refs = processor.batch_decode(pred.label_ids)
+        return {"wer": wer(refs, preds)}
+
+    args = TrainingArguments(
+        output_dir="./model",
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        fp16=True,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        num_train_epochs=5,
+        push_to_hub=True,
+        hub_model_id="your-username/telugu-asr-xlsr"
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=args,
+        train_dataset=train,
+        eval_dataset=test,
+        tokenizer=processor.feature_extractor,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    trainer.train()
+    trainer.push_to_hub()