Implement yt_dlp instead of pytube

This commit is contained in:
Damien A
2024-07-14 11:25:02 +02:00
parent 65e6483c9f
commit 4eb5f586d4
3 changed files with 16 additions and 11 deletions

1
.gitignore vendored
View File

@@ -3,3 +3,4 @@
# Local data # Local data
.DS_Store .DS_Store
tmp

View File

@@ -1,5 +1,4 @@
openai-whisper==20231117 openai-whisper==20231117
pytube==15.0.0
ollama==0.1.8 ollama==0.1.8
torch==2.5.0.dev20240712 torch==2.5.0.dev20240712
torchaudio==2.4.0.dev20240712 torchaudio==2.4.0.dev20240712

View File

@@ -1,24 +1,29 @@
import ollama import ollama
import argparse import argparse
import os
from pytube import YouTube
from pathlib import Path from pathlib import Path
from transformers import pipeline from transformers import pipeline
import yt_dlp
OLLAMA_MODEL = "llama3" OLLAMA_MODEL = "llama3"
# Function to download a video from YouTube # Function to download a video from YouTube using yt-dlp
def download_from_youtube(url: str, path: str): def download_from_youtube(url: str, path: str):
yt = YouTube(url) ydl_opts = {
# Filter streams to get the highest resolution progressive mp4 stream 'format': 'bestaudio/best',
stream = yt.streams.filter(file_extension="mp4", only_audio=True).first() 'outtmpl': str(Path(path) / 'to_transcribe.%(ext)s'),
# Download the video to the specified path 'postprocessors': [{
stream.download(Path(path), filename="to_transcribe.mp4") 'key': 'FFmpegExtractAudio',
'preferredcodec': 'mp3',
'preferredquality': '192',
}],
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
ydl.download([url])
# Function to transcribe an audio file using the transformers pipeline # Function to transcribe an audio file using the transformers pipeline
def transcribe_file(file_path: str, output_file: str) -> str: def transcribe_file(file_path: str, output_file: str) -> str:
# Load the pipeline model for automatic speech recognition with MPS # Load the pipeline model for automatic speech recognition with MPS
transcriber_gpu = pipeline("automatic-speech-recognition", model="openai/whisper-large-v3", device="mps") transcriber_gpu = pipeline("automatic-speech-recognition", model="openai/whisper-large-v2", device="mps")
# Transcribe the audio file # Transcribe the audio file
transcribe = transcriber_gpu(file_path) transcribe = transcriber_gpu(file_path)
@@ -82,7 +87,7 @@ def main():
# Download from YouTube # Download from YouTube
print(f"Downloading YouTube video from {args.from_youtube}") print(f"Downloading YouTube video from {args.from_youtube}")
download_from_youtube(args.from_youtube, str(data_directory)) download_from_youtube(args.from_youtube, str(data_directory))
file_path = data_directory / "to_transcribe.mp4" file_path = data_directory / "to_transcribe.mp3"
elif args.from_local: elif args.from_local:
# Use local file # Use local file
file_path = Path(args.from_local) file_path = Path(args.from_local)