From b17200147c72c415b81587a67af9d592cd60cd33 Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Mon, 24 Jun 2024 22:50:31 -0400 Subject: [PATCH] updated hubert model safetensors conversion --- lib/infer_pack/loaders.py | 110 ++++++++++++++++++++++++++++++++++---- lib/model_utils.py | 2 +- 2 files changed, 100 insertions(+), 12 deletions(-) diff --git a/lib/infer_pack/loaders.py b/lib/infer_pack/loaders.py index f282a83..3098cd1 100644 --- a/lib/infer_pack/loaders.py +++ b/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) \ No newline at end of file + 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 \ No newline at end of file diff --git a/lib/model_utils.py b/lib/model_utils.py index babb14f..905ae5b 100644 --- a/lib/model_utils.py +++ b/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="",)