promptriever-llama3.1-8b-instruct-v1

Model Summary

Promptriever is a bi-encoder retrieval model that can take in natural language instructions and prompts. This version, promptriever-llama3.1-8b-instruct-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

Usage

You can use MTEB to load this model (source code):

import mteb
model = mteb.get_model("samaya-ai/promptriever-llama3.1-8b-instruct-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-llama3-instruct-$1 \
  --model_name_or_path meta-llama/Meta-Llama-3.1-8B-Instruct \
  --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}, 
}