| from transformers import PreTrainedModel | |
| from .configuration_test import TestConfig | |
| import torch.nn as nn | |
| from transformers import AutoModelForMaskedLM, AutoConfig | |
| import librosa | |
| from huggingface_hub import hf_hub_download | |
| import os | |
| class TestModel(PreTrainedModel): | |
| config_class = TestConfig | |
| def __init__(self, config: TestConfig): | |
| super().__init__(config) | |
| self.input_dim = config.input_dim | |
| self.model1 = nn.Linear(config.input_dim, config.output_dim) | |
| self.model2 = AutoModelForMaskedLM.from_config( | |
| AutoConfig.from_pretrained("albert/albert-base-v2") | |
| ) | |
| self.path = config.name_or_path | |
| def get_audio_duration(self): | |
| audio_path = hf_hub_download( | |
| repo_id = self.path, | |
| filename = "output1.wav", | |
| repo_type = "model", | |
| local_dir = os.path.dirname(os.path.abspath(__file__)) | |
| ) | |
| arr, sr = librosa.load(audio_path) | |
| return librosa.get_duration(y=arr, sr=sr) | |
| def forward(self, tensor): | |
| return self.model1(tensor) | |