--- base_model: meta-llama/Meta-Llama-3.1-8B 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-llama3.1-8b-v1` was instruction-trained on a corpus of 490k MSMarco samples with instructions and 490k without instructions. See the [paper](https://arxiv.org/abs/2409.11136) for more details. - **Repository:** [orionw/Promptriever](https://github.com/orionw/promptriever) - **Paper:** [Promptriever: Instruction-Trained Retrievers Can Be Prompted Like Language Models](https://arxiv.org/abs/2409.11136) - **Instruction-Training Dataset:** [samaya-ai/msmarco-w-instructions](https://huggingface.co/datasets/samaya-ai/msmarco-w-instructions) # Other Links | Binary | Description | |:-------|:------------| | [samaya-ai/promptriever-llama2-7b-v1](https://huggingface.co/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](https://huggingface.co/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](https://huggingface.co/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](https://huggingface.co/samaya-ai/promptriever-mistral-v0.1-7b-v1) | A Promptriever bi-encoder model based on Mistral v0.1 (7B parameters). | | [samaya-ai/RepLLaMA-reproduced](https://huggingface.co/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](https://huggingface.co/datasets/Tevatron/msmarco-passage-aug) dataset. | | [samaya-ai/msmarco-w-instructions](https://huggingface.co/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](https://github.com/embeddings-benchmark/mteb/blob/main/mteb/models/promptriever_models.py)): ```python import mteb model = mteb.get_model("samaya-ai/promptriever-llama3.1-8b-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: ```python 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](https://github.com/orionw/tevatron) to fine-tune promptriever with the [samaya-ai/msmarco-w-instructions](https://huggingface.co/datasets/samaya-ai/msmarco-w-instructions) dataset. You can reproduce this with [this script](https://github.com/orionw/promptriever/blob/main/scripts/training/train_instruct_llama3.sh) (reproduced here for convenience). ```bash #!/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-$1 \ --model_name_or_path meta-llama/Meta-Llama-3.1-8B \ --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 ```bibtex @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}, } ```