You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: README.md
+58-24Lines changed: 58 additions & 24 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -1,6 +1,6 @@
1
1
# FIRST: Faster Improved Listwise Reranking with Single Token Decoding
2
2
3
-
This repository contains the code for the paper [FIRST: Faster Improved Listwise Reranking with Single Token Decoding](https://arxiv.org/pdf/2406.15657)
3
+
This repository contains the code for the paper [FIRST: Faster Improved Listwise Reranking with Single Token Decoding](https://arxiv.org/pdf/2406.15657) and the reranker code for the paper [CoRNStack: High-Quality Contrastive Data for Better Code Ranking](https://arxiv.org/abs/2412.01007).
4
4
5
5
FIRST is a novel listwise LLM reranking approach leveraging the output logits of the first generated identifier to obtain a ranked ordering of the input candidates directly. FIRST incorporates a learning-to-rank loss during training, prioritizing ranking accuracy for the more relevant passages.
6
6
@@ -22,59 +22,93 @@ pip install beir
22
22
23
23
Before running the scripts below, do
24
24
```
25
-
export REPO_DIR=<path to the llm-reranker directory
25
+
export REPO_DIR=<path to the llm-reranker directory
26
26
```
27
27
28
28
## 1. Retrieval
29
+
### 1a. Text Retrieval
29
30
We use [contriever](https://github.com/facebookresearch/contriever) as the underlying retrieval model. The precomputed query and passage embeddings for BEIR are available [here](https://huggingface.co/datasets/rryisthebest/Contreiever_BEIR_Embeddings).
30
31
31
32
**Note:** If you wish to not run the retrieval yourself, the retrieval results are provided [here](https://drive.google.com/drive/folders/1eMiqwiTVwJy_Zcss7LQF9hQ1aeTFMZUm?usp=sharing) and you can directly jump to [Reranking](#2-reranking)
32
33
33
-
34
34
To run the contriever retrieval using the precomputed encodings
35
35
36
36
```
37
-
bash bash/beir/run_1st_retrieval.sh <Path to folder with BEIR encodings>
37
+
bash bash/run_1st_retrieval.sh <Path to folder with BEIR encodings>
38
38
```
39
39
To get the retrieval scores, run:
40
40
41
41
```
42
-
bash bash/beir/run_eval.sh rank
42
+
bash bash/run_eval.sh rank
43
+
```
44
+
45
+
### 1b. Code Retrieval
46
+
**Note:** If you wish to not run the code retrieval yourself, the code retrieval results are provided [here](https://drive.google.com/drive/folders/1GYI4g7mTVOhsttwDSioOISBZe_KiFEFt?usp=sharing) and you can directly jump to [Reranking](#2-reranking)
47
+
48
+
To get the code retrieval scores, run:
49
+
50
+
```
51
+
bash bash/run_eval.sh rank code
43
52
```
44
53
45
54
## 2. Reranking
46
-
### 2a. Baseline Cross-encoder reranking
55
+
### 2a. Baseline Text Cross-encoder reranking
47
56
48
-
To run the baseline cross encoder re-ranking, run:
57
+
To run the baseline text cross encoder re-ranking, run:
49
58
```
50
-
bash bash/beir/run_rerank.sh
59
+
bash bash/run_rerank.sh
51
60
```
52
-
### 2b. FIRST LLM Reranking
61
+
### 2b. FIRST LLM Reranking - Text
53
62
54
-
To convert the retrieval results to input for LLM reranking, run:
63
+
To convert the retrieval results to input for Text LLM reranking, run:
55
64
56
65
```
57
-
bash bash/beir/run_convert_results.sh
66
+
bash bash/run_convert_results.sh text
58
67
```
59
68
60
69
We provide the trained FIRST reranker [here](https://huggingface.co/rryisthebest/First_Model).
61
70
62
-
To run the FIRST reranking, run:
71
+
To run the FIRST reranking, set RERANK_TYPE="text" in bash/run_rerank_llm.sh and run:
63
72
64
73
```
65
-
bash bash/beir/run_rerank_llm.sh
74
+
bash bash/run_rerank_llm.sh
66
75
```
67
76
68
77
To evaluate the reranking performance, run:
69
78
70
79
```
71
-
bash bash/run_eval.sh rerank
80
+
bash bash/run_eval.sh rerank text
72
81
73
82
```
74
83
**Note:** Set flag --suffix to "llm_FIRST_alpha" for FIRST reranker evaluation or "ce" for cross encoder reranker
75
84
85
+
### 2c. CodeRanker - Code Reranking
86
+
**Note:** CodeRanker currently does not support logit and alpha inference.
87
+
88
+
To convert the code retrieval results to input for Code LLM reranking, run:
89
+
90
+
```
91
+
bash bash/run_convert_results.sh code
92
+
```
93
+
94
+
We provide the trained FIRST reranker [here](https://huggingface.co/rryisthebest/First_Model).
95
+
96
+
To run the CodeRanker reranking, set RERANK_TYPE="code" and CODE_PROMPT_TYPE="docstring" (Codesearchnet) or "github_issue" (Swebench) in bash/run_rerank_llm.sh and run:
97
+
98
+
```
99
+
bash bash/run_rerank_llm.sh
100
+
```
101
+
102
+
To evaluate the reranking performance, run:
103
+
104
+
```
105
+
bash bash/run_eval.sh rerank code
106
+
```
107
+
76
108
## 3. Model Training
77
-
We also provide the data and scripts to train the LLM reranker by yourself if you wish to do so.
109
+
**Note:** Below is the training code for FIRST. We are still working on releasing the training code for the CodeRanker.
110
+
111
+
We provide the data and scripts to train the LLM reranker by yourself if you wish to do so.
78
112
### 3a. Training Dataset
79
113
Converted training dataset (alphabetic IDs) is on [HF](https://huggingface.co/datasets/rryisthebest/rank_zephyr_training_data_alpha). The standard numeric training dataset can be found [here](https://huggingface.co/datasets/castorini/rank_zephyr_training_data).
80
114
@@ -88,38 +122,38 @@ We support three training objectives:
88
122
89
123
To train the model, run:
90
124
```
91
-
bash bash/beir/run_train.sh
125
+
bash bash/run_train.sh
92
126
```
93
127
94
128
To train a gated model, login to Huggingface and get token access at huggingface.co/settings/tokens.
95
129
```
96
130
huggingface-cli login
97
131
```
98
-
## 4. Relevance Feedback
132
+
## 4. Relevance Feedback (not relevant for codeReranker)
99
133
We also provide scripts here to use the LLM reranker for a downstream task, such as relevance feedback. [Inference-time relevance feedback](https://arxiv.org/pdf/2305.11744) uses the reranker's output to distill the retriever's query embedding to improve recall.
100
134
### 4a. Dataset preparation for relevance feedback
101
135
To prepare dataset(s) for relevance feedback, run:
102
136
```
103
-
bash bash/beir/run_prepare_distill.sh <Path to folder with BEIR encodings>
137
+
bash bash/run_prepare_distill.sh <Path to folder with BEIR encodings>
104
138
```
105
-
### 4b. Distillation (Relevance Feedback Step)
139
+
### 4b. Distillation (Relevance Feedback Step) not relevant for codeReranker
106
140
You can choose to run distillation with either the cross encoder or the LLM reranker or both sequentially.
107
141
To perform the relevance feedback distillation step, run:
108
142
```
109
-
bash bash/beir/run_distill.sh
143
+
bash bash/run_distill.sh
110
144
```
111
145
This step creates new query embeddings after distillation.
112
146
113
-
### 4c. 2nd Retrieval
147
+
### 4c. 2nd Retrieval (not relevant for codeReranker)
114
148
To perform the retrieval step with the new query embedding after distillation, run:
115
149
```
116
-
bash bash/beir/run_2nd_retrieval.sh <Path to folder with BEIR encodings>
150
+
bash bash/run_2nd_retrieval.sh <Path to folder with BEIR encodings>
117
151
```
118
152
119
-
### 4d. Relevance feedback evaluation
153
+
### 4d. Relevance feedback evaluation (not relevant for codeReranker)
0 commit comments