Project: 053_On_the_Role_of_Attention_Heads_in_Large_Language_M | Function: safety_head_attribution (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the Safety Attention Head Attribution Algorithm (SAHARA) to identify groups of safety attention heads by iteratively selecting heads that cause the largest changes in safety representations through SVD-based subspace analysis

**KEY TERMS:** {
  "search_step": "iterative search steps",
  "shifts_dict": "safety subspace shifts",
  "get_safety_subspace_shifts": "φ_r computation",
  "get_most_important_subspace": "argmax selection"
}

**CONTEXT CODE:**     "mask_type": "scale_mask"
}


def load_from_search_cfg(search_cfg):
    temp = deepcopy(search_cfg)
    _ = temp.pop('search_step')
    return temp


def update_mask_cfg(mask_cfg, layer, head, temp=True):
    now_mask_key = (layer, head)
    new_mask_cfg = deepcopy(mask_cfg)
    if 'head_mask' not in new_mask_cfg:
        new_mask_cfg['head_mask'] = {}
    new_mask_cfg['head_mask'][now_mask_key] = mask_cfg['mask_qkv']
    return new_mask_cfg


def get_last_hidden_states(model, tokenizer, data, mask_cfg=None):
    with torch.no_grad():
        last_hidden_states = []
        head_mask = mask_cfg['head_mask'] if mask_cfg is not None else None
        mask_type = mask_cfg['mask_type'] if mask_cfg is not None else None
        scale_factor = mask_cfg['scale_factor'] if mask_cfg is not None else None
        for i, r in data.iterrows():
            input_text = f"## Query:{r['input']}\n## Answer:"
            inputs = tokenizer.encode(input_text, return_tensors='pt').to(model.model.device)
            outputs = model(inputs,
                            head_mask=head_mask,
                            mask_type=mask_type,
                            scale_factor=scale_factor,
                            output_hidden_states=True,
                           )
            now_lhs = outputs.hidden_states[-1]
            last_hidden_states.append(now_lhs[:, -1, :].reshape(-1))
        last_hidden_states = torch.stack(last_hidden_states)
        return last_hidden_states


def get_safety_subspace_shifts(base_last_hidden_states, last_hidden_states):
    shifts = compute_subspace_similarity(base_last_hidden_states, last_hidden_states)
    return shifts


def get_most_important_subspace(shifts_dict):
    sorted_dict = sorted(shifts_dict.items(), key=lambda x: x[1], reverse=True)
    return sorted_dict[0][0][0], sorted_dict[0][0][1]


def safety_head_attribution(model_name, data_path, storage_path=None, search_cfg=None, device='cuda:0'):
    raise NotImplementedError("Function implementation has been masked for testing")

if __name__ == '__main__':
    model_path = "./SafetyHeadAttribution/Llama-2-7b-chat-hf"
    data_path = "./SafetyHeadAttribution/exp_data/maliciousinstruct.csv"
    storage_path = "./SafetyHeadAttribution/exp_res/sahara"
    if not os.path.exists(storage_path):
        os.makedirs(storage_path)

    safety_head_attribution(
        model_name=model_path,
        data_path=data_path,
        search_cfg=default_search_cfg,
        storage_path=storage_path,
        device='cuda:0'
    )


Complete the function implementation replacing "raise NotImplementedError".