Update README.md

README.md

@@ -1,3 +1,27 @@
----
-license: mit
----
+---
+license: mit
+---
+Linear probe checkpoints for https://footprints.baulab.info
+
+To load a Llama-2-7b checkpoint at layer 0 and target index -3:
+
+```
+import torch 
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+
+class LinearModel(nn.Module):
+ def __init__(self, input_size, output_size, bias=False):
+ super(LinearModel, self).__init__()
+ self.fc = nn.Linear(input_size, output_size, bias=bias)
+ def forward(self, x):
+ output = self.fc(x)
+ return output
+
+checkpoint_path = hf_hub_download(
+ repo_id="sfeucht/footprints", 
+ filename="llama-2-7b/layer0_tgtidx-3.ckpt"
+)
+probe = LinearModel(4096, 32000)
+probe.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
+```