Spaces:
Runtime error
Runtime error
| #speech to text tool | |
| import os | |
| import subprocess | |
| import whisperx | |
| from pydub import AudioSegment | |
| from pytube import YouTube | |
| class SpeechToText: | |
| def __init__( | |
| self, | |
| video_url, | |
| audio_format='mp3', | |
| device='cuda', | |
| batch_size = 16, | |
| compute_type = "float16", | |
| hf_api_key = None | |
| ): | |
| """ | |
| # Example usage | |
| video_url = "url" | |
| speech_to_text = SpeechToText(video_url) | |
| transcription = speech_to_text.transcribe_youtube_video() | |
| print(transcription) | |
| """ | |
| self.video_url = video_url | |
| self.audio_format = audio_format | |
| self.device = device | |
| self.batch_size = batch_size | |
| self.compute_type = compute_type | |
| self.hf_api_key = hf_api_key | |
| def install(self): | |
| subprocess.run(["pip", "install", "whisperx"]) | |
| subprocess.run(["pip", "install", "pytube"]) | |
| subprocess.run(["pip", "install", "pydub"]) | |
| def download_youtube_video(self): | |
| audio_file = f'video.{self.audio_format}' | |
| # Download video 📥 | |
| yt = YouTube(self.video_url) | |
| yt_stream = yt.streams.filter(only_audio=True).first() | |
| yt_stream.download(filename='video.mp4') | |
| # Convert video to audio 🎧 | |
| video = AudioSegment.from_file("video.mp4", format="mp4") | |
| video.export(audio_file, format=self.audio_format) | |
| os.remove("video.mp4") | |
| return audio_file | |
| def transcribe_youtube_video(self): | |
| audio_file = self.download_youtube_video() | |
| device = "cuda" | |
| batch_size = 16 | |
| compute_type = "float16" | |
| # 1. Transcribe with original Whisper (batched) 🗣️ | |
| model = whisperx.load_model("large-v2", device, compute_type=compute_type) | |
| audio = whisperx.load_audio(audio_file) | |
| result = model.transcribe(audio, batch_size=batch_size) | |
| # 2. Align Whisper output 🔍 | |
| model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device) | |
| result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False) | |
| # 3. Assign speaker labels 🏷️ | |
| diarize_model = whisperx.DiarizationPipeline( | |
| use_auth_token=self.hf_api_key, | |
| device=device | |
| ) | |
| diarize_model(audio_file) | |
| try: | |
| segments = result["segments"] | |
| transcription = " ".join(segment['text'] for segment in segments) | |
| return transcription | |
| except KeyError: | |
| print("The key 'segments' is not found in the result.") | |
| def transcribe(self, audio_file): | |
| model = whisperx.load_model( | |
| "large-v2", | |
| self.device, | |
| self.compute_type | |
| ) | |
| audio = whisperx.load_audio(audio_file) | |
| result = model.transcribe( | |
| audio, | |
| batch_size=self.batch_size | |
| ) | |
| # 2. Align Whisper output 🔍 | |
| model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device) | |
| result = whisperx.align( | |
| result["segments"], | |
| model_a, | |
| metadata, | |
| audio, | |
| self.device, | |
| return_char_alignments=False | |
| ) | |
| # 3. Assign speaker labels 🏷️ | |
| diarize_model = whisperx.DiarizationPipeline( | |
| use_auth_token=self.hf_api_key, | |
| device=self.device | |
| ) | |
| diarize_model(audio_file) | |
| try: | |
| segments = result["segments"] | |
| transcription = " ".join(segment['text'] for segment in segments) | |
| return transcription | |
| except KeyError: | |
| print("The key 'segments' is not found in the result.") | |