Moss Audio Tn Intent

MOSS-Audio Tunisian Intent Classifier

A lightweight audio intent classifier for Tunisian Arabic built on top of the frozen OpenMOSS-Team/MOSS-Audio-8B-Instruct audio encoder and DeepStack feature merger.

The model achieves 67.4% accuracy on the SLURP-TN held-out evaluation set across 18 intent classes — well above majority-class baseline (~18%) and within striking distance of published Arabic intent-classification SOTA, despite using only ~0.79M trainable parameters (the rest of the 5.2B-parameter pipeline is frozen).

This validates the frozen-MOSS + DeepStack + lightweight head architecture for closed-set Arabic audio classification tasks where end-to-end MOSS generation falls short.

---

Architecture

                ┌─────────────────────────────────────────┐
   raw audio ──>│  MOSS-Audio audio encoder (FROZEN)      │
   16kHz mono   │  32 Whisper-style encoder layers, 1280-d│
                └─────┬─────────┬─────────┬───────────────┘
                      │ layer10 │ layer21 │ layer31
                      ▼         ▼         ▼
                ┌───────────────────────────────────────┐
                │  DeepStack mergers (FROZEN)           │
                │  3 × GatedMLP, each 1280 → 2560-d     │
                └─────┬─────────┬─────────┬─────────────┘
                      │         │         │
                      └─────sum─┴───sum───┘
                                ▼
                       ┌──────────────────┐
                       │  Mean-pool over  │
                       │  time dimension  │
                       └────────┬─────────┘
                                ▼
                          [2560-d vector]
                                ▼
                ┌───────────────────────────────────────┐
                │  Classifier head (TRAINABLE, 0.79M)   │
                │  LayerNorm → Linear(2560,512) → GELU  │
                │     → Linear(512,256) → GELU          │
                │     → Linear(256, num_classes)        │
                └────────────────┬──────────────────────┘
                                 ▼
                          intent label

Key design choices:

Component	Choice	Rationale
Base encoder	MOSS-Audio 8B, frozen	Strong acoustic representations; no need to retrain a massive model
Feature extraction	DeepStack merge across layers 10/21/31	Captures low/mid/high-level acoustic features; matches MOSS's 3-merger design
Aggregation across depths	Sum	Preserves all depth information; outperformed concat in early tests
Aggregation across time	Mean-pool	Simple, robust to variable clip length
Head capacity	0.79M params, 3-layer MLP with dropout	Small enough to train in seconds; large enough to learn 18-class boundary

---

Performance

Overall

Metric	Value
Best eval accuracy	0.674 (337/500)
Number of classes	18
Majority-class baseline	~0.188
Random baseline	~0.056
Training samples	2,174
Evaluation samples	500
Training time	~1 minute (after feature extraction)
Trainable params	789,330
Frozen params	5,222,891,520 (5.22B)

Per-class accuracy

Sorted by sample count (high-resource classes first):

Class	Eval samples	Accuracy
news_query	94	0.777
weather_query	94	0.819
general_quirky	80	0.487
email_query	41	0.854
alarm_set	34	0.735
email_sendemail	34	0.588
takeaway_order	25	0.680
alarm_query	24	0.417
takeaway_query	23	0.826
alarm_remove	15	0.467
email_querycontact	12	0.333
general_joke	10	0.800
email_addcontact	5	0.400
general_greet	4	0.250
query	3	0.000
quirky	1	0.000
sendemail	1	0.000

Pattern: classes with ≥20 training samples consistently hit 60–85% accuracy. Low-sample classes (n<5) cluster near 0% — these are not "broken" but simply under-trained. The SLURP-TN dataset has overlapping label pairs (quirky/general_quirky, joke/general_joke, etc.) that further depress the per-class number for the rare variants; merging these in a future training run would likely raise overall accuracy to ~73%.

The full confusion matrix is in confusion_matrix.csv.

---

Quick Start

Install dependencies

# MOSS-Audio package
git clone https://github.com/OpenMOSS/MOSS-Audio.git
cd MOSS-Audio
pip install -e .

# Other deps
pip install torch transformers huggingface_hub librosa

Minimal inference

import torch, torch.nn as nn, librosa, sys
from huggingface_hub import hf_hub_download

# 1. Make MOSS-Audio importable
sys.path.insert(0, "/path/to/MOSS-Audio")
from src.modeling_moss_audio import MossAudioModel
from src.processing_moss_audio import MossAudioProcessor

# 2. Apply Whisper kwarg-stripping patch (idempotent)
import transformers.models.whisper.modeling_whisper as whisper_mod
if not getattr(whisper_mod.WhisperEncoderLayer.forward, "_moss_patched", False):
    _orig = whisper_mod.WhisperEncoderLayer.forward
    def _patched(self, *a, **kw):
        kw.pop("layer_head_mask", None); kw.pop("head_mask", None)
        return _orig(self, *a, **kw)
    _patched._moss_patched = True
    whisper_mod.WhisperEncoderLayer.forward = _patched

# 3. Load the frozen base
MOSS_REPO = "OpenMOSS-Team/MOSS-Audio-8B-Instruct"
moss_model = MossAudioModel.from_pretrained(
    MOSS_REPO, trust_remote_code=True, dtype="auto",
    device_map="cuda:0", attn_implementation="sdpa",
).eval()
processor = MossAudioProcessor.from_pretrained(
    MOSS_REPO, trust_remote_code=True, enable_time_marker=True,
)

