Transcription model override

2024-10-16 09:52:53 -04:00
parent e8bdc317f5
commit 659cf3728a
2 changed files with 6 additions and 4 deletions
--- a/transcribe.py
+++ b/transcribe.py
@@ -1,14 +1,14 @@
 import torch
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
-def process(audio_path):
+def process(audio_path, model_name):
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
-        "openai/whisper-large-v3",
+        model_name,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
        use_safetensors=True
    )
    model.to("mps")
-    processor = AutoProcessor.from_pretrained("openai/whisper-large-v3")
+    processor = AutoProcessor.from_pretrained(model_name)
    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,