Upload folder using huggingface_hub

.gitattributes

@@ -35,3 +35,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 adviser/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 embedder/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+figures/intro.png filter=lfs diff=lfs merge=lfs -text
+figures/method.png filter=lfs diff=lfs merge=lfs -text
+figures/mmeb-v2.png filter=lfs diff=lfs merge=lfs -text
+figures/zero-shot.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,160 @@
+---
+license: mit
+datasets:
+- TIGER-Lab/MMEB-train
+language:
+- en
+base_model:
+- Qwen/Qwen3-VL-4B-Instruct
+library_name: transformers
+tags:
+- Retrieval
+- Multimodal
+- Embedding
+- Chain-of-Thought
+- Reinforcement-Learning
+pipeline_tag: image-text-to-text
+---
+
+<div align="center">
+
+<h1>UniME-R1-4B: Learning from Failures for Unified Multimodal Retrieval</h1>
+
+<a href="#">Zelong Sun*</a>,</span>
+<a href="#">Jun Wang*</a>,</span>
+<a href="#">Kaicheng Yang</a>,</span>
+<a href="#">Tiancheng Gu</a>,</span>
+Ziyong Feng,</span>
+<a href="#">Zhiwu Lu</a></span>
+
+Glint Lab
+
+[![GitHub](https://img.shields.io/badge/⭐-GitHub-black?logo=github)](https://github.com/deepglint/UniME-R1)
+[![Paper](https://img.shields.io/badge/📄-Paper-b31b1b.svg)]()
+[![Model](https://img.shields.io/badge/🤗-UniME--R1_Models-yellow)](https://huggingface.co/DeepGlint-AI)
+</div>
+
+UniME-R1 is an **Embedder–Adviser** framework that learns to reason over *retrieved candidates* (not the query alone) and generate **Retrieval-Centric Chain-of-Thought (RC-CoT)** to correct retrieval failures. This repository ships the **4B-scale** pair: a Qwen3-VL-4B dual-mode Embedder and the Qwen3-VL-4B retrieval-aware Adviser — the best-performing configuration in the paper.
+
+## 💡 Highlights
+
+- **Retrieval-Centric CoT (RC-CoT)** — The Adviser analyzes the *actual* top-k retrieved candidates to diagnose model-specific confusion, then emits `<cot_focus>` (which discriminative cues are missing) and `<cot_answer>` (a concise corrected query) to redirect retrieval.
+
+<div align="center">
+  <img src="figures/intro.png" width="95%">
+</div>
+
+- **Dual-Mode Embedder** — A single VLM backbone produces two embeddings via special tokens: `<dis_emb>` for fast discriminative retrieval and `<gen_emb>` for RC-CoT-enhanced re-retrieval. Candidates are encoded **once** with `<dis_emb>` and reused across both paths — no candidate-side CoT, no index rebuilding.
+- **Adaptive Rerank-or-Retrieve** — The Adviser predicts whether a match exists in the top-k set. If yes, it reranks; if not, it appends RC-CoT to the query and re-retrieves over the full corpus.
+- **Retrieval-Oriented GRPO** — The Adviser is optimized with a 4-component reward (format / NDCG rerank / CoT-embedding quality / judge decision) that calls a **frozen Embedder API** to score the Adviser's CoT against mined hard negatives, so the RL signal reflects real end-to-end retrieval quality.
+
+<div align="center">
+  <img src="figures/method.png" width="100%">
+</div>
+
+## 🧱 Model Components
+
+This release contains two sub-directories:
+
+| Component | Backbone | Format | Notes |
+|-----------|----------|--------|-------|
+| **Embedder** (`embedder/`) | Qwen3-VL-4B-Instruct | LoRA (DoRA, r=16, α=64) + `new_token_embeddings.pt` | Adds `<dis_emb>` / `<gen_emb>` tokens; pooling at the special-token position |
+| **Adviser** (`adviser/`) | Qwen3-VL-4B-Instruct | Full merged weights (bf16) | Outputs 5 structured XML fields; served with vLLM |
+
+```
+├── adviser/                       # Qwen3-VL-4B Adviser (full weights, vLLM-ready)
+│   ├── model-0000{1,2}-of-00002.safetensors
+│   ├── config.json / generation_config.json
+│   ├── preprocessor_config.json / video_preprocessor_config.json
+│   └── tokenizer.json / chat_template.jinja
+└── embedder/                      # Qwen3-VL-4B Embedder (LoRA adapter)
+    ├── adapter_config.json
+    ├── adapter_model.safetensors
+    ├── new_token_embeddings.pt    # trained <dis_emb>/<gen_emb> embeddings
+    ├── config.json
+    └── preprocessor_config.json / video_preprocessor_config.json
+```
+
+> The Embedder is released as a **PEFT/LoRA adapter** — load it on top of `Qwen/Qwen3-VL-4B-Instruct`. The two special tokens (`<dis_emb>`=151670, `<gen_emb>`=151669) and their embeddings are provided via `added_tokens.json` + `new_token_embeddings.pt`.
+
+The Adviser always emits five structured fields:
+
+| Field | Role |
+|-------|------|
+| `<rerank_think>` | Candidate-by-candidate relevance analysis |
+| `<rerank_list>` | Predicted candidate ordering (JSON array, 0-based) |
+| `<rerank_judge>` | Path decision: best-match ID, or `-1` if none matches |
+| `<cot_focus>` | Retrieval-failure diagnosis — which discriminative cues are confused |
+| `<cot_answer>` | Concise refined query text for re-retrieval |
+
+## 🚀 Quick Start
+
+```bash
+git clone https://github.com/deepglint/UniME-R1.git
+cd UniME-R1
+```
+
+```bash
+conda create -n vlm2vec python=3.10 -y
+conda activate vlm2vec
+
+# Install torch matching your CUDA first, e.g.:
+# pip install torch==2.5.1 torchvision --index-url https://download.pytorch.org/whl/cu121
+# pip install flash-attn==2.7.3 --no-build-isolation
+
+pip install -r Embedder/requirements.txt
+```
+
+### 🔍 Embedder-only evaluation (direct `<dis_emb>` retrieval)
+
+```bash
+cd Embedder
+bash shell/eval/eval.sh image "../models/UniME-R1-4B/embedder"
+```
+
+### 🎯 Full Adviser evaluation (rerank + RC-CoT iterative retrieval)
+
+Serve the Adviser via vLLM, then run the unified evaluation:
+
+```bash
+vllm serve models/UniME-R1-4B/adviser --tensor-parallel-size 8 --port 9000
+
+cd RL/eval
+export EMBEDDER_CHECKPOINT="../../models/UniME-R1-4B/embedder"
+export ADVISER_MODEL="Qwen3-VL-4B-Instruct"
+export ADVISER_URL="http://127.0.0.1:9000/v1"
+bash eval.sh image      # image | visdoc | video | uvrb | image_caption
+```
+
+## 📊 Results
+
+### 🏆 MMEB-V2
+
+UniME-R1 achieves the best overall performance at both model scales. Notably, the 2B model already outperforms all medium-size (4B–7B) baselines, indicating the gains stem from the framework rather than model scale alone.
+
+<div align="center">
+  <img src="figures/mmeb-v2.png" width="95%">
+</div>
+
+### 🌈 Zero-shot General Retrieval
+
+<div align="center">
+  <img src="figures/zero-shot.png" width="90%">
+</div>
+
+## 🖊️ Citation
+
+If you find this repository useful, please use the following BibTeX entry for citation.
+
+```bibtex
+@misc{unime-r1,
+      title  = {Learning from Failures: Retrieval-Centric CoT via Hard Negatives for Unified Multimodal Retrieval},
+      author = {Sun, Zelong and Wang, Jun and Yang, Kaicheng and Gu, Tiancheng and Feng, Ziyong and Lu, Zhiwu},
+      year   = {2025},
+      url    = {https://github.com/deepglint/UniME-R1}
+}
+```
+
+<div align="center">
+⭐ Don't forget to star this repository if you find it helpful!
+</div>