Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -14,33 +14,33 @@ from jiwer import wer, Compose, ToLowerCase, RemovePunctuation, RemoveMultipleSp
 class ASR_Diarization:
     def __init__(self, HF_TOKEN,
                  diar_model="pyannote/speaker-diarization-3.1",
-                 asr_model="Capstone04/TrainedWhisper"):
+                 asr_model="openai/whisper-large-v3"):
         self.HF_TOKEN = HF_TOKEN
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         # 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_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device=0 if torch.cuda.is_available() else -1,
+            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_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):
         diarization = self.diar_pipeline(audio_path)
         return [