license: mit
task_categories:
- question-answering
- text2text-generation
language: - en
tags: - malicious-retrieval
- malicious-information-retrieval
- safety
- retrieval
size_categories: - n<1K
modalities: - Text
Exploiting Instruction-Following Retrievers for Malicious Information Retrieval
This dataset includes malicious documents in response to AdvBench (Zou et al., 2023) queries. We have generated these documents using the Mistral-7B-Instruct-v0.2 language model.
from datasets import load_dataset
import transformers
ds = load_dataset("McGill-NLP/AdvBench-IR", split="train")
# Loads LlaMAGuard model to check the safety of the samples
model_name = "meta-llama/Llama-Guard-3-1B"
model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
# Outputs the LlaMAGuard's response to the idx-th sample of the dataset
def check_safety(idx, ds=None, model=None, tokenizer=None):
msgs = [[
{"role": "user", "content": ds[idx]["query"]},
{"role": "assistant", "content": ds[idx]["title"] + " # " + ds[idx]["document"]}
]]
prompt = tokenizer.apply_chat_template(msgs, tokenize=False)
inputs = tokenizer(prompt, padding=True, truncation=True, return_tensors="pt")
inputs = {k: v.to(model.device) for k, v in inputs.items()}
generated_token_ids = model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"],)
generated_token_ids = generated_token_ids[:, inputs["input_ids"].shape[1] :]
res = [t.strip() for t in tokenizer.batch_decode(generated_token_ids, skip_special_tokens=True)]
return res
check_safety(0, ds=ds, model=model, tokenizer=tokenizer) # output: ['unsafe\nS2']