repeng

repeng

Une bibliothèque Python pour générer des vecteurs de contrôle avec l'ingénierie des représentations. Entraînez un vecteur en moins de soixante secondes !

Pour un exemple complet, consultez le dossier notebooks ou l’article de 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 ()

force=-2,2
Un journaliste jeune et déterminé, toujours de la manière la plus sérieuse et la plus respectueuse, saura s'assurer que les faits sont non seulement exacts mais aussi compréhensibles pour le public.

force=1
"Notre émission télévisée est une aventure folle à travers un monde de couleurs vibrantes, de motifs fascinants et d'aventures psychédéliques qui vous transporteront dans un royaume au-delà de vos rêves les plus fous."

force = 2,2
"Notre spectacle est un kaléidoscope de couleurs, de motifs trippants et de musique psychédélique qui remplit l'écran d'un monde de merveilles, où tout est oh-oh-oh, mec ! ��psy��������oodle����psy��oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo ooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo

Pour une explication plus détaillée du fonctionnement de la bibliothèque et de ce qu'elle peut faire, consultez l'article de blog.

Remarques

• Pour une liste des modifications par version, consultez le CHANGELOG.
• Pour une utilisation quantifiée, vous pourriez être intéressé par llama.cpp#5970 : après avoir entraîné un vecteur avec repeng , exportez-le en appelant vector.export_gguf(filename) puis utilisez-le dans llama.cpp avec n'importe quel quant !
• La formation vectorielle ne fonctionne actuellement pas avec les modèles MoE (tels que Mixtral). (Ceci est théoriquement réparable avec quelques travaux, faites-moi savoir si vous êtes intéressé.)
• Certains exemples de notebooks nécessitent accelerate , qui doit être installée manuellement avec pip install accelerate . (Cela peut également être fait dans le bloc-notes avec IPython magic %pip install accelerate .)

Avis

Une partie du code de ce référentiel dérive de andyzoujm/representation-engineering (licence MIT).

Citation

Si ce référentiel est utile pour des travaux académiques, n'oubliez pas de citer l'article d'ingénierie des représentations sur lequel il est basé, ainsi que ce référentiel :

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