# Task-aware Retrieval with Instructions

Official repository: [github.com/facebookresearch/tart](https://github.com/facebookresearch/tart)

### Model descriptions

`facebook/tart-full-t0-3b` is a multi-task cross-encoder model trained via instruction-tuning on approximately 40 retrieval tasks, which is initialized with [bigscience/T0_3B](https://huggingface.co/bigscience/T0_3B).

TART-full is a 1.5 billion cross-necoder and it can rerank top documents given a query and natural language instruction (e.g., *find a Wikipedia paragraph that answers this question.*). 
Experimental results on widely-used [BEIR](https://github.com/beir-cellar/beir), [LOTTE](https://huggingface.co/datasets/colbertv2/lotte), and our new evaluation, [X^2-Retrieval](https://github.com/facebookresearch/tart/cross_task_cross_eval) show that TART-full outperforms previous state-of-the-art methods by levaraging natural language instructions. 

More details about modeling and training are in our paper: [Task-aware Retrieval with Instructions](https://arxiv.org/abs/2211.09260). 

### Installation
```sh
git clone https://github.com/facebookresearch/tart
pip install -r requirements.txt
cd tart/TART
```

### How to use?

TART-full can be loaded through our customized EncT5 model. 
```python
from src.modeling_enc_t5 import EncT5ForSequenceClassification
from src.tokenization_enc_t5 import EncT5Tokenizer
import torch
import torch.nn.functional as F
import numpy as np

# load TART full and tokenizer
model = EncT5ForSequenceClassification.from_pretrained("facebook/tart-full-t0-3b")
tokenizer =  EncT5Tokenizer.from_pretrained("facebook/tart-full-t0-3b")
model.eval()

q = "What is the population of Tokyo?"
in_answer = "retrieve a passage that answers this question from Wikipedia"

p_1 = "The population of Japan's capital, Tokyo, dropped by about 48,600 people to just under 14 million at the start of 2022, the first decline since 1996, the metropolitan government reported Monday."
p_2 = "Tokyo, officially the Tokyo Metropolis (東京都, Tōkyō-to), is the capital and largest city of Japan."

# 1. TART-full can identify more relevant paragraph. 
features = tokenizer(['{0} [SEP] {1}'.format(in_answer, q), '{0} [SEP] {1}'.format(in_answer, q)], [p_1, p_2], padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
    scores = model(**features).logits
    normalized_scores = [float(score[1]) for score in F.softmax(scores, dim=1)]

print([p_1, p_2][np.argmax(normalized_scores)]) # "The population of Japan's capital, Tokyo, dropped by about 48,600 people to just under 14 million ... "

# 2. TART-full can identify the document that is more relevant AND follows instructions.
q_1 = "How many people live in Tokyo?"
features = tokenizer(['{0} [SEP] {1}'.format(in_answer, q), '{0} [SEP] {1}'.format(in_answer, q)], [p_1, q_1], padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
    scores = model(**features).logits
    normalized_scores = [float(score[1]) for score in F.softmax(scores, dim=1)]

print([p_1, q_1][np.argmax(normalized_scores)]) #  "The population of Japan's capital, Tokyo, dropped by about 48,600 people to just under 14 million"
```