README.md
8.1 KB · 186 lines · markdown Raw
1 ---
2 license: apache-2.0
3 language:
4 - en
5 base_model: Qwen/Qwen2.5-7B-Instruct
6 tags:
7 - retrieval
8 - query-rewriting
9 - reinforcement-learning
10 ---
11
12 <h1 align="center">inf-retriever-v1-pro</h1>
13
14 <p align="center">
15 <a href="https://brightbenchmark.github.io/"><img src="https://img.shields.io/badge/BRIGHT_Benchmark-Rank_1st-8A2BE2" alt="Rank"></a>
16 <a href="https://huggingface.co/infly/inf-retrieve-v1-pro"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face-INF--Retrieve--v1--pro-blue" alt="Hugging Face"></a>
17 <a href="https://opensource.org/licenses/Apache-2.0"><img src="https://img.shields.io/badge/License-Apache--2.0-green.svg" alt="License"></a>
18 </p>
19
20 ## 📖 Overview
21
22 **inf-retriever-v1-pro** is a specialized retrieval component of the **INF-X-Retriever** framework, designed to distill the core retrieval intent from complex, verbose, or reasoning-intensive queries. Built upon the **inf-retriever-v1** foundation and further trained to serve as the retriever within a RAG (retrieval-augmented generation) system, it transforms raw user queries into concise, search-optimized queries for dense retrieval systems.
23
24 In our experiments, a single canonical query-writing prompt was applied across all datasets to ensure consistency and reproducibility.
25 ```bash
26 task = 'Given a web search query, retrieve relevant passages that answer the query'
27 ```
28
29 This model is a key enabler for **INF-X-Retriever**'s state-of-the-art performance, currently holding the **No. 1 position** on the [BRIGHT Benchmark](https://brightbenchmark.github.io/) (as of Dec 17, 2025).
30
31 For more details on the full framework, please visit the [INF-X-Retriever Repository](https://github.com/yaoyichen/INF-X-Retriever).
32
33 ---
34
35 ### Requirements
36
37 ```bash
38 transformers==4.51.0
39 ```
40
41 ### Usage
42
43 ### Sentence Transformers
44 ```python
45 from sentence_transformers import SentenceTransformer
46
47 model = SentenceTransformer("infly/inf-retriever-v1", trust_remote_code=True)
48 # In case you want to reduce the maximum length:
49 model.max_seq_length = 8192
50
51 queries = [
52 "how much protein should a female eat",
53 "summit define",
54 ]
55 documents = [
56 "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
57 "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
58 ]
59
60 query_embeddings = model.encode(queries, prompt_name="query")
61 document_embeddings = model.encode(documents)
62
63 scores = (query_embeddings @ document_embeddings.T) * 100
64 print(scores.tolist())
65 # [[91.46116638183594, 76.9832992553711], [70.7034683227539, 87.15817260742188]]
66 ```
67
68 ### Transformers
69 ```python
70 import torch
71 import torch.nn.functional as F
72
73 from torch import Tensor
74 from transformers import AutoTokenizer, AutoModel
75
76
77 def last_token_pool(last_hidden_states: Tensor,
78 attention_mask: Tensor) -> Tensor:
79 left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
80 if left_padding:
81 return last_hidden_states[:, -1]
82 else:
83 sequence_lengths = attention_mask.sum(dim=1) - 1
84 batch_size = last_hidden_states.shape[0]
85 return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
86
87
88 def get_detailed_instruct(task_description: str, query: str) -> str:
89 return f'Instruct: {task_description}\nQuery: {query}'
90
91
92 # Each query must come with a one-sentence instruction that describes the task
93 task = 'Given a web search query, retrieve relevant passages that answer the query'
94 queries = [
95 get_detailed_instruct(task, 'how much protein should a female eat'),
96 get_detailed_instruct(task, 'summit define')
97 ]
98 # No need to add instruction for retrieval documents
99 documents = [
100 "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
101 "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
102 ]
103 input_texts = queries + documents
104
105 tokenizer = AutoTokenizer.from_pretrained('infly/inf-retriever-v1', trust_remote_code=True)
106 model = AutoModel.from_pretrained('infly/inf-retriever-v1', trust_remote_code=True)
107
108 max_length = 8192
109
110 # Tokenize the input texts
111 batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt')
112 outputs = model(**batch_dict)
113 embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
114
115 # normalize embeddings
116 embeddings = F.normalize(embeddings, p=2, dim=1)
117 scores = (embeddings[:2] @ embeddings[2:].T) * 100
118 print(scores.tolist())
119 # [[91.46114349365234, 76.98332214355469], [70.7035140991211, 87.158203125]]
120 ```
121
122 ---
123
124 ## Performance
125
126 **INF-X-Retriever** achieves state-of-the-art results on the [BRIGHT Benchmark](https://brightbenchmark.github.io/) (as of Dec 20, 2025).
127
128 The **BRIGHT** (Benchmark for Reasoning-Intensive Grounded HT) is a rigorous text retrieval benchmark designed to evaluate the capability of retrieval models in handling questions that require intensive reasoning and cross-document synthesis. Collected from real-world sources such as StackExchange, competitive programming platforms, and mathematical competitions, it comprises complex queries spanning diverse domains like mathematics, coding, biology, economics, and robotics.
129
130 ### Short document
131
132 #### Overall & Category Performance
133
134 | Model | **Avg ALL** | **StackExchange** | **Coding** | **Theorem-based** |
135 |:---|:---:|:---:|:---:|:---:|
136 | **INF-X-Retriever** | **63.4** | **68.3** | **55.3** | **57.7** |
137 | DIVER (v3) | 46.8 | 51.8 | 39.9 | 39.7 |
138 | BGE-Reasoner-0928 | 46.4 | 52.0 | 35.3 | 40.7 |
139 | LATTICE | 42.1 | 51.6 | 26.9 | 30.0 |
140 | ReasonRank | 40.8 | 46.9 | 27.6 | 35.5 |
141 | XDR2 | 40.3 | 47.1 | 28.5 | 32.1 |
142
143 #### Detailed Results Across 12 Datasets
144
145 | Model | Avg | Bio. | Earth. | Econ. | Psy. | Rob. | Stack. | Sus. | Leet. | Pony | AoPS | TheoQ. | TheoT. |
146 | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
147 | **INF-X-Retriever** | **63.4** | **79.8** | **70.9** | **69.9** | **73.3** | **57.7** | **64.3** | **61.9** | **56.1** | **54.5** | **51.9** | **53.1** | **67.9** |
148 | DIVER (v3) | 46.8 | 66.0 | 63.7 | 42.4 | 55.0 | 40.6 | 44.7 | 50.4 | 32.5 | 47.3 | 17.2 | 46.4 | 55.6 |
149 | BGE-Reasoner-0928 | 46.4 | 68.5 | 66.4 | 40.6 | 53.1 | 43.2 | 44.1 | 47.8 | 29.0 | 41.6 | 17.2 | 46.5 | 58.4 |
150 | LATTICE | 42.1 | 64.4 | 62.4 | 45.4 | 57.4 | 47.6 | 37.6 | 46.4 | 19.9 | 34.0 | 12.0 | 30.1 | 47.8 |
151 | ReasonRank | 40.8 | 62.7 | 55.5 | 36.7 | 54.6 | 35.7 | 38.0 | 44.8 | 29.5 | 25.6 | 14.4 | 42.0 | 50.1 |
152 | XDR2 | 40.3 | 63.1 | 55.4 | 38.5 | 52.9 | 37.1 | 38.2 | 44.6 | 21.9 | 35.0 | 15.7 | 34.4 | 46.2 |
153
154 ### Long document
155
156 #### Detailed Results Across 8 Datasets
157
158 | Model | Avg | Bio. | Earth. | Econ. | Pony | Psy. | Rob. | Stack. | Sus. |
159 | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
160 | **INF-X-Retriever** | **54.6** | **73.2** | **59.6** | **69.3** | **12.1** | **74.3** | **55.9** | **27.8** | **64.8** |
161 | inf-retriever-v1-pro | 30.5 | 44.1 | 42.2 | 31.4 | 0.4 | 43.1 | 20.8 | 21.4 | 41.0 |
162
163
164 ---
165
166 ## 🖊️ Citation
167
168 If you find this model useful, please consider citing our work:
169
170 ```bibtex
171 @misc{inf-x-retriever-2025,
172 title = {INF-X-Retriever},
173 author = {Yichen Yao, Jiahe Wan, Yuxin Hong, Mengna Zhang, Junhan Yang, Zhouyu Jiang, Qing Xu, Kuan Lu, Yinghui Xu, Wei Chu, Emma Wang, Yuan Qi},
174 year = {2025},
175 url = {https://yaoyichen.github.io/INF-X-Retriever},
176 publisher = {GitHub repository}
177 }
178 ```
179
180 ---
181
182 ## 📬 Contact
183
184 Email: [eason.yyc@inftech.ai](mailto:eason.yyc@inftech.ai)
185
186