Update scripts/transcribe.py

scripts/transcribe.py

@@ -1,31 +1,16 @@
+import whisper
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-from pydub import AudioSegment
-import numpy as np
 
 class SpeechToText:
-    def __init__(self):
-        print("Loading model...")
-        self.processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
-        self.model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
-        print("Model loaded successfully.")
+    def __init__(self, model_size="base"):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Loading Whisper '{model_size}' model on {device}...")
+        self.model = whisper.load_model(model_size, device=device)
+        print("Whisper model loaded successfully.")
 
-    def convert_audio(self, audio_path):
-        print("Converting audio...")
-        audio = AudioSegment.from_file(audio_path)
-        audio = audio.set_channels(1).set_frame_rate(16000)
-        samples = np.array(audio.get_array_of_samples()).astype(np.float32)  # <-- fixed here
-        print("Audio conversion complete.")
-        return samples
-
-    def transcribe(self, audio_samples):
-        print("Starting transcription...")
-        inputs = self.processor(audio_samples, sampling_rate=16000, return_tensors="pt", padding=True)
-
-        with torch.no_grad():
-            logits = self.model(inputs.input_values).logits
-        
-        predicted_ids = torch.argmax(logits, dim=-1)
-        transcription = self.processor.decode(predicted_ids[0])
+    def transcribe(self, audio_path):
+        print("Starting transcription with Whisper...")
+        result = self.model.transcribe(audio_path)
+        transcript = result["text"]
         print("Transcription completed.")
-        return transcription
+        return transcript