Update README.md

README.md

@@ -222,6 +222,14 @@ tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
 ckpt = torch.load(model_path, map_location="cpu")
 state = ckpt["model_state"] if "model_state" in ckpt else ckpt
 
+head_state = {
+    k.replace("score.", ""): v
+    for k, v in state.items()
+    if k.startswith("score.")
+}
+
+assert set(head_state.keys()) == {"weight", "bias"}
+
 # Initialize scoring head (single linear layer)
 hidden_size = base_model.config.hidden_size
 scoring_head = torch.nn.Linear(hidden_size, 1)