Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -14,13 +14,21 @@ 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="Capstone04/TrainedWhisper",
+                 model_path = "None"):
         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)
 
+        if model_path and os.path.exists(model_path):
+            print(f"🔄 Loading custom ASR model from: {model_path}")
+            actual_asr_model = model_path
+        else:
+            print(f"🔄 Loading default ASR model: {asr_model}")
+            actual_asr_model = asr_model
+
         processor = WhisperProcessor.from_pretrained(asr_model, token=HF_TOKEN)
         model = WhisperForConditionalGeneration.from_pretrained(asr_model, token=HF_TOKEN).to(self.device)