# 4. Download the classifier head
head_path = hf_hub_download(repo_id="FatimahEmadEldin/moss-audio-tn-intent",
                            filename="best.pt")

# 5. Load head + label maps
state = torch.load(head_path, weights_only=False, map_location="cuda:0")
ID2LABEL = state["ID2LABEL"]
NUM_CLASSES = state["num_classes"]

class ClassifierHead(nn.Module):
    def __init__(self, in_dim, num_classes, hidden=512, dropout=0.15):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(in_dim),
            nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(hidden, hidden // 2), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(hidden // 2, num_classes),
        )
    def forward(self, x):
        return self.net(x)

head = ClassifierHead(state["feature_dim"], NUM_CLASSES).to("cuda:0")
head.load_state_dict(state["head_state_dict"])
head.eval()

# 6. Predict
@torch.no_grad()
def predict(audio_path):
    wav, _ = librosa.load(audio_path, sr=16000, mono=True)
    mel = processor(text="<|audio|>", audios=[wav], return_tensors="pt")
    audio_data = mel["audio_data"].to(moss_model.device, dtype=moss_model.dtype)

    enc_out = moss_model.audio_encoder(audio_data, output_hidden_states=True,
                                        return_dict=True)
    projected = []
    for merger_idx, layer_id in enumerate([10, 21, 31]):
        h = enc_out.hidden_states[layer_id]
        projected.append(moss_model.deepstack_audio_merger_list[merger_idx](h))
    agg = torch.stack(projected, dim=0).sum(dim=0)
    feat = agg.mean(dim=1).squeeze(0).float()

    logits = head(feat.unsqueeze(0))
    probs = torch.softmax(logits, dim=-1).squeeze(0)
    pred_id = int(probs.argmax())
    return {
        "intent": ID2LABEL[pred_id],
        "confidence": float(probs[pred_id]),
        "top3": [(ID2LABEL[i], float(probs[i]))
                 for i in probs.topk(3).indices.cpu().tolist()],
    }

result = predict("my_audio.wav")
print(result)
# Example output:
# {'intent': 'weather_query',
#  'confidence': 0.84,
#  'top3': [('weather_query', 0.84), ('news_query', 0.09), ('general_quirky', 0.03)]}

---

Detailed Inference (production-quality wrapper)

A more complete wrapper with audio validation, multiple input types, and structured output:

import torch
import torch.nn as nn
import librosa
import numpy as np
import sys, os
from pathlib import Path
from typing import Union, List, Dict
from huggingface_hub import hf_hub_download


class MossAudioIntentClassifier:
    """Tunisian Arabic intent classifier using frozen MOSS-Audio + DeepStack + head."""

    def __init__(
        self,
        moss_audio_repo_path: str = "/content/MOSS-Audio",
        classifier_repo: str = "FatimahEmadEldin/moss-audio-tn-intent",
        moss_base_repo: str = "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
        device: str = "cuda:0",
        moss_weights_dir: str = None,
    ):
        """
        Args:
            moss_audio_repo_path: Local path to a clone of github.com/OpenMOSS/MOSS-Audio
            classifier_repo: HF repo holding the trained classifier head (best.pt)
            moss_base_repo: HF repo for the MOSS-Audio base model
            device: 'cuda:0' / 'cuda' / 'cpu'
            moss_weights_dir: Optional local cache dir for MOSS weights
        """
        self.device = device

        # 1. Import MOSS-Audio modules from local clone
        if moss_audio_repo_path not in sys.path:
            sys.path.insert(0, moss_audio_repo_path)
        from src.modeling_moss_audio import MossAudioModel
        from src.processing_moss_audio import MossAudioProcessor

        # 2. Apply idempotent Whisper kwarg-stripping patch
        # (MOSS-Audio uses WhisperEncoderLayer but passes kwargs newer
        #  transformers versions don't accept)
        self._patch_whisper()

        # 3. Load (or cache) MOSS-Audio base model
        if moss_weights_dir and os.path.exists(os.path.join(moss_weights_dir, "config.json")):
            base_path = moss_weights_dir
        else:
            base_path = moss_base_repo
        print(f"📥 Loading frozen MOSS-Audio from {base_path}")
        self.moss_model = MossAudioModel.from_pretrained(
            base_path, trust_remote_code=True, dtype="auto",
            device_map=device, attn_implementation="sdpa",
        ).eval()
        for p in self.moss_model.parameters():
            p.requires_grad = False

        self.processor = MossAudioProcessor.from_pretrained(
            base_path, trust_remote_code=True, enable_time_marker=True,
        )

        # 4. Download + load classifier head
        print(f"📥 Loading classifier head from {classifier_repo}")
        head_path = hf_hub_download(repo_id=classifier_repo, filename="best.pt")
        state = torch.load(head_path, weights_only=False, map_location=device)

        self.feature_dim = state["feature_dim"]
        self.num_classes = state["num_classes"]
        self.ID2LABEL = state["ID2LABEL"]
        self.LABEL2ID = state["LABEL2ID"]
        self.deepstack_layer_ids = state["config"]["DEEPSTACK_LAYER_IDS"]
        self.aggregation = state["config"]["AGGREGATION"]

        self.head = self._build_head(self.feature_dim, self.num_classes).to(device)
        self.head.load_state_dict(state["head_state_dict"])
        self.head.eval()

        print(f"✅ Ready. {self.num_classes} classes, "
              f"DeepStack layers {self.deepstack_layer_ids}, agg={self.aggregation}")...