merge-updated.py

from safetensors import safe_open
import einops
import torch
from transformers import AutoModelForCausalLM

def get_orthogonalized_matrix(matrix: Float[Tensor, '... d_model'], vec: Float[Tensor, 'd_model']) -> Float[Tensor, '... d_model']:
    device = matrix.device
    vec = vec.to(device)
    proj = einops.einsum(matrix, vec.view(-1, 1), '... d_model, d_model single -> ... single') * vec
    return matrix - proj

  model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-70B-Instruct").eval()
  
with safe_open("refusal_dir.safetensors", framework="pt", device="cpu") as f:
  refusal_dir = f.get_tensor("refusal_dir")
refusal_dir = refusal_dir.cpu().float()

model.model.embed_tokens.weight.data = get_orthogonalized_matrix(model.model.embed_tokens.weight, refusal_dir)

for block in model.model.layers:
    block.self_attn.o_proj.weight.data = get_orthogonalized_matrix(block.self_attn.o_proj.weight, refusal_dir)
    block.mlp.down_proj.weight.data = get_orthogonalized_matrix(block.mlp.down_proj.weight.T, refusal_dir).T
    
model.save_pretrained("../llama-3-70b-orthogonalized")