Upload folder using huggingface_hub

model_onnx.py

@@ -29,8 +29,8 @@ class IndicASRModel(PreTrainedModel):
         self.models = {}
         names = ['encoder', 'ctc_decoder', 'rnnt_decoder', 'joint_enc', 'joint_pred', 'joint_pre_net'] + [f'joint_post_net_{z}' for z in ['as', 'bn', 'brx', 'doi', 'gu', 'hi', 'kn', 'kok', 'ks', 'mai', 'ml', 'mni', 'mr', 'ne', 'or', 'pa', 'sa', 'sat', 'sd', 'ta', 'te', 'ur']]
         self.models = {}
-        self.models['preprocessor'] = torch.jit.load(f'{config.ts_folder}/assets/preprocessor.ts', map_location=self.config.device)
-
+        self.d = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.models['preprocessor'] = torch.jit.load(f'{config.ts_folder}/assets/preprocessor.ts', map_location=self.d)
         for n in names:
             component_name = f'{config.ts_folder}/assets/{n}.onnx'
             if os.path.exists(config.ts_folder):
@@ -55,7 +55,7 @@ class IndicASRModel(PreTrainedModel):
 
     def encode(self, wav):
         # pass through preprocessor
-        audio_signal, length = self.models['preprocessor'](input_signal=wav.to(self.config.device), length=torch.tensor([wav.shape[-1]]).to(self.config.device))
+        audio_signal, length = self.models['preprocessor'](input_signal=wav.to(self.d), length=torch.tensor([wav.shape[-1]]).to(self.d))
         outputs, encoded_lengths = self.models['encoder'].run(['outputs', 'encoded_lengths'], {'audio_signal': audio_signal.cpu().numpy(), 'length': length.cpu().numpy()})
         return outputs, encoded_lengths