Найти тему
10,1 тыс подписчиков

🖥 pyvene: A Library for Understanding and Improving PyTorch Models via Interventions


Стэнфордская библиотека NLP для понимания и улучшения моделей на основе PyTorch.

Воздействие на внутренние состояния модели являются важной операцией во многих областях работы с ИИ, включая редактирование модели, управление, ее надежность и интерпретируемость.

Для облегчения таких задач исследователи Стэнфорда создали библиотеку Python с открытым исходным кодом, которая поддерживает сложные схемы взаимодействия с моделями в интуитивно понятном формате.

pip install pyvene

Например, вы можете использовать любую модель hf:

import torch
import pyvene as pv
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "meta-llama/Llama-2-7b-hf" # your HF model name.
model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype=torch.bfloat16, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name)

def zeroout_intervention_fn(b, s):
b[:,3] = 0. # 3rd position
return b

pv_model = pv.IntervenableModel({
"component": "model.layers[15].mlp.output", # string access
"intervention": zeroout_intervention_fn}, model=model)

# run the intervened forward pass
orig_outputs, intervened_outputs = pv_model(
tokenizer("The capital of Spain is", return_tensors="pt").to('cuda'),
output_original_output=True
)
print(intervened_outputs.logits - orig_outputs.logits)

Вернет:

tensor([[[ 0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.
[ 0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[ 0.4375, 1.0625, 0.3750, ..., -0.1562, 0.4844, 0.2969],
[ 0.0938, 0.1250, 0.1875, ..., 0.2031, 0.0625, 0.2188],
[ 0.0000, -0.0625, -0.0312, ..., 0.0000, 0.0000, -0.0156]]],
device='cuda:0')


1 минута