add train code

README.md

@@ -1,3 +1,92 @@
 ---
 license: mit
 ---
+# Speech to Text
+
+Fine-tune a [Wav2Vec2](https://huggingface.co/facebook/wav2vec2-base) acoustic model on the [LJSpeech](https://keithito.com/LJ-Speech-Dataset/) dataset using CTC, then export it to ONNX for inference.
+
+## Requirements
+
+- Python >= 3.10
+- A CUDA-capable GPU is recommended for training
+
+Install dependencies:
+
+```bash
+pip install -e .
+```
+
+## Training
+
+Fine-tune `facebook/wav2vec2-base` on LJSpeech (5% held out for eval). Training takes ~10 epochs by default and writes checkpoints to `wav2vec2-ljspeech/`.
+
+```bash
+python train.py
+```
+
+Key settings live at the top of `train.py`:
+
+| Constant | Default | Purpose |
+| --- | --- | --- |
+| `MODEL_ID` | `facebook/wav2vec2-base` | Pre-trained wav2vec2 checkpoint |
+| `DATASET_ID` | `lj_speech` | HuggingFace dataset id |
+
+Training hyperparameters (batch size, epochs, learning rate, etc.) are configured through `TrainingArguments` inside `train.py`.
+
+Monitor progress with TensorBoard:
+
+```bash
+tensorboard --logdir wav2vec2-ljspeech
+```
+
+## ONNX Export
+
+Export the trained checkpoint to ONNX and validate it with ONNX Runtime:
+
+```bash
+python export_onnx.py
+```
+
+Options:
+
+```
+--model-dir   Checkpoint directory (default: wav2vec2-ljspeech)
+--output      Output ONNX path   (default: wav2vec2-ljspeech.onnx)
+--opset       ONNX opset version (default: 17)
+```
+
+The exported model uses dynamic axes on batch and time, so it accepts audio of any length.
+
+## Inference
+
+```python
+import numpy as np
+import onnxruntime as ort
+from transformers import Wav2Vec2Processor
+
+processor = Wav2Vec2Processor.from_pretrained("wav2vec2-ljspeech")
+session = ort.InferenceSession("wav2vec2-ljspeech.onnx")
+
+# audio_array: 16 kHz mono float32 numpy array
+inputs = processor(audio_array, sampling_rate=16000, return_tensors="np")
+logits = session.run(None, {"input_values": inputs.input_values})[0]
+text = processor.tokenizer.batch_decode(np.argmax(logits, axis=-1))[0]
+print(text)
+```
+
+Notes:
+
+- Audio must be **16 kHz mono float32**.
+- The `Wav2Vec2Processor` handles waveform normalization and tokenization — always pass audio through it before the ONNX session.
+- This exports the **acoustic model only**. Add an external LM (e.g. KenLM) for language-model-rescored decoding if needed.
+
+## Project Layout
+
+```
+speech-to-text/
+├── train.py          # Wav2Vec2 + CTC fine-tuning on LJSpeech
+├── export_onnx.py    # ONNX export and ONNX Runtime validation
+├── main.py           # Placeholder entry point
+├── pyproject.toml    # Project metadata and dependencies
+└── README.md
+```

export_onnx.py

@@ -0,0 +1,69 @@
+import argparse
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+import torch
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+
+DEFAULT_MODEL_DIR = "wav2vec2-ljspeech"
+DEFAULT_OUTPUT = "wav2vec2-ljspeech.onnx"
+SAMPLE_RATE = 16000
+
+
+def export(model_dir: str, output: str, opset: int) -> None:
+    model_path = Path(model_dir)
+    processor = Wav2Vec2Processor.from_pretrained(model_path)
+    model = Wav2Vec2ForCTC.from_pretrained(model_path)
+    model.eval()
+
+    seq_len = SAMPLE_RATE
+    dummy = torch.zeros(1, seq_len, dtype=torch.float32)
+
+    dynamic_axes = {
+        "input_values": {0: "batch", 1: "time"},
+        "logits": {0: "batch", 1: "time"},
+    }
+
+    output_path = Path(output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with torch.no_grad():
+        torch.onnx.export(
+            model,
+            (dummy,),
+            str(output_path),
+            input_names=["input_values"],
+            output_names=["logits"],
+            dynamic_axes=dynamic_axes,
+            opset_version=opset,
+            do_constant_folding=True,
+        )
+
+    print(f"Exported ONNX model to {output_path}")
+    print(f"  opset: {opset}, vocab_size: {len(processor.tokenizer)}")
+    print(f"  size: {output_path.stat().st_size / (1024 * 1024):.1f} MB")
+
+    validate(output_path, processor, seq_len)
+
+
+def validate(onnx_path: Path, processor, seq_len: int) -> None:
+    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+
+    audio = np.zeros((1, seq_len), dtype=np.float32)
+    outputs = session.run(None, {"input_values": audio})
+    logits = outputs[0]
+
+    print(f"Validated with ONNX Runtime: output shape = {logits.shape}")
+    pred_ids = np.argmax(logits, axis=-1)
+    text = processor.tokenizer.batch_decode(pred_ids)[0]
+    print(f"Decoded dummy input -> {text!r}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Export Wav2Vec2 CTC to ONNX")
+    parser.add_argument("--model-dir", default=DEFAULT_MODEL_DIR)
+    parser.add_argument("--output", default=DEFAULT_OUTPUT)
+    parser.add_argument("--opset", type=int, default=17)
+    args = parser.parse_args()
+    export(args.model_dir, args.output, args.opset)

train.py

@@ -0,0 +1,128 @@
+import torch
+from dataclasses import dataclass
+from typing import Any
+
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    Wav2Vec2ForCTC,
+    Wav2Vec2Processor,
+)
+
+MODEL_ID = "facebook/wav2vec2-base"
+DATASET_ID = "lj_speech"
+
+
+@dataclass
+class DataCollatorCTCWithPadding:
+    processor: Wav2Vec2Processor
+    padding: bool | str = True
+
+    def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        input_features = [{"input_values": f["input_values"]} for f in features]
+        label_features = [{"input_ids": f["labels"]} for f in features]
+
+        batch = self.processor.feature_extractor.pad(
+            input_features, padding=self.padding, return_tensors="pt"
+        )
+        labels_batch = self.processor.tokenizer.pad(
+            label_features, padding=self.padding, return_tensors="pt"
+        )
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100
+        )
+        batch["labels"] = labels
+        return batch
+
+
+def prepare_dataset(batch, processor):
+    audio = batch["audio"]
+    batch["input_values"] = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"]
+    ).input_values[0]
+    batch["input_length"] = len(batch["input_values"])
+    batch["labels"] = processor(text=batch["normalized_text"]).input_ids
+    return batch
+
+
+def compute_metrics(pred):
+    from jiwer import wer
+
+    pred_logits = pred.predictions
+    pred_ids = torch.argmax(torch.tensor(pred_logits), dim=-1)
+
+    pred_str = [s if len(s) > 0 else "-" for s in tokenizer.batch_decode(pred_ids)]
+    label_ids = pred.label_ids.copy()
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+    label_str = tokenizer.batch_decode(label_ids, group_tokens=False)
+
+    error = wer(predictions=pred_str, references=label_str)
+    return {"wer": error}
+
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    processor = Wav2Vec2Processor.from_pretrained(MODEL_ID)
+
+    dataset = load_dataset(DATASET_ID, split="train")
+    dataset = dataset.train_test_split(test_size=0.05, seed=42)
+
+    dataset = dataset.map(
+        lambda batch: prepare_dataset(batch, processor),
+        remove_columns=dataset.column_names["train"],
+        num_proc=4,
+    )
+
+    dataset = dataset.filter(
+        lambda x: x < processor.tokenizer.model_max_length,
+        input_columns=["input_length"],
+        num_proc=4,
+    )
+
+    data_collator = DataCollatorCTCWithPadding(processor=processor)
+
+    model = Wav2Vec2ForCTC.from_pretrained(
+        MODEL_ID,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+        feat_proj_dropout=0.0,
+        mask_time_prob=0.0,
+        layerdrop=0.0,
+        ctc_loss_reduction="mean",
+        pad_token_id=processor.tokenizer.pad_token_id,
+        vocab_size=len(processor.tokenizer),
+    )
+    model.freeze_feature_extractor()
+
+    training_args = TrainingArguments(
+        output_dir="wav2vec2-ljspeech",
+        per_device_train_batch_size=8,
+        gradient_accumulation_steps=2,
+        learning_rate=3e-4,
+        num_train_epochs=10,
+        warmup_steps=500,
+        fp16=True,
+        save_steps=500,
+        eval_strategy="steps",
+        eval_steps=500,
+        logging_steps=100,
+        save_total_limit=3,
+        report_to=["tensorboard"],
+        gradient_checkpointing=True,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["test"],
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+        processing_class=processor.tokenizer,
+    )
+
+    trainer.train()
+    trainer.save_model(training_args.output_dir)
+    processor.save_pretrained(training_args.output_dir)