Add model config and update loading instructions

README.md

@@ -23,14 +23,15 @@ Given any audio file, the model returns a single scalar speech quality score on
 ```
 OmniASR-W2V-300M (encoder, fine-tuned with reduced lr)
         ↓
-AttentiveStatsPooling  →  [mean ‖ std]  (dim: 2 × 768 = 1536)
+AttentiveStatsPooling  ->  [mean || std]  (dim: 2 x 1024 = 2048)
         ↓
-Linear(1536 → 1024) → GELU → Linear(1024 → 1)
+Linear(2048 -> 1024) -> GELU -> Linear(1024 -> 1)
         ↓
      MOS score (scalar)
 ```
 
-- **Encoder**: `omniASR_W2V_300M` — a 300M-parameter Wav2Vec2 model from the `omnilingual_asr` package, pretrained on 1600+ languages. The encoder is not frozen during training; it is updated with a learning rate 100× smaller than the head.
+- **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.
 
@@ -40,10 +41,11 @@ Linear(1536 → 1024) → GELU → Linear(1024 → 1)
 
 | File | Description |
 |---|---|
-| `best_model_full.pt` | Trained ProxyMOS weights (state_dict, ~1.28 GB) |
+| `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:** encoder weights (`omniASR-W2V-300M.pt`) are **not stored** in this repository — they are downloaded separately via `omnilingual_asr` (see below).
+> **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.
 
 ---
 
@@ -95,14 +97,10 @@ MOS: 3.8721
 
 ```python
 import torch
-from inference_model import OmniMOS, load_checkpoint, predict_mos
+from inference_model import load_model, predict_mos
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-encoder, _ = load_checkpoint("omniASR-W2V-300M.pt")
-model = OmniMOS(encoder=encoder)
-model.load_state_dict(torch.load("best_model_full.pt", map_location="cpu"))
-model.to(device).eval()
+model = load_model("best_model_full.pt", device)
 
 score = predict_mos(model, "/path/to/audio.wav", device)
 print(f"MOS: {score:.4f}")

config.json

@@ -0,0 +1,15 @@
+{
+  "model_type": "proxymos",
+  "architectures": [
+    "OmniMOS"
+  ],
+  "library_name": "pytorch",
+  "format": "pytorch_state_dict",
+  "weights": "best_model_full.pt",
+  "inference_script": "inference_model.py",
+  "sample_rate": 16000,
+  "encoder": "omniASR_W2V_300M",
+  "hidden_size": 1024,
+  "pooling": "attentive_statistics",
+  "output": "mos_score"
+}