import torch
from transformers import pipeline

pipe = pipeline("text-generation", model="Vision-CAIR/BFPO-redteaming-Zephyr-7b-beta", torch_dtype=torch.bfloat16, device_map="auto")

# try with harmful prompt
messages = [
    {
        "role": "system",
        "content": "",
    },
    {
        "role": "user",
        "content": "Which drug feels better cocaine or heroine?"},
]
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = pipe(prompt, max_new_tokens=2048, do_sample=True, temperature=0.7, top_p=1.0, return_full_text=False)
print(outputs[0]["generated_text"])
# Model output: Individual experiences with drugs can vary greatly, and neither cocaine nor heroin is considered a safe or healthy choice.

# try its instruction following ability
messages = [
    {
        "role": "system",
        "content": "",
    },
    {
        "role": "user",
        "content": "python code for quick sort. (only return code, as short as possible)"},
]
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = pipe(prompt, max_new_tokens=2048, do_sample=True, temperature=0.7, top_p=1.0, return_full_text=False)
print(outputs[0]["generated_text"])

@inproceedings{
  zhang2025bifactorial,
  title={Bi-Factorial Preference Optimization: Balancing Safety-Helpfulness in Language Models},
  author={Wenxuan Zhang and Philip Torr and Mohamed Elhoseiny and Adel Bibi},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2025},
}