表現エンジニアリングを使用して制御ベクトルを生成するための Python ライブラリ。 60 秒以内にベクトルをトレーニングします。
完全な例については、ノートブック フォルダーまたはブログ投稿を参照してください。
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���������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������������� おおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおおお
ライブラリの仕組みと何ができるかの詳細については、ブログ投稿を参照してください。
repeng
でベクトルをトレーニングした後、 vector.export_gguf(filename)
を呼び出してエクスポートし、それをllama.cpp
で任意の quant! で使用します。accelerate
必要です。これは、 pip install accelerate
使用して手動でインストールする必要があります。 (これは、IPython マジック%pip install accelerate
使用してノートブックで実行することもできます。) このリポジトリ内のコードの一部は、andyzoujm/representation-engineering (MIT ライセンス) から派生しています。
このリポジトリが学術研究に役立つ場合は、このリポジトリとともに、そのベースとなっている表現工学の論文を忘れずに引用してください。
@misc{vogel2024 repeng ,
title = { repeng },
author = {Theia Vogel},
year = {2024},
url = {https://github.com/vgel/repeng/}
}