updated hubert model safetensors conversion

lib/infer_pack/loaders.py

@@ -1,3 +1,6 @@
+from collections import defaultdict
+import os
+from typing import Dict, List, Optional
 import torch
 from transformers import HubertModel, HubertConfig
 import safetensors.torch as st
@@ -14,28 +17,113 @@ class HubertModelWithFinalProj(HubertModel):
         self.final_proj = torch.nn.Linear(config.hidden_size, config.classifier_proj_size)
 
     @staticmethod
-    def load_safetensors(path: str, device="cpu"):
+    def from_safetensors(path: str, device="cpu", framework="pt"):
         assert path.endswith(".safetensors"), f"{path} must end with '.safetensors'"
         
-        with safe_open(path, framework="pt", device="cpu") as f:
+        with safe_open(path, framework=framework, device="cpu") as f:
             metadata = f.metadata()
             state_dict = {}
             for key in f.keys():
                 state_dict[key] = f.get_tensor(key)
         model = HubertModelWithFinalProj(HubertConfig.from_dict(json.loads(metadata["config"])))
         model.load_state_dict(state_dict=state_dict)
+        model.eval()
         return model.to(device)
-    
-    def save_safetensors(self, path: str):
-        assert path.endswith(".safetensors"), f"{path} must end with '.safetensors'"
-        
-        with open(path,"wb") as f:
-            state_dict = self.state_dict()
-            f.write(st.save(state_dict,dict(config=json.dumps(self.config.to_dict()))))
+
+    def to_safetensors(self, sf_filename: str, discard_names: Optional[List[str]]=[]):
+        assert hasattr(self, "state_dict"), f"Model does not have state_dict"
+
+        state_dict = self.state_dict()
+        to_removes = _remove_duplicate_names(state_dict, discard_names=discard_names)
+
+        metadata = dict(format="pt", config=json.dumps(self.config.to_dict()))
+        for kept_name, to_remove_group in to_removes.items():
+            for to_remove in to_remove_group:
+                if to_remove not in metadata:
+                    metadata[to_remove] = kept_name
+                del state_dict[to_remove]
+        # Force tensors to be contiguous
+        state_dict = {k: v.contiguous() for k, v in state_dict.items()}
+
+        dirname = os.path.dirname(sf_filename)
+        os.makedirs(dirname, exist_ok=True)
+        st.save_file(state_dict, sf_filename, metadata=metadata)
+        reloaded = HubertModelWithFinalProj.from_safetensors(sf_filename, device=self.device)
+        assert self.equals(reloaded), f"{sf_filename} does not equal source!"
+        del reloaded
 
     def extract_features(self, source: torch.Tensor, version="v2", **kwargs):
         with torch.no_grad():
             output_layer = 9 if version == "v1" else 12
-            output = self(source.to(self.config.torch_dtype), output_hidden_states=True)["hidden_states"][output_layer]
+            # https://github.com/facebookresearch/fairseq/blob/main/fairseq/models/hubert/hubert.py#L476
+            output = self(source.to(self.config.torch_dtype), output_hidden_states=True)["hidden_states"][output_layer-1]
             features = self.final_proj(output) if version == "v1" else output
-        return features.to(source.dtype)
+        return features.to(source.dtype)
+    
+    def equals(self, other: "HubertModelWithFinalProj"):
+        # check config
+        if self.config.to_dict()!=other.config.to_dict(): return False
+
+        # check parameters
+        self_dict = self.state_dict()
+        other_dict = other.state_dict()
+        for k in other_dict:
+            if not torch.equal(self_dict[k], other_dict[k]):
+                print(f"{k} is not the same")
+                return False
+        # for p1, p2 in zip(self.parameters(), other.parameters()):
+        #     if p1.data.ne(p2.data).sum() > 0: return False
+            
+        # check output
+        inputs = torch.ones((1,16000))
+        if self(inputs).last_hidden_state.data.ne(other(inputs).last_hidden_state.data).sum() > 0: return False
+
+        return True
+    
+# from: https://github.com/huggingface/safetensors/blob/main/bindings/python/convert.py#L36
+def _remove_duplicate_names(
+    state_dict: Dict[str, torch.Tensor],
+    *,
+    preferred_names: List[str] = None,
+    discard_names: List[str] = None,
+) -> Dict[str, List[str]]:
+    if preferred_names is None:
+        preferred_names = []
+    preferred_names = set(preferred_names)
+    if discard_names is None:
+        discard_names = []
+    discard_names = set(discard_names)
+
+    shareds = st._find_shared_tensors(state_dict)
+    to_remove = defaultdict(list)
+    for shared in shareds:
+        complete_names = set([name for name in shared if st._is_complete(state_dict[name])])
+        if not complete_names:
+            if len(shared) == 1:
+                # Force contiguous
+                name = list(shared)[0]
+                state_dict[name] = state_dict[name].clone()
+                complete_names = {name}
+            else:
+                raise RuntimeError(
+                    f"Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue."
+                )
+
+        keep_name = sorted(list(complete_names))[0]
+
+        # Mecanism to preferentially select keys to keep
+        # coming from the on-disk file to allow
+        # loading models saved with a different choice
+        # of keep_name
+        preferred = complete_names.difference(discard_names)
+        if preferred:
+            keep_name = sorted(list(preferred))[0]
+
+        if preferred_names:
+            preferred = preferred_names.intersection(complete_names)
+            if preferred:
+                keep_name = sorted(list(preferred))[0]
+        for name in sorted(shared):
+            if name != keep_name:
+                to_remove[keep_name].append(name)
+    return to_remove

lib/model_utils.py

@@ -19,7 +19,7 @@ def get_hash(model_path):
 def load_hubert(model_path: str, config):
     try:
         if model_path.endswith(".safetensors"):
-            return HubertModelWithFinalProj.load_safetensors(model_path, device=config.device)
+            return HubertModelWithFinalProj.from_safetensors(model_path, device=config.device)
         else:
             from fairseq import checkpoint_utils
             models, _, _ = checkpoint_utils.load_model_ensemble_and_task([model_path],suffix="",)