README.md

---
datasets:
- lab260/mos260
- urgent-challenge/urgent2024_mos
language:
- ru
- en
tags:
- pytorch
- audio
- mos
---
# ProxyMOS

**ProxyMOS** is a lightweight Mean Opinion Score (MOS) prediction model built on top of the [OmniASR W2V 300M](https://github.com/facebookresearch/omnilingual-asr) encoder (Wav2Vec2 architecture via `fairseq2`).

Given any audio file, the model returns a single scalar speech quality score on the MOS scale. It is the distilled model from the paper **"ProxyMOS: Training Speech Quality Models via Ensemble-Derived MOS Targets"** (Borodin & Trokunov, MTUCI, 2026).

---

## Architecture

```
OmniASR-W2V-300M (encoder, fine-tuned with reduced lr)
        ↓
AttentiveStatsPooling  ->  [mean || std]  (dim: 2 x 1024 = 2048)
        ↓
Linear(2048 -> 1024) -> GELU -> Linear(1024 -> 1)
        ↓
     MOS score (scalar)
```

- **Encoder**: `omniASR_W2V_300M` — a 300M-parameter Wav2Vec2-style encoder instantiated via `fairseq2`.
- **Weights**: `best_model_full.pt` contains the trained ProxyMOS state dict, including the encoder and MOS head weights. No separate `omniASR-W2V-300M.pt` checkpoint is required.
- **Pooling**: Attentive Statistics Pooling — computes attention-weighted mean and standard deviation over the time dimension.
- **Head**: Two-layer MLP with GELU activation and a scalar output.

---

## Repository files

| File | Description |
|---|---|
| `config.json` | Model metadata used by Hugging Face Hub tooling |
| `best_model_full.pt` | Trained ProxyMOS state dict, including encoder and MOS head weights (~1.28 GB) |
| `inference_model.py` | Model architecture + inference script |

> **Note:** inference uses `fairseq2` to instantiate the `omniASR_W2V_300M` architecture, then loads all trained weights from `best_model_full.pt`. You do not need a separate `omniASR-W2V-300M.pt` file.

---

## Quick start

### 1. Clone the repository

```bash
git lfs install
git clone https://huggingface.co/lab260/ProxyMos
cd ProxyMos
```

### 2. Install dependencies

Requires Python 3.10+.

```bash
pip install torch torchaudio
pip install fairseq2
pip install accelerate scipy scikit-learn tqdm
```

> **GPU:** if CUDA is available, the model will use it automatically.



### 3. Run inference

```bash
python inference_model.py /path/to/your/audio.wav
```

**Example:**

```bash
python inference_model.py /home/ae_samples_ae_finetuned_vocoder_segment_28.wav
```

**Output:**

```
MOS: 3.8721
```

---

## Python API

```python
import torch
from inference_model import load_model, predict_mos

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = load_model("best_model_full.pt", device)

score = predict_mos(model, "/path/to/audio.wav", device)
print(f"MOS: {score:.4f}")
```

### Batch processing

```python
from pathlib import Path

audio_dir = Path("/path/to/audio_folder")

for wav_file in sorted(audio_dir.glob("*.wav")):
    score = predict_mos(model, str(wav_file), device)
    print(f"{wav_file.name}: {score:.4f}")
```

---
## Model Performance

### URGENT Dataset

| Model | Pearson | Spearman | Kendall | RMSE | MAE |
|-------|---------|----------|---------|------|-----|
| [Proxymos](https://huggingface.co/lab260/ProxyMos) | 0.802 | 0.799 | 0.617 | 0.474 | 0.370 |
| [Proxymosonnx](https://huggingface.co/lab260/ProxyMos) | 0.779 | 0.778 | 0.591 | 0.502 | 0.393 |
| [Proxymosonnx16](https://huggingface.co/lab260/ProxyMos) | 0.779 | 0.778 | 0.591 | 0.502 | 0.393 |
| [Whisqa](https://github.com/leto19/WhiSQA) | 0.780 | 0.773 | 0.588 | 0.502 | 0.397 |
| [Distillmos](https://github.com/microsoft/Distill-MOS) | 0.759 | 0.748 | 0.568 | 0.525 | 0.411 |
| [UTMOS](https://github.com/sarulab-speech/UTMOS22) | 0.721 | 0.748 | 0.560 | 0.565 | 0.450 |
| [XLS](https://huggingface.co/facebook/wav2vec2-xls-r-300m) | 0.743 | 0.745 | 0.561 | 0.542 | 0.426 |
| [MosNet](https://github.com/lochenchou/MOSNet) | 0.279 | 0.110 | 0.073 | 0.908 | 0.702 |
| [HuBERT](https://pytorch.org/audio/stable/pipelines.html#torchaudio.pipelines.HUBERT_BASE) | -0.005 | -0.002 | -0.001 | 1.073 | 0.855 |
| [NISQA](https://github.com/gabrielmittag/NISQA) | 0.028 | -0.027 | -0.021 | 1.055 | 0.863 |
| [Dnsmos](https://github.com/microsoft/DNS-Challenge/tree/master/DNSMOS) | -0.010 | -0.041 | -0.028 | 1.075 | 0.876 |

### MOS260 Dataset

| Model | Pearson | Spearman | Kendall | RMSE | MAE |
|-------|---------|----------|---------|------|-----|
| [Proxymosonnx](https://huggingface.co/lab260/ProxyMos) | 0.700 | 0.647 | 0.481 | 0.883 | 0.687 |
| [Proxymosonnx16](https://huggingface.co/lab260/ProxyMos) | 0.700 | 0.647 | 0.481 | 0.883 | 0.687 |
| [Proxymos](https://huggingface.co/lab260/ProxyMos) | 0.691 | 0.636 | 0.474 | 0.897 | 0.692 |
| [Distillmos](https://github.com/microsoft/Distill-MOS) | 0.670 | 0.613 | 0.456 | 0.927 | 0.712 |
| [UTMOS](https://github.com/sarulab-speech/UTMOS22) | 0.622 | 0.555 | 0.403 | 0.991 | 0.781 |
| [XLS](https://huggingface.co/facebook/wav2vec2-xls-r-300m) | 0.583 | 0.473 | 0.337 | 1.041 | 0.834 |
| [Whisqa](https://github.com/leto19/WhiSQA) | 0.576 | 0.466 | 0.332 | 1.050 | 0.841 |
| [NISQA](https://github.com/gabrielmittag/NISQA) | 0.261 | 0.226 | 0.158 | 1.386 | 1.105 |
| [Dnsmos](https://github.com/microsoft/DNS-Challenge/tree/master/DNSMOS) | 0.075 | 0.056 | 0.039 | 1.551 | 1.254 |
| [HuBERT](https://pytorch.org/audio/stable/pipelines.html#torchaudio.pipelines.HUBERT_BASE) | -0.003 | -0.009 | -0.006 | 1.614 | 1.294 |
| [MosNet](https://github.com/lochenchou/MOSNet) | -0.210 | -0.209 | -0.146 | 1.773 | 1.384 |

## Requirements

| Package | Version |
|---|---|
| Python | ≥ 3.10 |
| PyTorch | ≥ 2.0 |
| torchaudio | compatible with torch |
| fairseq2 | ≥ 0.3 |
| accelerate | ≥ 0.20 |
| scipy | ≥ 1.10 |
| scikit-learn | ≥ 1.2 |

---

## Technical details

- Audio is automatically resampled to **16,000 Hz**.
- Stereo input is averaged to mono before processing.
- Input is normalized to shape `[B, T]` internally.
- Inference runs correctly on CPU when no GPU is available.

---

## Results

| Benchmark | Spearman ρ | Pearson r | RMSE |
|---|---|---|---|
| URGENT | 0.802 | 0.806 | 0.471 |
| mos260 (Russian TTS) | 0.647 | 0.700 | 0.883 |

ProxyMOS outperforms all individual teacher models (WhiSQA, DistillMOS, UTMOS, XLS-R) on both benchmarks.

---

## Citation

```bibtex
@article{borodin2026proxymos,
  title     = {ProxyMOS: Training Speech Quality Models via Ensemble-Derived MOS Targets},
  author    = {Borodin, Kirill and Trokunov, Maksim},
  year      = {2026},
  institution = {MTUCI, Moscow}
}
```

---

## License

Please refer to the [fairseq2](https://github.com/facebookresearch/fairseq2) and [omnilingual-asr](https://github.com/facebookresearch/omnilingual-asr) repositories for encoder licensing terms.

## Contact

- Email: kborodin.research@gmail.com
- Telegram: [@korallll_ai](https://t.me/korallll_ai)