Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -21,16 +21,24 @@ class ASR_Diarization:
         # Load diarization model
         self.diar_pipeline = Pipeline.from_pretrained(diar_model, use_auth_token=HF_TOKEN)
 
-        processor = WhisperProcessor.from_pretrained(asr_model, token=HF_TOKEN)
-        model = WhisperForConditionalGeneration.from_pretrained(asr_model, token=HF_TOKEN).to(self.device)
-
+        # processor = WhisperProcessor.from_pretrained(asr_model, token=HF_TOKEN)
+        # model = WhisperForConditionalGeneration.from_pretrained(asr_model, token=HF_TOKEN).to(self.device)
+
+        # self.asr_pipeline = hf_pipeline(
+        #     "automatic-speech-recognition",
+        #     model=model,
+        #     tokenizer=processor.tokenizer,
+        #     feature_extractor=processor.feature_extractor,
+        #     device=0 if self.device == "cuda" else -1,
+        #     return_timestamps=True
+        # )
+
+        model_id = "Capstone04/TrainedWhisper"
         self.asr_pipeline = hf_pipeline(
             "automatic-speech-recognition",
-            model=model,
-            tokenizer=processor.tokenizer,
-            feature_extractor=processor.feature_extractor,
-            device=0 if self.device == "cuda" else -1,
-            return_timestamps=True
+            model=model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device=0 if torch.cuda.is_available() else -1,
         )
 
     def run_diarization(self, audio_path):