| import argparse |
| import json |
| import numpy as np |
| import torch |
| from common import inside, save_json, permit_compute |
| from weights import load, project_writers, save_edit |
|
|
|
|
| def matched_direction(productive, repetitive): |
| a, b = np.asarray(productive), np.asarray(repetitive) |
| if a.shape != b.shape or a.ndim != 2: |
| raise ValueError('Нужны попарные средние активации одинаковой формы [contexts, width]') |
| direction = (b - a).mean(0) |
| norm = np.linalg.norm(direction) |
| if not np.isfinite(norm) or norm == 0: |
| raise ValueError('Направление не определено') |
| return direction / norm |
|
|
|
|
| def writer_scores(weight, direction, productive_mlp, repetitive_mlp): |
| coefficients = np.asarray(direction) @ weight.float().numpy() |
| deltas = np.asarray(repetitive_mlp) - np.asarray(productive_mlp) |
| values = deltas * coefficients |
| return values.mean(0), (values > 0).mean(0) |
|
|
|
|
| if __name__ == '__main__': |
| p = argparse.ArgumentParser() |
| p.add_argument('--model', required=True) |
| p.add_argument('--tensor', required=True) |
| p.add_argument('--productive', required=True) |
| p.add_argument('--repetitive', required=True) |
| p.add_argument('--productive-mlp', required=True) |
| p.add_argument('--repetitive-mlp', required=True) |
| p.add_argument('--columns', type=int, required=True) |
| p.add_argument('--alpha', type=float, required=True) |
| p.add_argument('--out', required=True) |
| p.add_argument('--hypothesis', required=True) |
| a = p.parse_args() |
| permit_compute() |
| direction = matched_direction(np.load(inside(a.productive)), np.load(inside(a.repetitive))) |
| weight = load(a.model, a.tensor) |
| score, consistency = writer_scores(weight, direction, np.load(inside(a.productive_mlp)), np.load(inside(a.repetitive_mlp))) |
| if not 1 <= a.columns < weight.shape[1]: |
| raise ValueError('Неверное количество столбцов') |
| columns = np.argsort(score)[-a.columns:] |
| edited = project_writers(weight, direction, columns, a.alpha) |
| record = save_edit(a.model, a.out, a.tensor, edited, a.hypothesis) |
| save_json(inside(a.out) / (a.tensor + '.analysis.json'), {'columns': columns.tolist(), 'alpha': a.alpha, |
| 'scores': score[columns].tolist(), |
| 'consistency': consistency[columns].tolist(), |
| 'input_files': [a.productive, a.repetitive, a.productive_mlp, a.repetitive_mlp]}) |
| np.save(inside(a.out) / (a.tensor + '.direction.npy'), direction) |
| print(json.dumps(record, ensure_ascii=False)) |
|
|
