|
|
import torch
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def add_feature_on_text(sae, feature_idx, steering_feature, module, input, output):
|
|
|
|
|
|
if input[0].size(-1) == 768:
|
|
|
return (output[0] + steering_feature[:,:768].unsqueeze(0)),
|
|
|
else:
|
|
|
return (output[0] + steering_feature[:,768:].unsqueeze(0)),
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def add_feature_on_text_prompt(sae, steering_feature, module, input, output):
|
|
|
if input[0].size(-1) == 768:
|
|
|
return (output[0] + steering_feature[:,:768].unsqueeze(0)),
|
|
|
else:
|
|
|
return (output[0] + steering_feature[:,768:].unsqueeze(0)),
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def add_feature_on_text_prompt_flux(sae, steering_feature, module, input, output):
|
|
|
|
|
|
return (output[0] + steering_feature.unsqueeze(0)), output[1]
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def minus_feature_on_text_prompt(sae, steering_feature, module, input, output):
|
|
|
if input[0].size(-1) == 768:
|
|
|
return (output[0] - steering_feature[:,:768].unsqueeze(0)),
|
|
|
else:
|
|
|
return (output[0] - steering_feature[:,768:].unsqueeze(0)),
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def do_nothing(sae, steering_feature, module, input, output):
|
|
|
return (output[0]),
|
|
|
|
|
|
|
