Update README.md

README.md

@@ -68,7 +68,7 @@ Input Text (Reasoning Trace)
     ↓
 [Frozen Base LM Encoder]  ← Pre-trained, frozen during training
     ↓
-[Final Token (EOS) Pooling]
+[Final Non-Padding Token Pooling (attention-mask aware)]
     ↓
 [Lightweight Linear Head]    ← Only these parameters are trained
     ↓
@@ -219,12 +219,13 @@ base_model = AutoModel.from_pretrained("facebook/opt-1.3b")
 tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
 
 # Load the trained scoring head weights
-scoring_head_weights = torch.load(model_path, map_location="cpu")
+ckpt = torch.load(model_path, map_location="cpu")
+state = ckpt["model_state"] if "model_state" in ckpt else ckpt
 
 # Initialize scoring head (single linear layer)
 hidden_size = base_model.config.hidden_size
 scoring_head = torch.nn.Linear(hidden_size, 1)
-scoring_head.load_state_dict(scoring_head_weights)
+scoring_head.load_state_dict(state)
 
 # Move to device
 device = "cuda" if torch.cuda.is_available() else "cpu"