repeng

repeng

Una biblioteca de Python para generar vectores de control con ingeniería de representación. ¡Entrena un vector en menos de sesenta segundos!

Para ver un ejemplo completo, consulte la carpeta de cuadernos o la publicación del 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 ()

fuerza = -2.2
Un periodista joven y decidido, siempre de la manera más seria y respetuosa, podrá asegurarse de que los hechos no sólo sean exactos sino también comprensibles para el público.

fuerza = 1
"Nuestro programa de televisión es un viaje salvaje a través de un mundo de colores vibrantes, patrones fascinantes y aventuras psicodélicas que te transportarán a un reino más allá de tus sueños más locos".

fuerza = 2.2
"Nuestro programa es un caleidoscopio de colores, patrones alucinantes y música psicodélica que llena la pantalla con un mundo de maravillas, donde todo es ¡oh-oh-oh, hombre! ��psy����������oodle����psy��oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo ooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo

Para obtener una explicación más detallada de cómo funciona la biblioteca y qué puede hacer, consulte la publicación del blog.

Notas

• Para obtener una lista de cambios por versión, consulte CHANGELOG.
• Para uso cuantificado, puede que le interese llama.cpp#5970. Después de entrenar un vector con repeng , expórtelo llamando vector.export_gguf(filename) y luego úselo en llama.cpp con cualquier cantidad.
• El entrenamiento vectorial actualmente no funciona con modelos MoE (como Mixtral). (Esto, en teoría, se puede solucionar con algo de trabajo; avíseme si está interesado).
• Algunos portátiles de ejemplo requieren accelerate , que debe instalarse manualmente con pip install accelerate . (Esto también se puede hacer en la computadora portátil con la %pip install accelerate de IPython).

Aviso

Parte del código de este repositorio deriva de andyzoujm/representation-engineering (licencia MIT).

Citación

Si este repositorio es útil para el trabajo académico, recuerde citar el artículo de ingeniería de representación en el que se basa, junto con este repositorio:

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