Update README.md

README.md

@@ -8,6 +8,7 @@ base_model:
 ---
 
 Multilingual speaker-invariant speech content encoder based on [SPIN method](https://arxiv.org/abs/2305.11072) taken from [FreeSVC](https://huggingface.co/alefiury/free-svc), converted to transformers library format.
+Usage:
 ```py
 from transformers import AutoProcessor, HubertModel
 processor = AutoProcessor.from_pretrained("hr16/multilingual_spin")
@@ -19,4 +20,332 @@ with torch.no_grad():
     outputs = model(**inputs)
 
 print(outputs.last_hidden_state.shape)
+```
+
+Conversion code:
+```py
+import re
+import torch
+def convert_hubert_s3prl_to_transformers(model, checkpoint_path: str):
+    # 1. Load checkpoint
+    state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    if "model" in state_dict:
+        state_dict = state_dict["model"]
+    elif "state_dict" in state_dict:
+        state_dict = state_dict["state_dict"]
+
+    hf_state_dict = {}
+
+    # 2. Map weights
+    for key, value in state_dict.items():
+        # Skip unused pretraining components
+        if (
+            key.startswith("pred_head")
+            or key.startswith("loss_module")
+            or "label_embs_concat" in key
+        ):
+            continue
+
+        # Mask embedding
+        if key == "encoder.model.mask_emb":
+            hf_state_dict["masked_spec_embed"] = value
+            continue
+
+        # Feature extractor group norm / conv layers
+        if "feature_extractor.conv_layers" in key:
+            # e.g.: encoder.model.feature_extractor.conv_layers.0.0.weight -> feature_extractor.conv_layers.0.conv.weight
+            # e.g.: encoder.model.feature_extractor.conv_layers.0.2.weight -> feature_extractor.conv_layers.0.layer_norm.weight
+            match = re.search(r"conv_layers\.(\d+)\.(\d+)\.(weight|bias)", key)
+            if match:
+                layer_idx, sub_idx, param_name = match.groups()
+                if sub_idx == "0":
+                    hf_key = f"feature_extractor.conv_layers.{layer_idx}.conv.{param_name}"
+                elif sub_idx == "2":
+                    hf_key = f"feature_extractor.conv_layers.{layer_idx}.layer_norm.{param_name}"
+                hf_state_dict[hf_key] = value
+            continue
+
+        # Feature projection layer norm
+        if key in ("encoder.model.layer_norm.weight", "encoder.model.layer_norm.bias"):
+            param_name = key.split(".")[-1]
+            hf_state_dict[f"feature_projection.layer_norm.{param_name}"] = value
+            continue
+
+        # Feature projection linear layer
+        if "post_extract_proj" in key:
+            param_name = key.split(".")[-1]
+            hf_state_dict[f"feature_projection.projection.{param_name}"] = value
+            continue
+
+        # Positional Convolution Embedding
+        if "pos_conv.0" in key:
+            param_name = key.split(".")[-1]
+            # Handles weight_v, weight_g, bias
+            hf_state_dict[f"encoder.pos_conv_embed.conv.{param_name}"] = value
+            continue
+
+        # Encoder pre-LayerNorm
+        if key in (
+            "encoder.model.encoder.layer_norm.weight",
+            "encoder.model.encoder.layer_norm.bias",
+        ):
+            param_name = key.split(".")[-1]
+            hf_state_dict[f"encoder.layer_norm.{param_name}"] = value
+            continue
+
+        # Encoder Transformer Layers
+        if "encoder.model.encoder.layers." in key:
+            # Map layer components
+            m = re.match(
+                r"encoder\.model\.encoder\.layers\.(\d+)\.(.+)", key
+            )
+            if m:
+                layer_idx, rest = m.groups()
+                prefix = f"encoder.layers.{layer_idx}"
+
+                if "self_attn." in rest:
+                    # self_attn.{q,k,v,out}_proj.{weight,bias}
+                    attn_component = rest.replace("self_attn.", "attention.")
+                    hf_state_dict[f"{prefix}.{attn_component}"] = value
+                elif "self_attn_layer_norm." in rest:
+                    param_name = rest.split(".")[-1]
+                    hf_state_dict[f"{prefix}.layer_norm.{param_name}"] = value
+                elif "fc1." in rest:
+                    param_name = rest.split(".")[-1]
+                    hf_state_dict[
+                        f"{prefix}.feed_forward.intermediate_dense.{param_name}"
+                    ] = value
+                elif "fc2." in rest:
+                    param_name = rest.split(".")[-1]
+                    hf_state_dict[
+                        f"{prefix}.feed_forward.output_dense.{param_name}"
+                    ] = value
+                elif "final_layer_norm." in rest:
+                    param_name = rest.split(".")[-1]
+                    hf_state_dict[f"{prefix}.final_layer_norm.{param_name}"] = value
+            continue
+
+        print(f"Unmapped key: {key}")
+    model.load_state_dict(hf_state_dict)
+    return model
+
+from transformers import HubertModel
+orig_hubert = HubertModel.from_pretrained("facebook/hubert-base-ls960")
+spin = convert_hubert_s3prl_to_transformers(HubertModel.from_pretrained("facebook/hubert-base-ls960"), "spin.ckpt")
+spin.push_to_hub("hr16/multilingual_spin")
+```
+
+State dict mean-square differences:
+```
+masked_spec_embed diff: 0.0
+feature_extractor.conv_layers.0.conv.weight diff: 0.0
+feature_extractor.conv_layers.0.layer_norm.weight diff: 0.0
+feature_extractor.conv_layers.0.layer_norm.bias diff: 0.0
+feature_extractor.conv_layers.1.conv.weight diff: 0.0
+feature_extractor.conv_layers.2.conv.weight diff: 0.0
+feature_extractor.conv_layers.3.conv.weight diff: 0.0
+feature_extractor.conv_layers.4.conv.weight diff: 0.0
+feature_extractor.conv_layers.5.conv.weight diff: 0.0
+feature_extractor.conv_layers.6.conv.weight diff: 0.0
+feature_projection.layer_norm.weight diff: 0.0
+feature_projection.layer_norm.bias diff: 0.0
+feature_projection.projection.weight diff: 0.0
+feature_projection.projection.bias diff: 0.0
+encoder.pos_conv_embed.conv.bias diff: 0.0
+encoder.pos_conv_embed.conv.parametrizations.weight.original0 diff: 0.0
+encoder.pos_conv_embed.conv.parametrizations.weight.original1 diff: 0.0
+encoder.layer_norm.weight diff: 0.0
+encoder.layer_norm.bias diff: 0.0
+encoder.layers.0.attention.k_proj.weight diff: 0.0
+encoder.layers.0.attention.k_proj.bias diff: 0.0
+encoder.layers.0.attention.v_proj.weight diff: 0.0
+encoder.layers.0.attention.v_proj.bias diff: 0.0
+encoder.layers.0.attention.q_proj.weight diff: 0.0
+encoder.layers.0.attention.q_proj.bias diff: 0.0
+encoder.layers.0.attention.out_proj.weight diff: 0.0
+encoder.layers.0.attention.out_proj.bias diff: 0.0
+encoder.layers.0.layer_norm.weight diff: 0.0
+encoder.layers.0.layer_norm.bias diff: 0.0
+encoder.layers.0.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.0.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.0.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.0.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.0.final_layer_norm.weight diff: 0.0
+encoder.layers.0.final_layer_norm.bias diff: 0.0
+encoder.layers.1.attention.k_proj.weight diff: 0.0
+encoder.layers.1.attention.k_proj.bias diff: 0.0
+encoder.layers.1.attention.v_proj.weight diff: 0.0
+encoder.layers.1.attention.v_proj.bias diff: 0.0
+encoder.layers.1.attention.q_proj.weight diff: 0.0
+encoder.layers.1.attention.q_proj.bias diff: 0.0
+encoder.layers.1.attention.out_proj.weight diff: 0.0
+encoder.layers.1.attention.out_proj.bias diff: 0.0
+encoder.layers.1.layer_norm.weight diff: 0.0
+encoder.layers.1.layer_norm.bias diff: 0.0
+encoder.layers.1.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.1.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.1.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.1.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.1.final_layer_norm.weight diff: 0.0
+encoder.layers.1.final_layer_norm.bias diff: 0.0
+encoder.layers.2.attention.k_proj.weight diff: 0.0
+encoder.layers.2.attention.k_proj.bias diff: 0.0
+encoder.layers.2.attention.v_proj.weight diff: 0.0
+encoder.layers.2.attention.v_proj.bias diff: 0.0
+encoder.layers.2.attention.q_proj.weight diff: 0.0
+encoder.layers.2.attention.q_proj.bias diff: 0.0
+encoder.layers.2.attention.out_proj.weight diff: 0.0
+encoder.layers.2.attention.out_proj.bias diff: 0.0
+encoder.layers.2.layer_norm.weight diff: 0.0
+encoder.layers.2.layer_norm.bias diff: 0.0
+encoder.layers.2.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.2.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.2.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.2.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.2.final_layer_norm.weight diff: 0.0
+encoder.layers.2.final_layer_norm.bias diff: 0.0
+encoder.layers.3.attention.k_proj.weight diff: 0.0
+encoder.layers.3.attention.k_proj.bias diff: 0.0
+encoder.layers.3.attention.v_proj.weight diff: 0.0
+encoder.layers.3.attention.v_proj.bias diff: 0.0
+encoder.layers.3.attention.q_proj.weight diff: 0.0
+encoder.layers.3.attention.q_proj.bias diff: 0.0
+encoder.layers.3.attention.out_proj.weight diff: 0.0
+encoder.layers.3.attention.out_proj.bias diff: 0.0
+encoder.layers.3.layer_norm.weight diff: 0.0
+encoder.layers.3.layer_norm.bias diff: 0.0
+encoder.layers.3.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.3.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.3.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.3.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.3.final_layer_norm.weight diff: 0.0
+encoder.layers.3.final_layer_norm.bias diff: 0.0
+encoder.layers.4.attention.k_proj.weight diff: 0.0
+encoder.layers.4.attention.k_proj.bias diff: 0.0
+encoder.layers.4.attention.v_proj.weight diff: 0.0
+encoder.layers.4.attention.v_proj.bias diff: 0.0
+encoder.layers.4.attention.q_proj.weight diff: 0.0
+encoder.layers.4.attention.q_proj.bias diff: 0.0
+encoder.layers.4.attention.out_proj.weight diff: 0.0
+encoder.layers.4.attention.out_proj.bias diff: 0.0
+encoder.layers.4.layer_norm.weight diff: 0.0
+encoder.layers.4.layer_norm.bias diff: 0.0
+encoder.layers.4.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.4.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.4.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.4.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.4.final_layer_norm.weight diff: 0.0
+encoder.layers.4.final_layer_norm.bias diff: 0.0
+encoder.layers.5.attention.k_proj.weight diff: 0.0
+encoder.layers.5.attention.k_proj.bias diff: 0.0
+encoder.layers.5.attention.v_proj.weight diff: 0.0
+encoder.layers.5.attention.v_proj.bias diff: 0.0
+encoder.layers.5.attention.q_proj.weight diff: 0.0
+encoder.layers.5.attention.q_proj.bias diff: 0.0
+encoder.layers.5.attention.out_proj.weight diff: 0.0
+encoder.layers.5.attention.out_proj.bias diff: 0.0
+encoder.layers.5.layer_norm.weight diff: 0.0
+encoder.layers.5.layer_norm.bias diff: 0.0
+encoder.layers.5.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.5.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.5.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.5.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.5.final_layer_norm.weight diff: 0.0
+encoder.layers.5.final_layer_norm.bias diff: 0.0
+encoder.layers.6.attention.k_proj.weight diff: 0.0
+encoder.layers.6.attention.k_proj.bias diff: 0.0
+encoder.layers.6.attention.v_proj.weight diff: 0.0
+encoder.layers.6.attention.v_proj.bias diff: 0.0
+encoder.layers.6.attention.q_proj.weight diff: 0.0
+encoder.layers.6.attention.q_proj.bias diff: 0.0
+encoder.layers.6.attention.out_proj.weight diff: 0.0
+encoder.layers.6.attention.out_proj.bias diff: 0.0
+encoder.layers.6.layer_norm.weight diff: 0.0
+encoder.layers.6.layer_norm.bias diff: 0.0
+encoder.layers.6.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.6.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.6.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.6.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.6.final_layer_norm.weight diff: 0.0
+encoder.layers.6.final_layer_norm.bias diff: 0.0
+encoder.layers.7.attention.k_proj.weight diff: 0.0
+encoder.layers.7.attention.k_proj.bias diff: 0.0
+encoder.layers.7.attention.v_proj.weight diff: 0.0
+encoder.layers.7.attention.v_proj.bias diff: 0.0
+encoder.layers.7.attention.q_proj.weight diff: 0.0
+encoder.layers.7.attention.q_proj.bias diff: 0.0
+encoder.layers.7.attention.out_proj.weight diff: 0.0
+encoder.layers.7.attention.out_proj.bias diff: 0.0
+encoder.layers.7.layer_norm.weight diff: 0.0
+encoder.layers.7.layer_norm.bias diff: 0.0
+encoder.layers.7.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.7.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.7.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.7.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.7.final_layer_norm.weight diff: 0.0
+encoder.layers.7.final_layer_norm.bias diff: 0.0
+encoder.layers.8.attention.k_proj.weight diff: 0.0
+encoder.layers.8.attention.k_proj.bias diff: 0.0
+encoder.layers.8.attention.v_proj.weight diff: 0.0
+encoder.layers.8.attention.v_proj.bias diff: 0.0
+encoder.layers.8.attention.q_proj.weight diff: 0.0
+encoder.layers.8.attention.q_proj.bias diff: 0.0
+encoder.layers.8.attention.out_proj.weight diff: 0.0
+encoder.layers.8.attention.out_proj.bias diff: 0.0
+encoder.layers.8.layer_norm.weight diff: 0.0
+encoder.layers.8.layer_norm.bias diff: 0.0
+encoder.layers.8.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.8.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.8.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.8.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.8.final_layer_norm.weight diff: 0.0
+encoder.layers.8.final_layer_norm.bias diff: 0.0
+encoder.layers.9.attention.k_proj.weight diff: 0.0
+encoder.layers.9.attention.k_proj.bias diff: 0.0
+encoder.layers.9.attention.v_proj.weight diff: 0.0
+encoder.layers.9.attention.v_proj.bias diff: 0.0
+encoder.layers.9.attention.q_proj.weight diff: 0.0
+encoder.layers.9.attention.q_proj.bias diff: 0.0
+encoder.layers.9.attention.out_proj.weight diff: 0.0
+encoder.layers.9.attention.out_proj.bias diff: 0.0
+encoder.layers.9.layer_norm.weight diff: 0.0
+encoder.layers.9.layer_norm.bias diff: 0.0
+encoder.layers.9.feed_forward.intermediate_dense.weight diff: 0.0
+encoder.layers.9.feed_forward.intermediate_dense.bias diff: 0.0
+encoder.layers.9.feed_forward.output_dense.weight diff: 0.0
+encoder.layers.9.feed_forward.output_dense.bias diff: 0.0
+encoder.layers.9.final_layer_norm.weight diff: 0.0
+encoder.layers.9.final_layer_norm.bias diff: 0.0
+encoder.layers.10.attention.k_proj.weight diff: 0.0026198839768767357
+encoder.layers.10.attention.k_proj.bias diff: 0.0010964443208649755
+encoder.layers.10.attention.v_proj.weight diff: 0.001084218267351389
+encoder.layers.10.attention.v_proj.bias diff: 8.861951209837571e-05
+encoder.layers.10.attention.q_proj.weight diff: 0.002644478576257825
+encoder.layers.10.attention.q_proj.bias diff: 0.0009331119363196194
+encoder.layers.10.attention.out_proj.weight diff: 0.0009169431868940592
+encoder.layers.10.attention.out_proj.bias diff: 0.00011753084982046857
+encoder.layers.10.layer_norm.weight diff: 0.004704533610492945
+encoder.layers.10.layer_norm.bias diff: 0.00040463244658894837
+encoder.layers.10.feed_forward.intermediate_dense.weight diff: 0.002256407169625163
+encoder.layers.10.feed_forward.intermediate_dense.bias diff: 8.606249321019277e-05
+encoder.layers.10.feed_forward.output_dense.weight diff: 0.0017423474928364158
+encoder.layers.10.feed_forward.output_dense.bias diff: 7.794254634063691e-05
+encoder.layers.10.final_layer_norm.weight diff: 0.003364397445693612
+encoder.layers.10.final_layer_norm.bias diff: 0.000663460697978735
+encoder.layers.11.attention.k_proj.weight diff: 0.0019817741122096777
+encoder.layers.11.attention.k_proj.bias diff: 0.032148730009794235
+encoder.layers.11.attention.v_proj.weight diff: 0.0010857016313821077
+encoder.layers.11.attention.v_proj.bias diff: 0.00020938993839081377
+encoder.layers.11.attention.q_proj.weight diff: 0.002357392804697156
+encoder.layers.11.attention.q_proj.bias diff: 0.0008044671267271042
+encoder.layers.11.attention.out_proj.weight diff: 0.0011331378482282162
+encoder.layers.11.attention.out_proj.bias diff: 0.00014901244139764458
+encoder.layers.11.layer_norm.weight diff: 0.0027753293979912996
+encoder.layers.11.layer_norm.bias diff: 0.0056587993167340755
+encoder.layers.11.feed_forward.intermediate_dense.weight diff: 0.0032594590447843075
+encoder.layers.11.feed_forward.intermediate_dense.bias diff: 0.0002475830551702529
+encoder.layers.11.feed_forward.output_dense.weight diff: 0.002621536375954747
+encoder.layers.11.feed_forward.output_dense.bias diff: 2.9405724490061402e-05
+encoder.layers.11.final_layer_norm.weight diff: 0.002755063585937023
+encoder.layers.11.final_layer_norm.bias diff: 0.00021274278697092086
 ```