--- language: en license: mit library_name: transformers pipeline_tag: feature-extraction tags: - audio-text-retrieval - audio-retrieval - clap - multimodal-llm - lora - acl2026 base_model: - nvidia/omni-embed-nemotron-3b --- # OEA-Nemo3B-AC `OEA-Nemo3B-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: > Web demo: ## Model summary | Field | Value | |-------|-------| | Base model | [`nvidia/omni-embed-nemotron-3b`](https://huggingface.co/nvidia/omni-embed-nemotron-3b) (Omni-Embed-Nemotron-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 Omni-Embed-Nemotron-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 [`nvidia/omni-embed-nemotron-3b`](https://huggingface.co/nvidia/omni-embed-nemotron-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-Nemo3B-AC", filename="step_40.pt") adapter = OmniEmbedAdapter( repo_id="nvidia/omni-embed-nemotron-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 Omni-Embed-Nemotron-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 (`nvidia/omni-embed-nemotron-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.