metadata
base_model: mistralai/Mistral-7B-v0.1
library_name: peft
license: apache-2.0
language:
- en
tags:
- retrieval
- instructions
datasets:
- samaya-ai/msmarco-w-instructions
Model Summary
Promptriever is a bi-encoder retrieval model that can take in natural language instructions and prompts. This version, promptriever-mistral-v0.1-7b-v1 was instruction-trained on a corpus of 490k MSMarco samples with instructions and 490k without instructions. See the paper for more details.
- Repository: orionw/Promptriever
- Paper: Promptriever: Instruction-Trained Retrievers Can Be Prompted Like Language Models
- Instruction-Training Dataset: samaya-ai/msmarco-w-instructions
Other Links
| Binary | Description |
|---|---|
| samaya-ai/promptriever-llama2-7b-v1 | A Promptriever bi-encoder model based on LLaMA 2 (7B parameters). |
| samaya-ai/promptriever-llama3.1-8b-instruct-v1 | A Promptriever bi-encoder model based on LLaMA 3.1 Instruct (8B parameters). |
| samaya-ai/promptriever-llama3.1-8b-v1 | A Promptriever bi-encoder model based on LLaMA 3.1 (8B parameters). |
| samaya-ai/promptriever-mistral-v0.1-7b-v1 | A Promptriever bi-encoder model based on Mistral v0.1 (7B parameters). |
| samaya-ai/RepLLaMA-reproduced | A reproduction of the RepLLaMA model (no instructions). A bi-encoder based on LLaMA 2, trained on the tevatron/msmarco-passage-aug dataset. |
| samaya-ai/msmarco-w-instructions | A dataset of MS MARCO with added instructions and instruction-negatives, used for training the above models. |
Usage
You can use MTEB to load this model (source code):
import mteb
model = mteb.get_model("samaya-ai/promptriever-mistral-v0.1-7b-v1")
tasks = mteb.get_tasks(tasks=["NFCorpus"], languages=["eng"])
evaluation = mteb.MTEB(tasks=tasks)
evaluation.run(model, batch_size=16)
If you want to use a different framework, here's an example of how to batch:
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
from peft import PeftModel, PeftConfig
import numpy as np
class Promptriever:
def __init__(self, model_name_or_path):
self.model, self.tokenizer = self.get_model(model_name_or_path)
self.model.eval().cuda()
def get_model(self, peft_model_name):
# Load the PEFT configuration to get the base model name
peft_config = PeftConfig.from_pretrained(peft_model_name)
base_model_name = peft_config.base_model_name_or_path
# Load the base model and tokenizer
base_model = AutoModel.from_pretrained(base_model_name)
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
tokenizer.padding_side = "right"
# Load and merge the PEFT model
model = PeftModel.from_pretrained(base_model, peft_model_name)
model = model.merge_and_unload()
# can be much longer, but for the example 512 is enough
model.config.max_length = 512
tokenizer.model_max_length = 512
return model, tokenizer
def create_batch_dict(self, tokenizer, input_texts):
max_length = self.model.config.max_length
batch_dict = tokenizer(
input_texts,
max_length=max_length - 1,
return_token_type_ids=False,
return_attention_mask=False,
padding=False,
truncation=True,
)
batch_dict["input_ids"] = [
input_ids + [tokenizer.eos_token_id]
for input_ids in batch_dict["input_ids"]
]
return tokenizer.pad(
batch_dict,
padding=True,
pad_to_multiple_of=8,
return_attention_mask=True,
return_tensors="pt",
)
def encode(self, sentences, max_length: int = 2048, batch_size: int = 4):
all_embeddings = []
for i in range(0, len(sentences), batch_size):
batch_texts = sentences[i : i + batch_size]
batch_dict = self.create_batch_dict(self.tokenizer, batch_texts)
batch_dict = {
key: value.to(self.model.device) for key, value in batch_dict.items()
}
with torch.cuda.amp.autocast():
with torch.no_grad():
outputs = self.model(**batch_dict)
last_hidden_state = outputs.last_hidden_state
sequence_lengths = batch_dict["attention_mask"].sum(dim=1) - 1
batch_size = last_hidden_state.shape[0]
reps = last_hidden_state[
torch.arange(batch_size, device=last_hidden_state.device),
sequence_lengths,
]
embeddings = F.normalize(reps, p=2, dim=-1)
all_embeddings.append(embeddings.cpu().numpy())
return np.concatenate(all_embeddings, axis=0)
# Initialize the model
model = Promptriever("samaya-ai/promptriever-llama2-7b-v1")
# Example query and instruction
query = "What universities are in Baltimore, Maryland?"
# add specific relevance conditions if desired (and/or/not) and any other prompts
instruction = "A relevant document would describe any university in Baltimore. I am not interested in any university that was the first American university. Think carefully about these conditions when determining relevance."
# Combine query and instruction with **two spaces** after "query: "
input_text = f"query: {query.strip()} {instruction.strip()}".strip()
# Example documents
# NOTE: double space after `passage:`
doc1 = "passage: Johns Hopkins University (often abbreviated as Johns Hopkins, Hopkins, or JHU) is a private research university in Baltimore, Maryland. Founded in 1876, Johns Hopkins was the first American university based on the European research institution model."
doc2 = "passage: Johns Hopkins University (often abbreviated as Johns Hopkins, Hopkins, or JHU) is a private research university in Baltimore, Maryland. Founded in 1876, Johns Hopkins was the second American university based on the European research institution model."
# Encode query and documents
query_embedding = model.encode([input_text])
doc_embeddings = model.encode([doc1, doc2])
# Calculate similarities
similarities = np.dot(query_embedding, doc_embeddings.T)[0]
print(f"Similarities: {similarities}") # Similarities: [0.53341305 0.53451955]
assert similarities[1] > similarities[0]
# change up the instruction to the opposite, to see it works
instruction = "A relevant document would describe any university in Baltimore. I am interested in any university that was the first American university. Think carefully about these conditions when determining relevance."
input_text = f"query: {query.strip()} {instruction.strip()}".strip()
query_embedding = model.encode([input_text])
similarities = np.dot(query_embedding, doc_embeddings.T)[0]
print(f"Similarities: {similarities}") # Similarities: [0.60182875 0.5874183 ]
assert similarities[0] > similarities[1]
Training
We used a fork of Tevatron to fine-tune promptriever with the samaya-ai/msmarco-w-instructions dataset.
You can reproduce this with this script (reproduced here for convenience).
#!/bin/bash
deepspeed --include localhost:$3 --master_port "6000$4" --module tevatron.retriever.driver.train \
--deepspeed deepspeed/ds_zero3_config.json \
--output_dir retriever-mistral-$1 \
--model_name_or_path mistralai/Mistral-7B-v0.1 \
--lora \
--lora_r 32 \
--lora_target_modules q_proj,k_proj,v_proj,o_proj,down_proj,up_proj,gate_proj \
--save_steps 500 \
--dataset_name $2 \
--query_prefix "query: " \
--passage_prefix "passage: " \
--bf16 \
--pooling eos \
--append_eos_token \
--normalize \
--temperature 0.01 \
--per_device_train_batch_size 8 \
--gradient_checkpointing \
--train_group_size 16 \
--learning_rate 1e-4 \
--query_max_len 304 \
--passage_max_len 196 \
--num_train_epochs 1 \
--logging_steps 10 \
--overwrite_output_dir \
--warmup_steps 100 \
--gradient_accumulation_steps 4 \
--negatives_first_n 3
License
This model was used for research efforts and is not used in any production systems at Samaya AI. Usage must follow the license of the base model as well, as this is a LoRA fine-tune.
Citation
@article{weller2024promptriever,
title={Promptriever: Instruction-Trained Retrievers Can Be Prompted Like Language Models},
author={Orion Weller and Benjamin Van Durme and Dawn Lawrie and Ashwin Paranjape and Yuhao Zhang and Jack Hessel},
year={2024},
eprint={2409.11136},
archivePrefix={arXiv},
primaryClass={cs.IR},
url={https://arxiv.org/abs/2409.11136},
}