add pooler weights to sdadapter

Signed-off-by: HuiyingLi <[email protected]>

Author: HuiyingLi

nemo_automodel/components/models/biencoder/state_dict_adapter.py

@@ -48,6 +48,8 @@ def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
             if key.startswith("lm_q."):
                 new_key = key.replace("lm_q.", "model.")
                 hf_state_dict[new_key] = value
+            elif key.startswith("linear_pooler."):
+                hf_state_dict[key] = value
 
         return hf_state_dict
 
@@ -76,6 +78,8 @@ def from_hf(
                 biencoder_state_dict[new_key_q] = value
                 new_key_p = key.replace("model.", "lm_p.")
                 biencoder_state_dict[new_key_p] = value
+            elif key.startswith("linear_pooler."):
+                biencoder_state_dict[key] = value
 
         return biencoder_state_dict
 
@@ -94,6 +98,8 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
         if fqn.startswith("lm_q."):
             new_fqn = fqn.replace("lm_q.", "model.")
             return [(new_fqn, tensor)]
+        if fqn.startswith("linear_pooler."):
+            return [(fqn, tensor)]
 
         # Skip tensors that are not part of lm_q
         return []