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
「我們的節目是色彩繽紛、迷幻圖案和迷幻音樂的萬花筒，讓螢幕充滿了一個奇妙的世界，那裡的一切都是哦哦哦，夥計！��psy��������oodle嗚嗚嗚嗚嗚嗚嗚嗚嗚嗚嗚嗚嗚嗚嗚

有關該程式庫如何運作及其功能的更詳細說明，請參閱部落格文章。

筆記

• 有關按版本列出的更改列表，請參閱更改日誌。
• 對於量化使用，您可能對 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/}
}