GPT-2 J-lens: use local Shakespeare + fix summary

1 files changed, 11 insertions, 6 deletions

diff --git a/src/gpt2_jlens.py b/src/gpt2_jlens.py
index b55a65f..232cf4a 100644
--- a/src/gpt2_jlens.py
+++ b/src/gpt2_jlens.py
@@ -37,12 +37,11 @@ def main():
     vocab_size = model.config.vocab_size
     print(f"GPT-2 Small: {n_layers} layers, d={d_model}, V={vocab_size}, V/d={vocab_size/d_model:.1f}x")
     
-    # Load corpus
-    text = requests.get(
-        "https://raw.githubusercontent.com/karpathy/nanoGPT/master/data/shakespeare_char/input.txt"
-    ).text[:200000]
+    # Load corpus from local Shakespeare file
+    with open('data/shakespeare_char/input.txt') as f:
+        text = f.read()[:500000]
     tokens = tokenizer(text, return_tensors='np', truncation=True, max_length=2000)['input_ids'][0]
-    print(f"Corpus: {len(tokens)} tokens")
+    print(f"Corpus: {len(tokens)} tokens, {len(set(tokens))} unique")
     
     # Estimate token frequencies in our corpus
     from collections import Counter
@@ -133,11 +132,17 @@ def main():
             print(f"    '{tok_str}' (freq={freqs[tid]:.3f}%): norm={n:.4f}")
         print(f"  Pearson r(norm, freq): {corr:.3f}")
     
+    # Summary
+    avg_corr = np.mean([np.corrcoef(
+        np.array(list(accum[l][t][0]/max(accum[l][t][1],1) for t in sample_tokens)),
+        np.array([freq.get(t,0)/total*100 for t in sample_tokens])
+    )[0,1] for l in layers_to_test])
+    
     print(f"\n{'='*60}")
     print("RESULTS SUMMARY")
     print(f"{'='*60}")
     print(f"  nanoGPT (V/d=0.2x): full-rank J-space, r=-0.65 freq correlation")
-    print(f"  GPT-2   (V/d=65x):  {'r=' + str(corr)}")
+    print(f"  GPT-2   (V/d=65x):  avg r={avg_corr:.3f}")
     print(f"  HYPOTHESIS: frequency anti-correlation persists at scale")