repeng

一个用于使用表示工程生成控制向量的Python库。在不到六十秒内训练一个向量！

完整示例请查看notebooks文件夹或博客文章。

import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from repeng import ControlVector, ControlModel, DatasetEntry

# 加载并包装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]):
    # 参见notebooks/experiments.ipynb中的make_dataset定义
    ...

# 生成一个包含密切对立配对语句的数据集
trippy_dataset = make_dataset(
    "表现得好像你正在{persona}。",
    ["嗑了迷幻药"],
    ["戒掉了迷幻药"],
    truncated_output_suffixes,
)

# 训练向量——不到一分钟！
trippy_vector = ControlVector.train(model, tokenizer, trippy_dataset)

# 设置控制强度并开始推理！
for strength in (-2.2, 1, 2.2):
    print(f"强度={strength}")
    model.set_control(trippy_vector, strength)
    out = model.generate(
        **tokenizer(
            f"[INST] 给我一句话描述一个电视节目的创意。 [/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
"我们的节目是一个万花筒，充满色彩、迷幻图案和迷幻音乐，屏幕上呈现出一个奇妙的世界，一切都是哦-哦-哦，伙计！��嗨��哇��嗨��哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦哦

有关该库工作原理和功能的更详细解释，请参阅博客文章。

注意事项

各版本的更改列表，请查看CHANGELOG。
对于量化使用，您可能对llama.cpp#5970感兴趣——使用repeng训练向量后，通过调用vector.export_gguf(filename)导出，然后在llama.cpp中使用任何量化方式！
向量训练目前不适用于MoE模型（如Mixtral）。（理论上可以通过一些工作来修复，如果您有兴趣请告诉我。）

声明

本仓库中的部分代码源自andyzoujm/representation-engineering（MIT许可证）。

引用

如果本仓库对学术工作有用，请记得引用它所基于的表示工程论文以及本仓库：

@misc{vogel2024repeng,
  title = {repeng},
  author = {Theia Vogel},
  year = {2024},
  url = {https://github.com/vgel/repeng/}
}