Update model_hf.py

model_hf.py

@@ -9,6 +9,7 @@ from torch import Tensor
 from .config_ssl import SSLConfig
 from huggingface_hub import hf_hub_download
 from transformers import Wav2Vec2ForPreTraining
+import fairseq
 
 ___author__ = "Hemlata Tak"
 __email__ = "tak@eurecom.fr"
@@ -23,11 +24,11 @@ class SSLModel(nn.Module):
         super(SSLModel, self).__init__()
         # eliminate fairseq dependency 
         # facebook/wav2vec2-xls-r-300m
-        # repo_id = "facebook/wav2vec2-xlsr-300m"
-        model = Wav2Vec2.from_pretrained("facebook/wav2vec2-xls-r-300m")
-        # cp_path = hf_hub_download(repo_id=repo_id, filename=fname)   # Change the pre-trained XLSR model path. 
-        # model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([cp_path])
-        self.model = model
+        repo_id = "ash56/ssl-aasist"
+        fname = "xlsr2_300m.pt"
+        cp_path = hf_hub_download(repo_id=repo_id, filename=fname)   # Change the pre-trained XLSR model path. 
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([cp_path])
+        self.model = model[0]
         self.model_device=device
         self.out_dim = 1024
         return