Update transcription pipeline to use AutoModelForCTC, closes #238

neuml · Mar 3, 2022 · 6816a63 · 6816a63
1 parent 728e3dd
commit 6816a63
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/python/txtai/pipeline/audio/transcription.py b/src/python/txtai/pipeline/audio/transcription.py
@@ -9,7 +9,7 @@
 except (ImportError, OSError):
     SOUNDFILE = False
 
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import AutoModelForCTC, Wav2Vec2Processor
 
 from ..hfmodel import HFModel
 
@@ -38,7 +38,7 @@ def __init__(self, path="facebook/wav2vec2-base-960h", quantize=False, gpu=True,
             raise ImportError("SoundFile library not installed or libsndfile not found")
 
         # load model and processor
-        self.model = Wav2Vec2ForCTC.from_pretrained(self.path)
+        self.model = AutoModelForCTC.from_pretrained(self.path)
         self.processor = Wav2Vec2Processor.from_pretrained(self.path)
 
         # Move model to device