| import argparse |
| import json |
| import numpy as np |
| from common import ROOT, inside, save_json |
| from weights import load |
|
|
|
|
| def mean_window(directory, metadata, text, signal, window): |
| total, count = None, 0 |
| for file in metadata['files']: |
| if file['text'] != text or file['signal'] != signal: |
| continue |
| start, end = file['token_start'], file['token_end'] |
| left, right = max(start, window[0]), min(end, window[1]) |
| if right <= left: |
| continue |
| array = np.load(directory / file['name'], mmap_mode='r')[0, left-start:right-start] |
| value = array.sum(0, dtype=np.float64) |
| total = value if total is None else total + value |
| count += right-left |
| if count != window[1]-window[0]: |
| raise ValueError(f'Incomplete {signal} window for text {text}: {count}') |
| return (total/count).astype(np.float32) |
|
|
|
|
| def branch_logps(directory, metadata, labels): |
| result = [] |
| for row in labels: |
| values = [] |
| for file in metadata['files']: |
| if file['text'] != row['text'] or file['signal'] != 'token_logprobs': |
| continue |
| array = np.load(directory / file['name']) |
| start = max(0, row['branch_start']-1-file['token_start']) |
| end = min(len(array), row.get('branch_end',file['token_start']+len(array)+1)-1-file['token_start']) |
| if end > start: |
| values.extend(array[start:end].tolist()) |
| if not values: |
| raise ValueError('Missing branch scores') |
| result.append(dict(row, mean_logp=float(np.mean(values)), logp_sum=float(np.sum(values)), scored_tokens=len(values))) |
| return result |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument('--probe', required=True) |
| p.add_argument('--labels', default=str(ROOT/'configs/matched_probe_001_labels.json')) |
| p.add_argument('--out', required=True) |
| p.add_argument('--model', default=str(ROOT/'models/original')) |
| p.add_argument('--logp-only', action='store_true') |
| p.add_argument('--protect-productive', action='store_true') |
| args = p.parse_args() |
| directory, out = inside(args.probe), inside(args.out) |
| metadata = json.loads((directory/'meta.json').read_text()) |
| if not metadata.get('complete'): |
| raise ValueError('Probe is incomplete') |
| labels = json.loads(inside(args.labels).read_text()) |
| out.mkdir(parents=True, exist_ok=True) |
| save_json(out/'branch_logps.json', branch_logps(directory, metadata, labels)) |
| if args.logp_only: |
| return |
| arrays, ranking = {}, [] |
| discovery = [r for r in labels if r['split']=='discovery'] |
| pairs = sorted({r['pair'] for r in discovery}) |
| for layer in range(24): |
| hidden, mlp = {}, {} |
| for row in discovery: |
| key = (row['pair'], row['condition']) |
| hidden[key] = mean_window(directory, metadata, row['text'], f'hidden_{layer+1}', row['mean_window']) |
| mlp[key] = mean_window(directory, metadata, row['text'], f'mlp_act_{layer}', row['mean_window']) |
| hp = np.stack([hidden[(pair,'productive')] for pair in pairs]) |
| hr = np.stack([hidden[(pair,'repetitive')] for pair in pairs]) |
| dh = hr-hp |
| dm = np.stack([mlp[(pair,'repetitive')]-mlp[(pair,'productive')] for pair in pairs]) |
| direction = dh.mean(0) |
| if args.protect_productive: |
| _, singular, basis = np.linalg.svd(hp,full_matrices=False) |
| basis = basis[singular>singular[0]*1e-6] |
| direction = direction-basis.T@(basis@direction) |
| magnitude = np.linalg.norm(direction) |
| direction = direction/magnitude |
| weight = load(args.model, f'model.language_model.layers.{layer}.mlp.down_proj.weight').float().numpy() |
| projection = direction @ weight |
| scores = dm * projection |
| mean = scores.mean(0) |
| columns = np.argsort(mean)[-64:][::-1].copy() |
| arrays[f'u_{layer}'] = direction |
| arrays[f'columns_{layer}'] = columns |
| arrays[f'scores_{layer}'] = mean |
| arrays[f'pair_delta_mlp_{layer}'] = dm |
| arrays[f'productive_mlp_{layer}'] = np.stack([mlp[(pair,'productive')] for pair in pairs]) |
| arrays[f'repetitive_mlp_{layer}'] = np.stack([mlp[(pair,'repetitive')] for pair in pairs]) |
| arrays[f'pair_delta_hidden_{layer}'] = dh |
| arrays[f'productive_hidden_{layer}'] = hp |
| arrays[f'repetitive_hidden_{layer}'] = hr |
| ranking.append({'layer':layer,'direction_norm':float(magnitude),'top64_writer_contrast':float(mean[columns].sum()), |
| 'top64_fraction_positive_contrast':float(mean[columns].sum()/np.maximum(mean,0).sum()), |
| 'direction_pair_cosines':(dh@direction/np.linalg.norm(dh,axis=1)).tolist(), |
| 'top_neurons':[{'neuron':int(j),'contrast':float(mean[j]),'positive_pairs':int((scores[:,j]>0).sum())} for j in columns[:12]], |
| 'top64_consistent_pairs_mean':float((scores[:,columns]>0).mean()), |
| 'productive_protection':args.protect_productive, |
| 'max_productive_projection':float(np.max(np.abs(hp@direction)))}) |
| np.savez(out/'directions_and_writers.npz', **arrays) |
| save_json(out/'layer_ranking.json', sorted(ranking,key=lambda x:x['top64_writer_contrast'],reverse=True)) |
| print(json.dumps(sorted(ranking,key=lambda x:x['top64_writer_contrast'],reverse=True)[:5],indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
