Create README.md

README.md

@@ -0,0 +1,82 @@
+---
+datasets:
+- openslr/librispeech_asr
+- ghananlpcommunity/twi-speech-text-multispeaker-16k
+- ghananlpcommunity/twi-english-paragraph-dataset_news
+---
+
+# MMT-JEPA
+
+A multimodal machine translation model for English ↔ Twi using a JEPA (Joint Embedding Predictive Architecture) objective.
+
+## What it does
+
+Learns a shared latent space across text and audio in both languages by training a predictor to anticipate target representations from context — no reconstruction loss, no cascaded pipeline.
+
+Three training objectives:
+- **A** — Audio → Text (both languages)
+- **B** — Text → Text (translation)
+- **C** — Text → Audio (both languages)
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `model.py` | `MMT_JEPA` model + EMA target encoder |
+| `dataset.py` | `ObjA`, `ObjB`, `ObjC` dataset classes |
+| `tokenizer.py` | Trains a joint BPE tokenizer on all text data |
+| `train.py` | Training loop (all objectives) |
+| `train_b.py` | Training loop (Objective B only) |
+
+## Setup
+
+```bash
+pip install torch librosa soundfile sentencepiece datasets
+```
+
+## Usage
+
+**1. Train the tokenizer**
+```bash
+python tokenizer.py
+# outputs: tokenizer.model, tokenizer.vocab
+```
+
+**2. Train the model**
+```bash
+python train.py
+```
+
+Checkpoints saved to `checkpoints/epoch{N}.pt` after each epoch.
+
+## Data
+
+| Objective | Dataset |
+|---|---|
+| A + C (English audio) | [LibriSpeech train-clean-100](https://huggingface.co/datasets/openslr/librispeech_asr) |
+| A + C (Twi audio) | [twi-speech-text-multispeaker-16k](https://huggingface.co/datasets/ghananlpcommunity/twi-speech-text-multispeaker-16k) |
+| B (translation) | [twi-english-paragraph-dataset_news](https://huggingface.co/datasets/ghananlpcommunity/twi-english-paragraph-dataset_news) · [english-twi-sentences-non-nouns](https://huggingface.co/datasets/ghananlpcommunity/english-twi-sentences-non-nouns) · [english-twi-nouns-v2](https://huggingface.co/datasets/ghananlpcommunity/english-twi-nouns-v2) |
+
+All datasets load automatically via HuggingFace on first run.
+
+## Model config
+
+Edit `ModelConfig` in `model.py` to change capacity:
+
+```python
+d_model      = 512    # embedding dimension
+trunk_layers = 6      # shared transformer depth
+vocab_size   = 16_000
+n_mels       = 80
+sample_rate  = 16_000
+```
+
+## Training notes
+
+- First 5 epochs run text-only (ObjB) to warm up representations before audio is introduced
+- L2 normalization applied to both sides before MSE loss to keep scale stable across modalities
+- EMA target encoder uses cosine-annealed decay (0.990 → 0.996)
+- Collapse logged as `COLLAPSE` when `std < 0.01` or `cos_sim > 0.99`
+
+
+Get code on github [MMT-JEPA](https://github.com/etornam45/mmt-jepa)