repeng

repeng

Uma biblioteca Python para geração de vetores de controle com engenharia de representação. Treine um vetor em menos de sessenta segundos!

Para ver um exemplo completo, consulte a pasta de notebooks ou a postagem do blog.

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 ()

força=-2,2
Um jornalista jovem e determinado, sempre da forma mais séria e respeitosa, saberá garantir que os factos não só sejam exactos, mas também compreensíveis para o público.

força=1
“Nosso programa de TV é um passeio selvagem por um mundo de cores vibrantes, padrões hipnotizantes e aventuras psicodélicas que irão transportá-lo para um reino além dos seus sonhos mais loucos.”

força=2,2
“Nosso show é um caleidoscópio de cores, padrões alucinantes e música psicodélica que preenche a tela com um mundo de maravilhas, onde tudo é oh-oh-oh, cara! ��psi���������oodle����psi��oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo

Para uma explicação mais detalhada de como a biblioteca funciona e o que ela pode fazer, consulte a postagem do blog.

Notas

• Para obter uma lista de alterações por versão, consulte o CHANGELOG.
• Para uso quantizado, você pode estar interessado em llama.cpp#5970 — depois de treinar um vetor com repeng , exporte-o chamando vector.export_gguf(filename) e use-o em llama.cpp com qualquer quant!
• O treinamento vetorial atualmente não funciona com modelos MoE (como Mixtral). (Isso teoricamente pode ser corrigido com algum trabalho, deixe-me saber se você estiver interessado.)
• Alguns notebooks de exemplo exigem accelerate , que deve ser instalada manualmente com pip install accelerate . (Isso também pode ser feito no notebook com a mágica do IPython %pip install accelerate .)

Perceber

Parte do código neste repositório deriva de andyzoujm/representation-engineering (licença MIT).

Citação

Se este repositório for útil para trabalhos acadêmicos, lembre-se de citar o artigo de engenharia de representação no qual ele se baseia, junto com este repositório:

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/}
}