README.md

---
language: en
license: mit
library_name: transformers
pipeline_tag: feature-extraction
tags:
  - audio-text-retrieval
  - audio-retrieval
  - clap
  - multimodal-llm
  - lora
  - acl2026
base_model:
  - Qwen/Qwen2.5-Omni-3B
---

# OEA-Qwen3B-AC

`OEA-Qwen3B-AC` is a **LoRA + projection-head checkpoint** for the **Omni-Embed-Audio (OEA)** retrieval encoder presented in:

> **Omni-Embed-Audio: Leveraging Multimodal LLMs for Robust Audio-Text Retrieval**
> HaeJun Yoo, Yongseop Shin, Insung Lee, Myoung-Wan Koo, Du-Seong Chang (Sogang University)
> *Proceedings of ACL 2026 (Oral)*
> Code & UIQ benchmark: <https://github.com/JudeJiwoo/Omni-Embed-Audio>
> Web demo: <https://omni-embed-audio.github.io>

## Model summary

| Field | Value |
|-------|-------|
| Base model | [`Qwen/Qwen2.5-Omni-3B`](https://huggingface.co/Qwen/Qwen2.5-Omni-3B) (Qwen2.5-Omni-3B, ~3B params) |
| Trained on | AudioCaps (train split) |
| Embedding dim | 512 (L2-normalized) |
| Audio sample rate | 16 kHz mono |
| Trainable parameters | LoRA adapters + 2 projection heads (~11–16M) |
| Backbone | **Frozen** |
| Checkpoint file | `step_40.pt` (PyTorch state dict) |
| License (this checkpoint) | MIT — the underlying base model is governed by its own license |

## What's inside `step_40.pt`

```python
ckpt = torch.load("step_40.pt", map_location="cpu")
ckpt.keys()
# dict_keys(['lora_state_dict', 'audio_head', 'text_head', ...])
```

- `lora_state_dict` — LoRA adapters (r=16, α=32, dropout=0.05) attached to
  `q_proj, k_proj, v_proj, o_proj, qkv, out_proj` of the Qwen2.5-Omni-3B backbone.
- `audio_head`, `text_head` — modality-specific 512-d `ProjectionHead`
  (`Linear → Dropout → LayerNorm → L2-normalize`).

The base model weights are **not** redistributed here; they are loaded from
[`Qwen/Qwen2.5-Omni-3B`](https://huggingface.co/Qwen/Qwen2.5-Omni-3B) on first use.

## Quick start

Install dependencies (see the repository for the full pinned list):

```bash
pip install torch torchaudio transformers peft huggingface_hub soundfile librosa
```

Minimal encoding example (text & audio → 512-d L2-normalized embeddings):

```python
import torch
from types import SimpleNamespace
from huggingface_hub import hf_hub_download

# Cloned from https://github.com/JudeJiwoo/Omni-Embed-Audio
from AudioRetrieval.models.omni_embed_adapter import OmniEmbedAdapter
from AudioRetrieval.training.oea.train_omniembed_lora import ProjectionHead, attach_lora

device = "cuda" if torch.cuda.is_available() else "cpu"

ckpt_path = hf_hub_download("JudeJiwoo/OEA-Qwen3B-AC", filename="step_40.pt")

adapter = OmniEmbedAdapter(
    repo_id="Qwen/Qwen2.5-Omni-3B",
    device=device,
    passage_prefix="passage:",
    query_prefix="query:",
)

lora_cfg = SimpleNamespace(
    lora_rank=16, lora_alpha=32, lora_dropout=0.05,
    lora_targets=["q_proj","k_proj","v_proj","o_proj","qkv","out_proj"],
)
peft_model = attach_lora(adapter.get_underlying_model(), lora_cfg)
adapter.set_underlying_model(peft_model)

ckpt = torch.load(ckpt_path, map_location=device)
peft_model.load_state_dict(ckpt["lora_state_dict"], strict=False)

hidden = peft_model.config.text_config.hidden_size
audio_head = ProjectionHead(hidden, 512, 0.1).to(device).eval()
text_head  = ProjectionHead(hidden, 512, 0.1).to(device).eval()
audio_head.load_state_dict(ckpt["audio_head"])
text_head.load_state_dict(ckpt["text_head"])

with torch.inference_mode():
    a = adapter.encode_audio(["sample.wav"])
    t = adapter.encode_text(["A clock ticks once a second as it runs."])
    a_emb = audio_head(torch.from_numpy(a).to(device)).cpu().numpy()
    t_emb = text_head(torch.from_numpy(t).to(device)).cpu().numpy()
print(a_emb.shape, t_emb.shape)   # (1, 512) (1, 512)
print((t_emb @ a_emb.T).item())   # cosine similarity (already L2-normalized)
```

A ready-to-run version is at
[`examples/encode_example.py`](https://github.com/JudeJiwoo/Omni-Embed-Audio/blob/main/examples/encode_example.py)
(or the matching `.ipynb`) in the public repo — it encodes five bundled Clotho
clips and one query per UIQ type out of the box.

## Related checkpoints

| Repo | Base model | Trained on |
|------|-----------|-----------|
| [JudeJiwoo/OEA-Qwen3B-Cl](https://huggingface.co/JudeJiwoo/OEA-Qwen3B-Cl) | Qwen2.5-Omni-3B | Clotho |
| [JudeJiwoo/OEA-Qwen3B-AC](https://huggingface.co/JudeJiwoo/OEA-Qwen3B-AC) | Qwen2.5-Omni-3B | AudioCaps |
| [JudeJiwoo/OEA-Qwen7B-Cl](https://huggingface.co/JudeJiwoo/OEA-Qwen7B-Cl) | Qwen2.5-Omni-7B | Clotho |
| [JudeJiwoo/OEA-Qwen7B-AC](https://huggingface.co/JudeJiwoo/OEA-Qwen7B-AC) | Qwen2.5-Omni-7B | AudioCaps |
| [JudeJiwoo/OEA-Nemo3B-Cl](https://huggingface.co/JudeJiwoo/OEA-Nemo3B-Cl) | Omni-Embed-Nemotron-3B | Clotho |
| [JudeJiwoo/OEA-Nemo3B-AC](https://huggingface.co/JudeJiwoo/OEA-Nemo3B-AC) | Omni-Embed-Nemotron-3B | AudioCaps |

`-Cl` checkpoints are recommended for evaluation on **Clotho**, and `-AC`
checkpoints for **AudioCaps** (Section 4 of the paper).

## Intended use & limitations

**Intended use.** Open-vocabulary text-to-audio / audio-to-text / text-to-text
retrieval; UIQ-style robustness studies (questions, imperatives, paraphrases,
keyword tags, exclusion queries). Both encoders share a single multimodal-LLM
backbone, so OEA is also useful for embedding-space analyses across modalities.

**Out of scope.** Generative captioning, ASR, music transcription, speaker
identification, or any safety-critical downstream task. The model was trained
on environmental-sound captioning data; it is **not** evaluated on speech-,
music-, or biomedical-audio retrieval.

**Limitations.** Inherits biases of the Qwen2.5-Omni-3B backbone and of the
AudioCaps caption distribution. Performance on languages other than English
and on out-of-distribution audio (long-form speech, music, multilingual
content) is not measured.

## License

This LoRA + projection-head checkpoint is released under the **MIT License**.
The underlying base model (`Qwen/Qwen2.5-Omni-3B`) is governed by **its own license** —
review the upstream model card before redistribution.

## Citation

```bibtex
@inproceedings{yoo2026omniembedaudio,
  title     = {Omni-Embed-Audio: Leveraging Multimodal LLMs for Robust Audio-Text Retrieval},
  author    = {Yoo, HaeJun and Shin, Yongseop and Lee, Insung and Koo, Myoung-Wan and Chang, Du-Seong},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (ACL)},
  note      = {Oral presentation},
  year      = {2026}
}
```

## Acknowledgments

This work was partly supported by the Institute of Information & Communications
Technology Planning & Evaluation (IITP) grant funded by the Korea government
(MSIT) (No. RS-2025-25441313, Professional AI Talent Development Program for
Multimodal AI Agents, Contribution: 50%). This research was also supported by
the MSIT, Korea, under the Top-Tier AI Global HRD invitation program
(RS-2025-25461932) supervised by the IITP.