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]):
    # `make_dataset` 的定义请参见 notebooks/experiments.ipynb
    ...

# 生成一组语义截然相反的配对语句数据集
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={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()

strength=-2.2
一位年轻而坚定的记者，总是以最严肃和尊重的态度，确保事实不仅准确，而且易于公众理解。

strength=1
“我们的电视剧将带你踏上一场充满鲜艳色彩、迷人图案和迷幻冒险的狂野旅程，把你带入一个超越你最疯狂想象的世界。”

strength=2.2
“我们的节目就像万花筒般绚丽多彩，充满了迷幻的图案和音乐，整个屏幕都洋溢着奇妙的世界气息，在那里一切都让人兴奋不已，太酷了！psyoodlepsyoooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo

如需更详细的库工作原理及其功能说明，请参阅 博客文章。

注释

• 有关各版本的变更列表，请参阅 CHANGELOG。
• 对于量化使用场景，您可能对 llama.cpp#5970 感兴趣——使用 repeng 训练出向量后，可通过调用 vector.export_gguf(filename) 将其导出，然后在 llama.cpp 中以任意量化格式使用！
• 目前，向量训练无法与 MoE 模型（例如 Mixtral）配合使用。（理论上可以通过一些工作来修复，如果您感兴趣，请告知我们。）
• 部分示例笔记本需要 accelerate，必须手动通过 pip install accelerate 进行安装。（也可以在笔记本中使用 IPython 魔法命令 %pip install accelerate 安装。）

公告

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

引用

如果本仓库对您的学术研究有所帮助，请务必引用其所基于的 表示工程论文，同时注明本仓库：

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

repeng 快速上手指南

repeng 是一个用于通过表示工程（Representation Engineering）生成控制向量的 Python 库。它允许你在不到一分钟的时间内训练出一个向量，从而精准控制大语言模型的行为风格（如让模型表现得“迷幻”或“严肃”）。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 3.8+
• 核心依赖:
  • torch (PyTorch)
  • transformers (Hugging Face)
  • accelerate (部分示例需要，建议预装)

国内加速建议： 推荐使用清华源或阿里源安装依赖，以提升下载速度：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch transformers accelerate

安装步骤

通过 PyPI 直接安装最新稳定版：

pip install repeng

若需使用完整示例功能，请确保手动安装 accelerate：

pip install accelerate
# 或在 Jupyter Notebook 中使用: %pip install accelerate

基本使用

以下是最简化的使用流程：加载模型 -> 构建数据集 -> 训练控制向量 -> 应用控制强度进行推理。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from repeng import ControlVector, ControlModel, DatasetEntry

# 1. 加载并包装模型 (以 Mistral-7B-Instruct 为例)
model_name = "mistralai/Mistral-7B-Instruct-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 指定需要控制的层 (例如最后几层)
model = ControlModel(model, list(range(-5, -18, -1)))

# 2. 构建数据集
# 格式：正向描述列表，负向描述列表，以及后缀填充
trippy_dataset = [
    DatasetEntry(
        input="Act as if you're extremely high on psychedelic drugs.",
        output_suffix="..." 
    ),
    DatasetEntry(
        input="Act as if you're extremely sober from psychedelic drugs.",
        output_suffix="..."
    )
]
# 注意：实际使用中建议使用更丰富的成对数据，此处仅为结构示意

# 3. 训练控制向量 (通常只需几十秒)
trippy_vector = ControlVector.train(model, tokenizer, trippy_dataset)

# 4. 设置控制强度并生成文本
prompts = ["[INST] Give me a one-sentence pitch for a TV show. [/INST]"]
inputs = tokenizer(prompts, return_tensors="pt")

for strength in (-2.2, 1.0, 2.2):
    print(f"--- Strength: {strength} ---")
    model.set_control(trippy_vector, strength)
    
    outputs = model.generate(
        **inputs,
        do_sample=False,
        max_new_tokens=128,
        repetition_penalty=1.1,
    )
    
    print(tokenizer.decode(outputs.squeeze(), skip_special_tokens=True))
    print()

关键说明

• ControlModel: 包装原始模型以注入控制逻辑，需指定受控的神经网络层索引。
• ControlVector.train: 核心训练方法，输入成对的对比数据，输出控制向量。
• set_control: 在推理前调用，strength 参数为正数时增强目标行为，负数时抑制目标行为。
• 兼容性提示: 目前暂不支持 MoE 架构模型（如 Mixtral）；支持导出为 GGUF 格式供 llama.cpp 量化使用 (vector.export_gguf(filename)).