diff options
Diffstat (limited to 'src/jlens_v3.py')
| -rw-r--r-- | src/jlens_v3.py | 9 |
1 files changed, 6 insertions, 3 deletions
diff --git a/src/jlens_v3.py b/src/jlens_v3.py index 7316fe8..b5f59ed 100644 --- a/src/jlens_v3.py +++ b/src/jlens_v3.py @@ -169,6 +169,8 @@ def main(): ap.add_argument('--n_prompts', type=int, default=20) ap.add_argument('--batch_size', type=int, default=16) ap.add_argument('--layers', default='2,3,4', help='comma-separated layer indices') + ap.add_argument('--output_dir', default='outputs/jlens_v3', + help='directory for per-layer result artifacts') ap.add_argument('--device', default='cuda') ap.add_argument('--chunk', type=int, default=32) args = ap.parse_args() @@ -210,7 +212,7 @@ def main(): if layer_idx == n_layer - 1: # Validation: J_{L-1} should be identity, so faithful vectors == W_U rows sims = torch.nn.functional.cosine_similarity( - faithful_vecs.float(), W_U.float(), dim=1) + faithful_vecs.float(), W_U.cpu().float(), dim=1) print(f" [validation] last layer: mean cos-sim(faithful, W_U rows) = " f"{sims.mean().item():.4f} (expect ~1.0 if J=identity)") @@ -239,10 +241,11 @@ def main(): for tid, n in srt[-5:]: print(f" '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}") + os.makedirs(args.output_dir, exist_ok=True) torch.save({'faithful_vecs': faithful_vecs, 'faithful_norms': faithful_norms, 'proxy_norms': proxy_norms}, - f'outputs/jlens_v3_layer{layer_idx}.pt') - print(f" Saved outputs/jlens_v3_layer{layer_idx}.pt") + os.path.join(args.output_dir, f'layer{layer_idx}.pt')) + print(f" Saved {os.path.join(args.output_dir, f'layer{layer_idx}.pt')}") print("\nDONE.") |