Upload folder using huggingface_hub

modeling_spear.py

@@ -1,10 +1,8 @@
 # modeling_spear.py
 
-import torch
-
 from transformers import PreTrainedModel
-from configuration_spear import SpearConfig
-from spear_model import SpearModel as model
+from .configuration_spear import SpearConfig
+from .spear_model import SpearModel as model
 
 
 class SpearModel(PreTrainedModel):
@@ -18,60 +16,4 @@ class SpearModel(PreTrainedModel):
         return self.model(*args, **kwargs)
 
     def load_audio(self, audio_path):
-        return self.model.load_audio(audio_path)
-    
-    @classmethod
-    def from_legacy_checkpoint(cls, path, config):
-        model = cls(config)
-        ckpt = torch.load(path)["model"]
-        info = model.model.model.load_state_dict(ckpt, strict=False)
-        print(info)
-        return model
-    
-def export_to_hf():
-    ckpt = "/mnt/shared-storage-user/housiyuan/xiaoyu/models/spear_encoders/94m-uni-v2-dual-domain-mvq/iter-400000-avg-4.pt"
-    config = SpearConfig()
-    my_model = SpearModel.from_legacy_checkpoint(ckpt, config)
-    my_model.save_pretrained("/mnt/shared-storage-user/housiyuan/xiaoyu/models/spear_encoders_hf/spear_base_speech_audio")
-
-
-def _test_from_pretrained():
-    device = torch.device("cpu")
-    if torch.cuda.is_available():
-        device = torch.device("cuda")
-    
-    audio_file = [
-        "common_voice_af_39597042.wav",
-        # "1284-1180-0027.flac",
-    ]
-    
-    
-    config= SpearConfig.from_pretrained("/mnt/shared-storage-user/housiyuan/xiaoyu/models/spear_encoders_hf/spear_base_speech_audio")
-    my_model= SpearModel.from_pretrained("/mnt/shared-storage-user/housiyuan/xiaoyu/models/spear_encoders_hf/spear_base_speech_audio", config=config)
-    my_model.eval()
-    my_model.to(device)
-    
-    num_params = sum([p.numel() for p in my_model.parameters()])
-    print(f"A total of {num_params} model parameters.")
-    
-    audio, audio_len = my_model.load_audio(audio_file)
-    audio = audio.to(device)
-    audio_len = audio_len.to(device)
-    with torch.no_grad():
-        outputs = my_model(audio, audio_len)
-    
-    encoder_out = outputs["encoder_out"] # (N,T,C)
-    encoder_out_lens = outputs["encoder_out_lens"] # (N)
-    middle_out = outputs["hidden_states"] # list of (N,T,C)
-    
-    print(encoder_out)
-    print(encoder_out_lens)
-    print(middle_out[0].shape)
-    
-    
-if __name__=="__main__":
-    export_to_hf()
-    _test_from_pretrained()
-    
-    
-    
+        return self.model.load_audio(audio_path)

spear_model.py

@@ -26,8 +26,8 @@ import torch.nn as nn
 from torch.nn.utils.rnn import pad_sequence
 from torchaudio.compliance.kaldi import fbank as torch_fbank
 
-from configuration_spear import SpearConfig
-from zipformer import Zipformer2, Conv2dSubsampling
+from .configuration_spear import SpearConfig
+from .zipformer import Zipformer2, Conv2dSubsampling
 
 LOG_EPS=math.log(1e-10)
 SAMPLING_RATE=16000
@@ -237,11 +237,10 @@ class SpearEncoder(nn.Module):
         self.distillation_delta = distillation_delta
         
         if num_codebooks > 0:
-            from multi_quantization.prediction import JointCodebookLoss
+            from .spear_modules import JointCodebookLoss
             self.codebook_loss_net = JointCodebookLoss(
-                predictor_channels=encoder_dim,
+                input_dim=encoder_dim,
                 num_codebooks=num_codebooks * self.teacher_frame_ratio,
-                is_joint=False,
                 reduction="none",
             )
         else:

spear_modules.py

@@ -32,6 +32,58 @@ def logaddexp_onnx(x: Tensor, y: Tensor) -> Tensor:
     diff = torch.abs(x - y)
     return max_value + torch.log1p(torch.exp(-diff))
 
+class JointCodebookLoss(torch.nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 512,
+        num_codebooks: int = 16,
+        codebook_size: int = 256,
+        ignore_index: int = -100,
+        reduction: str = "none"
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.num_codebooks = num_codebooks
+        self.codebook_size = codebook_size
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+        
+        self.proj = nn.Linear(input_dim, num_codebooks * codebook_size)
+    
+    def forward_logprobs(self, input: torch.Tensor):
+        B,T,_ = input.shape
+        logits = self.proj(input)
+        logits = logits.view(B, T, self.num_codebooks, self.codebook_size) # (B,T,N,256)
+        log_probs = F.log_softmax(logits, dim=-1) # (B,T,N,256)
+        return log_probs
+    
+    
+    def forward(self, input, target, return_log_probs: bool = False):
+        # input: (B,T,C)
+        # target: (B,T,num_codebooks)
+
+        B,T,_ = input.shape
+        logits = self.proj(input)
+        logits = logits.view(B, T, self.num_codebooks, self.codebook_size) # (B,T,N,256)
+        
+        loss = F.cross_entropy(
+            logits.reshape(-1, self.codebook_size), 
+            target.reshape(-1),
+            ignore_index=self.ignore_index,
+            reduction=self.reduction
+        )
+        log_probs = None
+        if return_log_probs:
+            log_probs = F.log_softmax(logits, dim=-1)
+        
+        if self.reduction == "none":
+            loss = loss.view(B, T, self.num_codebooks)
+
+        if return_log_probs:
+            return loss, log_probs
+        
+        return loss
+
 
 # RuntimeError: Exporting the operator logaddexp to ONNX opset version
 # 14 is not supported. Please feel free to request support or submit

zipformer.py

@@ -23,7 +23,7 @@ from typing import List, Optional, Tuple, Union
 import logging
 import torch
 import random
-from spear_modules import (
+from .spear_modules import (
     Balancer,
     BiasNorm,
     Dropout2,