Update scripts/transcribe.py

scripts/transcribe.py

@@ -0,0 +1,31 @@
+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 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())
+        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])
+        print("Transcription completed.")
+        return transcription