# Balochi STT (Whisper)

Fine-tuned **Whisper-small** speech-to-text model for **Balochi (Latin script)**.

| Item | Value |
| --- | --- |
| Task | Automatic speech recognition (ASR / STT) |
| Language | Balochi (Latin orthography) |
| Base model | Whisper-small |
| Sample rate | 16 kHz, mono |
| Best eval WER | ~4.5% |
| Training data | 5,237 clips (~3 hours) from `wavs/wavs` + `wavs/wavs.txt` |

---

## Folder layout

```text
balochi_whisper/
├── README.md
├── best_model/
│   ├── model/          # Whisper weights + config
│   ├── processor/      # tokenizer + feature extractor
│   └── eval_metrics.txt
└── checkpoints/        # training checkpoints (optional)
```

Use **`best_model/`** for inference.

---

## Requirements

```bash
source ~/venvs/torch/bin/activate
```

Needs: `torch`, `transformers`, `soundfile`, `torchaudio`, `numpy`.

GPU is recommended (tested on NVIDIA RTX 3060).

---

## Quick inference (Python)

```python
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
import torchaudio
from transformers import WhisperForConditionalGeneration, WhisperProcessor

ROOT = Path("balochi_whisper/best_model")
device = "cuda" if torch.cuda.is_available() else "cpu"

processor = WhisperProcessor.from_pretrained(ROOT / "processor")
model = WhisperForConditionalGeneration.from_pretrained(ROOT / "model").to(device)
model.eval()
model.config.forced_decoder_ids = None
model.generation_config.forced_decoder_ids = None
model.generation_config.pad_token_id = processor.tokenizer.pad_token_id

def load_16k(path: str) -> np.ndarray:
    wav, sr = sf.read(path, always_2d=False)
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    wav = np.asarray(wav, dtype=np.float32)
    if sr != 16000:
        wav = (
            torchaudio.functional.resample(
                torch.from_numpy(wav).unsqueeze(0), sr, 16000
            )
            .squeeze(0)
            .numpy()
        )
    return wav

wav = load_16k("your_audio.wav")
inputs = processor(
    wav, sampling_rate=16000, return_tensors="pt", return_attention_mask=True
)
with torch.no_grad():
    ids = model.generate(
        inputs.input_features.to(device),
        attention_mask=inputs.attention_mask.to(device),
        max_new_tokens=224,
    )
text = processor.batch_decode(ids.cpu(), skip_special_tokens=True)[0].strip()
print(text)
```

---

## CLI (from project root)

```bash
cd ~/Balochi_tts
source ~/venvs/torch/bin/activate
python infer_stt.py /path/to/audio.wav
```

Desktop GUI:

```bash
python stt_gui.py
```

---

## Training data format

```text
wavs/wavs/1.wav … N.wav
wavs/wavs.txt          # line i → transcript for (i+1).wav
```

Transcripts use Balochi **Latin** script (`á`, `é`, `ó` allowed).

Retrain / resume from project root:

```bash
python train_stt.py
```

Checkpoints are saved under `balochi_whisper/checkpoints/`. The best model (lowest WER) is copied to `balochi_whisper/best_model/`.

---

## Hugging Face

Upload this model:

```bash
# from Balochi_tts project root
bash scripts/push_models_to_hf.sh YOUR_HF_USER
```

That publishes `YOUR_HF_USER/balochi-whisper-stt` with `model/` and `processor/`.

Load from the Hub:

```python
from transformers import WhisperForConditionalGeneration, WhisperProcessor

repo = "YOUR_HF_USER/balochi-whisper-stt"
processor = WhisperProcessor.from_pretrained(repo, subfolder="processor")
model = WhisperForConditionalGeneration.from_pretrained(repo, subfolder="model")
```

If you upload files at the repo root (no subfolders), omit `subfolder=...`.

---

## Notes

- Input audio is converted to **16 kHz mono** before recognition.
- Audio longer than **30 seconds** should be chunked (the project CLI/GUI do this).
- Output is Balochi **Latin** text, not Arabic script.
- Do not force an English language token; Balochi is not a built-in Whisper language.

---

## Files to ship

Minimum files for inference:

```text
best_model/model/config.json
best_model/model/generation_config.json
best_model/model/model.safetensors
best_model/processor/   # full processor directory
```

`checkpoints/` is only needed to resume training.