Update transcription_diarization.py

transcription_diarization.py

@@ -4,12 +4,15 @@ import gc
 import math
 from moviepy.editor import VideoFileClip
 from pyannote.audio import Pipeline
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 import librosa
 import soundfile as sf
 import datetime
 from collections import defaultdict
 import numpy as np
+import openai
+from config import openai_api_key
+
+openai.api_key = openai_api_key
 
 class LazyDiarizationPipeline:
     def __init__(self):
@@ -25,81 +28,37 @@ class LazyDiarizationPipeline:
             gc.collect()
         return self.pipeline
 
-
-class LazyTranscriptionPipeline:
-    def __init__(self):
-        self.model = None
-        self.processor = None
-        self.pipe = None
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    def get_pipeline(self):
-        if self.pipe is None:
-            model_id = "openai/whisper-large-v3"
-            torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-            self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
-                model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
-            )
-            self.model.to(self.device)
-            self.processor = AutoProcessor.from_pretrained(model_id)
-            self.pipe = pipeline(
-                "automatic-speech-recognition",
-                model=self.model,
-                tokenizer=self.processor.tokenizer,
-                feature_extractor=self.processor.feature_extractor,
-                chunk_length_s=30,
-                return_timestamps=True,
-                device=self.device
-            )
-        return self.pipe
-
-
 lazy_diarization_pipeline = LazyDiarizationPipeline()
-lazy_transcription_pipeline = LazyTranscriptionPipeline()
 
 def extract_audio(video_path, audio_path):
     video = VideoFileClip(video_path)
     audio = video.audio
     audio.write_audiofile(audio_path, codec='pcm_s16le', fps=16000)
 
-
 def format_timestamp(seconds):
     return str(datetime.timedelta(seconds=seconds)).split('.')[0]
 
-
 def transcribe_audio(audio_path, language):
-    pipe = lazy_transcription_pipeline.get_pipeline()
-
-    audio, sr = librosa.load(audio_path, sr=16000)
-    duration = len(audio) / sr
-    n_chunks = math.ceil(duration / 30)
-    transcription_txt = ""
+    with open(audio_path, "rb") as audio_file:
+        transcript = openai.Audio.transcribe(
+            file=audio_file,
+            model="whisper-1",
+            language=language,
+            response_format="verbose_json"
+        )
+
+    transcription_txt = transcript["text"]
     transcription_chunks = []
 
-    for i in range(n_chunks):
-        start = i * 30 * sr
-        end = min((i + 1) * 30 * sr, len(audio))
-        audio_chunk = audio[start:end]
-        audio_chunk = (audio_chunk * 32767).astype(np.float32)
-
-        result = pipe(audio_chunk, generate_kwargs={"language": language, "task": "transcribe"})
-
-        transcription_txt += result["text"]
-        for chunk in result["chunks"]:
-            start_time, end_time = chunk["timestamp"]
-            if start_time is None:
-                start_time = 0
-            if end_time is None:
-                end_time = 0
-            transcription_chunks.append({
-                "start": start_time + i * 30,
-                "end": end_time + i * 30,
-                "text": chunk["text"]
-            })
+    for segment in transcript["segments"]:
+        transcription_chunks.append({
+            "start": segment["start"],
+            "end": segment["end"],
+            "text": segment["text"]
+        })
 
     return transcription_txt, transcription_chunks
 
-
 def diarize_audio(audio_path, pipeline, max_speakers):
     # Load the entire audio file
     audio, sr = librosa.load(audio_path, sr=16000)
@@ -118,7 +77,6 @@ def diarize_audio(audio_path, pipeline, max_speakers):
 
     return diarization
 
-
 def create_combined_srt(transcription_chunks, diarization, output_path, max_speakers):
     speaker_segments = []
     speaker_durations = defaultdict(float)
@@ -155,7 +113,6 @@ def create_combined_srt(transcription_chunks, diarization, output_path, max_spea
             duration_str = format_timestamp(duration).split('.')[0].lstrip('0')
             srt_file.write(f"Speaker {i} (originally {speaker}): total duration {duration_str}\n")
 
-
 def process_video(video_path, hf_token, language, max_speakers=3):
     base_name = os.path.splitext(video_path)[0]
     audio_path = f"{base_name}.wav"
@@ -183,4 +140,4 @@ def process_video(video_path, hf_token, language, max_speakers=3):
     torch.cuda.empty_cache()
     gc.collect()
 
-    return combined_srt_path
+    return combined_srt_path