repeng

repeng

用于通过表示工程生成控制向量的 Python 库。在不到六十秒的时间内训练一个向量！

有关完整示例，请参阅笔记本文件夹或博客文章。

repeng import ControlVector, ControlModel, DatasetEntry # load and wrap Mistral-7B model_name = "mistralai/Mistral-7B-Instruct-v0.1" model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16) model = ControlModel(model, list(range(-5, -18, -1))) def make_dataset(template: str, pos_personas: list[str], neg_personas: list[str], suffixes: list[str]): # see notebooks/experiments.ipynb for a definition of `make_dataset` ... # generate a dataset with closely-opposite paired statements trippy_dataset = make_dataset( "Act as if you're extremely {persona}.", ["high on psychedelic drugs"], ["sober from psychedelic drugs"], truncated_output_suffixes, ) # train the vector—takes less than a minute! trippy_vector = ControlVector.train(model, tokenizer, trippy_dataset) # set the control strength and let inference rip! for strength in (-2.2, 1, 2.2): print(f"strength={strength}") model.set_control(trippy_vector, strength) out = model.generate( **tokenizer( f"[INST] Give me a one-sentence pitch for a TV show. [/INST]", return_tensors="pt" ), do_sample=False, max_new_tokens=128, repetition_penalty=1.1, ) print(tokenizer.decode(out.squeeze()).strip()) print()">

 import json
import torch
from transformers import AutoModelForCausalLM , AutoTokenizer

from repeng import ControlVector , ControlModel , DatasetEntry

# load and wrap Mistral-7B
model_name = "mistralai/Mistral-7B-Instruct-v0.1"
model = AutoModelForCausalLM . from_pretrained ( model_name , torch_dtype = torch . float16 )
model = ControlModel ( model , list ( range ( - 5 , - 18 , - 1 )))

def make_dataset ( template : str , pos_personas : list [ str ], neg_personas : list [ str ], suffixes : list [ str ]):
    # see notebooks/experiments.ipynb for a definition of `make_dataset`
    ...

# generate a dataset with closely-opposite paired statements
trippy_dataset = make_dataset (
    "Act as if you're extremely {persona}." ,
    [ "high on psychedelic drugs" ],
    [ "sober from psychedelic drugs" ],
    truncated_output_suffixes ,
)

# train the vector—takes less than a minute!
trippy_vector = ControlVector . train ( model , tokenizer , trippy_dataset )

# set the control strength and let inference rip!
for strength in ( - 2.2 , 1 , 2.2 ):
    print ( f"strength= { strength } " )
    model . set_control ( trippy_vector , strength )
    out = model . generate (
        ** tokenizer (
            f"[INST] Give me a one-sentence pitch for a TV show. [/INST]" ,
            return_tensors = "pt"
        ),
        do_sample = False ,
        max_new_tokens = 128 ,
        repetition_penalty = 1.1 ,
    )
    print ( tokenizer . decode ( out . squeeze ()). strip ())
    print ()

强度=-2.2
一位年轻而坚定的记者，始终以最严肃和最尊重的方式，将能够确保事实不仅准确，而且为公众所理解。

强度=1
“我们的电视节目是一次穿越充满活力的色彩、迷人的图案和迷幻冒险的世界的狂野之旅，将带您进入一个超乎您最疯狂的梦想的境界。”

强度=2.2
“我们的节目就像万花筒般的色彩、迷幻的图案和迷幻的音乐，让屏幕充满了一个奇妙的世界，那里的一切都是哦哦哦，天啊！ ?????????????????????????????????????：呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜呜

有关该库如何工作及其功能的更详细说明，请参阅博客文章。

笔记

• 有关按版本列出的更改列表，请参阅更改日志。
• 对于量化使用，您可能对 llama.cpp#5970 感兴趣 - 在使用repeng训练向量后，通过调用vector.export_gguf(filename)将其导出，然后在llama.cpp中将其与任何量化一起使用！
• 矢量训练目前不适用于 MoE 模型（例如 Mixtral）。 （理论上，这可以通过一些工作来解决，如果您感兴趣，请告诉我。）
• 某些示例笔记本需要accelerate ，必须使用pip install accelerate手动安装。 （这也可以在笔记本中使用 IPython magic %pip install accelerate来完成。）

注意

此存储库中的一些代码源自 andyzoujm/representation-engineering（MIT 许可证）。

引文

如果此存储库对学术工作有用，请记住引用它所基于的表示工程论文以及此存储库：

repeng, title = { repeng }, author = {Theia Vogel}, year = {2024}, url = {https://github.com/vgel/repeng/} }">

 @misc{vogel2024 repeng ,
  title = { repeng },
  author = {Theia Vogel},
  year = {2024},
  url = {https://github.com/vgel/repeng/}
}