Create README.md

README.md

@@ -0,0 +1,396 @@
+---
+license: apache-2.0
+base_model:
+- Qwen/Qwen3-0.6B
+tags:
+- transformers
+- sentence-similarity
+- feature-extraction
+- text-embeddings-inference
+- reranking
+pipeline_tag: feature-extraction
+---
+
+<!-- <div align="center">
+<p align="center">
+  <img src="assets/overall.jpg" width="50%" height="50%" />
+</p>
+</div> -->
+
+<div align="center">
+<h1>E2Rank: Your Text Embedding can Also be an Effective and Efficient Listwise Reranker</h1>
+
+<!-- <a href='https://arxiv.org/pdf/2501.07572'><img src='https://img.shields.io/badge/Paper-arXiv-red'></a>
+<img src="https://img.shields.io/github/stars/Alibaba-NLP/E2Rank?color=yellow" alt="Stars">
+<a href='https://huggingface.co/papers/2501.07572'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Discussion-orange'></a>
+<a href='https://huggingface.co/collections/callanwu/webwalker-677f6527407edfda44098b09'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Colloectionss-blue'></a>
+<a href='https://huggingface.co/datasets/callanwu/WebWalkerQA'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Datasets-green'></a> -->
+
+<!-- **Authors:** -->
+
+<div class="is-size-5 publication-authors">
+    <span class="author-block">
+        <a href="https://liuqi6777.github.io">Qi Liu</a><sup>1,2</sup>,</span>
+    <span class="author-block">
+        Yanzhao Zhang<sup>2</sup>,</span>
+    <span class="author-block">
+        Mingxin Li<sup>2</sup>,
+    </span>
+    <span class="author-block">
+        Dingkun Long<sup>2</sup>,
+    </span>
+    <span class="author-block">
+        Pengjun Xie<sup>2</sup>,
+    </span>
+    <span class="author-block">
+        Jiaxin Mao<sup>1</sup>,
+    </span>
+</div>
+<!-- **Affiliations:** -->
+<div class="is-size-5 publication-authors">
+    <span><sup>1</sup><i>GSAI, Renmin University of China,</i></span>
+    <span><sup>2</sup><i>Tongyi Lab, Alibaba Group</i></span>
+</div>
+
+
+<a href="https://Alibaba-NLP.github.io/E2Rank/">[🤖 Website]</a> | 
+<a href="https://arxiv.org/pdf/2501.07572">[📄 Arxiv Paper]</a> | 
+<a href="https://huggingface.co/collections/Alibaba-NLP/e2rank">[🤗 Huggingface Collection]</a> |
+<a href="# 🚩Citation">[🚩 Citation]</a>
+
+</div>
+
+# 📌 Introduction
+
+We introduce E2Rank, 
+meaning **E**fficient **E**mbedding-based **Rank**ing
+(also meaning **Embedding-to-Rank**), 
+which extends a single text embedding model
+to perform both high-quality retrieval and listwise reranking, 
+thereby achieving strong effectiveness with remarkable efficiency. 
+
+By applying cosine similarity between the query and
+document embeddings as a unified ranking function, the listwise ranking prompt,
+which is constructed from the original query and its candidate documents, serves
+as an enhanced query enriched with signals from the top-K documents, akin to
+pseudo-relevance feedback (PRF) in traditional retrieval models. This design 
+preserves the efficiency and representational quality of the base embedding model
+while significantly improving its reranking performance. 
+
+Empirically, E2Rank achieves state-of-the-art results on the BEIR reranking benchmark 
+and demonstrates competitive performance on the reasoning-intensive BRIGHT benchmark,
+with very low reranking latency. We also show that the ranking training process
+improves embedding performance on the MTEB benchmark. 
+Our findings indicate that a single embedding model can effectively unify retrieval and reranking,
+offering both computational efficiency and competitive ranking accuracy.
+
+**Our work highlights the potential of single embedding models to serve as unified retrieval-reranking engines, offering a practical, efficient, and accurate alternative to complex multi-stage ranking systems.**
+
+
+# 🚀 Quick Start
+
+## Model List
+
+| Supported Task              | Model Name           | Size | Layers | Sequence Length | Embedding Dimension | Instruction Aware |
+|-----------------------------|----------------------|------|--------|-----------------|---------------------|-------------------|
+| **Embedding + Reranking**   | [Alibaba-NLP/E2Rank-0.6B](https://huggingface.co/Alibaba-NLP/E2Rank-0.6B) | 0.6B | 28     | 32K             | 1024                | Yes            |
+| **Embedding + Reranking**   | [Alibaba-NLP/E2Rank-4B](https://huggingface.co/Alibaba-NLP/E2Rank-4B)     | 4B   | 36     | 32K             | 2560                | Yes            |
+| **Embedding + Reranking**   | [Alibaba-NLP/E2Rank-8B](https://huggingface.co/Alibaba-NLP/E2Rank-8B)     | 8B   | 36     | 32K             | 4096                | Yes            |
+| Embedding Only              | [Alibaba-NLP/E2Rank-0.6B-Embedding-Only](https://huggingface.co/Alibaba-NLP/E2Rank-0.6B-Embedding-Only) | 0.6B | 28     | 32K             | 1024                | Yes         |
+| Embedding Only              | [Alibaba-NLP/E2Rank-0.6B-Embedding-Only](https://huggingface.co/Alibaba-NLP/E2Rank-4B-Embedding-Only)   | 4B   | 36     | 32K             | 2560                | Yes         |
+| Embedding Only              | [Alibaba-NLP/E2Rank-0.6B-Embedding-Only](https://huggingface.co/Alibaba-NLP/E2Rank-8B-Embedding-Only)   | 8B   | 36     | 32K             | 4096                | Yes         |
+
+
+> **Note**:
+> - `Embedding Only` indicates that the model is trained only with the constrative learning and support embedding tasks, while `Embedding + Reranking` indicates the **full E2Rank model** trained with both embedding and reranking objectives (for more detals, please refer to the [paper]()). 
+> - `Instruction Aware` notes whether the model supports customizing the input instruction according to different tasks.
+<!-- > - For `Listwise Reranking` models, they are supervised fine-tuned from the Qwen3 Models in the paradigm of RankGPT and support only the reranking task. -->
+
+## Usage
+
+### Embedding Model
+
+The usage of E2Rank as an embedding model is similar to [Qwen3-Embedding](https://github.com/QwenLM/Qwen3-Embedding). The only difference is that Qwen3-Embedding will automatically append an EOS token, while E2Rank requires users to manully append the special token `<|endoftext|>` at the end of each input text.
+
+
+**vLLM Usage (recommended)**
+
+```python
+# Requires vllm>=0.8.5
+import torch
+import vllm
+from vllm import LLM
+from vllm.config import PoolerConfig
+
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'Instruct: {task_description}\nQuery:{query}'
+
+# Each query must come with a one-sentence instruction that describes the task
+task = 'Given a web search query, retrieve relevant passages that answer the query'
+
+queries = [
+    get_detailed_instruct(task, 'What is the capital of China?'),
+    get_detailed_instruct(task, 'Explain gravity')
+]
+# No need to add instruction for retrieval documents
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+]
+input_texts = queries + documents
+input_texts = [t + "<|endoftext|>" for t in input_texts]
+
+model = LLM(
+    model="Alibaba-NLP/E2Rank-0.6B",
+    task="embed", 
+    override_pooler_config=PoolerConfig(pooling_type="LAST", normalize=True)
+)
+
+outputs = model.embed(input_texts)
+embeddings = torch.tensor([o.outputs.embedding for o in outputs])
+scores = (embeddings[:2] @ embeddings[2:].T)
+print(scores.tolist())
+# [[0.5958386659622192, 0.030148349702358246], [0.060259245336055756, 0.5595865249633789]]
+```
+
+<details>
+<summary><b>Transformers Usage</b></summary>
+
+```python
+# Requires transformers>=4.51.0
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+
+
+def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'Instruct: {task_description}\nQuery:{query}'
+
+# Each query must come with a one-sentence instruction that describes the task
+task = 'Given a web search query, retrieve relevant passages that answer the query'
+
+queries = [
+    get_detailed_instruct(task, 'What is the capital of China?'),
+    get_detailed_instruct(task, 'Explain gravity')
+]
+# No need to add instruction for retrieval documents
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+]
+input_texts = queries + documents
+input_texts = [t + "<|endoftext|>" for t in input_texts]
+
+tokenizer = AutoTokenizer.from_pretrained('Alibaba-NLP/E2Rank-0.6B', padding_side='left')
+model = AutoModel.from_pretrained('Alibaba-NLP/E2Rank-0.6B')
+
+max_length = 8192
+
+# Tokenize the input texts
+batch_dict = tokenizer(
+    input_texts,
+    padding=True,
+    truncation=True,
+    max_length=max_length,
+    return_tensors="pt",
+)
+batch_dict.to(model.device)
+with torch.no_grad():
+    outputs = model(**batch_dict)
+    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+
+    # normalize embeddings
+    embeddings = F.normalize(embeddings, p=2, dim=1)
+    scores = (embeddings[:2] @ embeddings[2:].T)
+
+print(scores.tolist())
+# [[0.5950675010681152, 0.030417663976550102], [0.061970409005880356, 0.562691330909729]]
+```
+</details>
+
+
+### Reranking
+
+For using E2Rank as a reranker, you only need to perform additional processing on the query by adding (part of) the docs that needs to be reranked to the *listwise prompt*, while the rest is the same as using the embedding model.
+
+**vLLM Usage (recommended)**
+
+```python
+# Requires vllm>=0.8.5
+import torch
+import vllm
+from vllm import LLM
+from vllm.config import PoolerConfig
+
+model = LLM(
+    model="./checkpoints/E2Rank-0.6B",
+    task="embed", 
+    override_pooler_config=PoolerConfig(pooling_type="LAST", normalize=True)
+)
+tokenizer = model.get_tokenizer()
+
+def get_listwise_prompt(task_description: str, query: str, documents: list[str], num_input_docs: int = 20) -> str:
+    input_docs = documents[:num_input_docs]
+    input_docs = "\n".join([f"[{i}] {doc}" for i, doc in enumerate(input_docs, start=1)])
+    messages = [{
+        "role": "user", 
+        "content": f'{task_description}\nDocuments:\n{input_docs}Search Query:{query}'
+    }]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False,
+    )
+    return text
+
+task = 'Given a web search query and some relevant documents, rerank the documents that answer the query:'
+
+queries = [
+    'What is the capital of China?',
+    'Explain gravity'
+]
+
+# No need to add instruction for retrieval documents
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+]
+documents = [doc + "<|endoftext|>" for doc in documents]
+
+pseudo_queries = [
+    get_listwise_prompt(task, queries[0], documents),
+    get_listwise_prompt(task, queries[1], documents)
+]  # no need to add the EOS token here
+
+input_texts = pseudo_queries + documents
+
+outputs = model.embed(input_texts)
+embeddings = torch.tensor([o.outputs.embedding for o in outputs])
+scores = (embeddings[:2] @ embeddings[2:].T)
+print(scores.tolist())
+# [[0.8516960144042969, 0.24043934047222137], [0.33099934458732605, 0.7905282974243164]]
+```
+
+<details>
+<summary><b>Transformers Usage</b></summary>
+
+```python
+# Requires transformers>=4.51.0
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+
+
+tokenizer = AutoTokenizer.from_pretrained('./checkpoints/E2Rank-0.6B', padding_side='left')
+model = AutoModel.from_pretrained('./checkpoints/E2Rank-0.6B')
+
+
+def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+def get_listwise_prompt(task_description: str, query: str, documents: list[str], num_input_docs: int = 20) -> str:
+    input_docs = documents[:num_input_docs]
+    input_docs = "\n".join([f"[{i}] {doc}" for i, doc in enumerate(input_docs, start=1)])
+    messages = [{
+        "role": "user", 
+        "content": f'{task_description}\nDocuments:\n{input_docs}Search Query:{query}'
+    }]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False,
+    )
+    return text
+
+task = 'Given a web search query and some relevant documents, rerank the documents that answer the query:'
+
+queries = [
+    'What is the capital of China?',
+    'Explain gravity'
+]
+
+# No need to add instruction for retrieval documents
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+]
+documents = [doc + "<|endoftext|>" for doc in documents]
+
+pseudo_queries = [
+    get_listwise_prompt(task, queries[0], documents),
+    get_listwise_prompt(task, queries[1], documents)
+]  # no need to add the EOS token here
+
+input_texts = pseudo_queries + documents
+
+
+max_length = 8192
+# Tokenize the input texts
+batch_dict = tokenizer(
+    input_texts,
+    padding=True,
+    truncation=True,
+    max_length=max_length,
+    return_tensors="pt",
+)
+batch_dict.to(model.device)
+with torch.no_grad():
+    outputs = model(**batch_dict)
+    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+
+    # normalize embeddings
+    embeddings = F.normalize(embeddings, p=2, dim=1)
+    scores = (embeddings[:2] @ embeddings[2:].T)
+
+print(scores.tolist())
+# [[0.8513513207435608, 0.24268491566181183], [0.33154672384262085, 0.7923378944396973]]
+```
+</details>
+
+### End-to-end search
+
+Since E2Rank extends a single text embedding model to perform both high-quality retrieval and listwise reranking, you can directly use it to build an end-to-end search system. By reusing the embeddings computed during the retrieval stage, E2Rank only need to compute the pseudo query's embedding and can efficiently rerank the retrieved documents with minimal additional computational overhead.
+
+Example code is comming soon.
+
+
+# 🚩 Citation
+
+If this work is helpful, please kindly cite as:
+
+```bibtext
+
+```
+
+if you have any questions, please feel free to contact via qiliu6777[AT]gmail.com or create an issue.
+
+<!-- ## Star History
+
+<div align="center">
+
+[![Star History Chart](https://api.star-history.com/svg?repos=Alibaba-NLP/E2Rank&type=Date)](https://www.star-history.com/#Alibaba-NLP/WebAgent&Date)
+
+</div> -->