Develop (#3)

* first commit

* black code

* settings added

* readme completed

* workflow test

* workflow test

* workflow test

* workflow test

* codecov status

* readme updated

* fix pytube bug

* fix imports

* demo video added

* demo video added

* demo video added

* demo video removed

* fix config error issue #2

Author: rafleze

plugin.json

@@ -1,6 +1,6 @@
 {
-    "name": "Youtube 2 text",
-    "version": "0.0.1",
+    "name": "Youtube 2 Text",
+    "version": "0.0.2",
     "description": "A plugin to transcript a youtube video to text. Just past the video link and wait for the response!",
     "author_name": "rafleze",
     "author_url": "https://github.com/rafleze",

requirements.txt

@@ -1,3 +1,3 @@
 faster-whisper==1.0.1
-pytube==15.0.0
+yt-dlp==2025.2.19
 ffmpeg-python==0.2.0

tests/test_transcriber.py

@@ -4,17 +4,15 @@
 from ..transcriber import transcribe
 
 
-@patch("youtube2text.transcriber.YouTube")
+@patch("youtube2text.transcriber.download_audio")
 @patch("youtube2text.transcriber.WhisperModel")
-def test_transcribe_with_settings(mock_whisper_model, mock_youtube):
+def test_transcribe_with_settings(mock_whisper_model, mock_download_audio):
     """Test transcribe function with settings."""
 
     with open("tests/media/youtube/test.mp4", "w") as f:
         f.write("test")
 
-    mock_youtube.return_value.streams.filter.return_value.first.return_value.download.return_value = (
-        "tests/media/youtube/test.mp4"
-    )
+    mock_download_audio.return_value = "tests/media/youtube/test.mp4"
     mock_whisper_model.return_value.transcribe.return_value = [
         {"text": "Hello", "start": 0.0, "end": 1.0},
         {"text": "World", "start": 1.0, "end": 2.0},
@@ -36,16 +34,14 @@ def test_transcribe_with_settings(mock_whisper_model, mock_youtube):
     )
 
 
-@patch("youtube2text.transcriber.YouTube")
-def test_transcribe_without_settings(mock_youtube):
+@patch("youtube2text.transcriber.download_audio")
+def test_transcribe_without_settings(mock_download_audio):
     """Test transcribe function with settings."""
 
     with open("tests/media/youtube/test.mp4", "w") as f:
         f.write("test")
 
-    mock_youtube.return_value.streams.filter.return_value.first.return_value.download.return_value = (
-        "tests/media/youtube/test.mp4"
-    )
+    mock_download_audio.return_value = "tests/media/youtube/test.mp4"
     settings = {}
     with pytest.raises(Exception):
         transcribe(
@@ -55,14 +51,12 @@ def test_transcribe_without_settings(mock_youtube):
         )
 
 
-@patch("youtube2text.transcriber.YouTube")
-def test_transcribe_with_download_exception(mock_youtube):
+@patch("youtube2text.transcriber.download_audio")
+def test_transcribe_with_download_exception(mock_download_audio):
     """Test transcribe function with settings."""
 
     # mock youtube stream filter first download as exception
-    mock_youtube.return_value.streams.filter.return_value.first.return_value.download.side_effect = Exception(
-        "Download failed"
-    )
+    mock_download_audio.side_effect = Exception("Download failed")
     settings = {
         "model_size_or_path": "large-v3",
         "device": "cpu",

transcriber.py

@@ -1,37 +1,48 @@
 """Transcribe a youtube video."""
-
-from pytube import YouTube
 from faster_whisper import WhisperModel
 import tempfile
 import os
-from .custom_cipher import Cipher as CustomCipher
-from unittest.mock import patch
+import yt_dlp
 
 
-@patch(
-    "pytube.extract.Cipher", CustomCipher
-)  # Patch the Cipher class because it's not working: look at the issue https://github.com/pytube/pytube/issues/1918
 def transcribe(link, language="en", settings={}):
     """Transcribe a youtube video."""
     if not settings:
         raise Exception("No configuration found for Youtube2Text.")
     model_size_or_path = settings.get("model_size_or_path", "large-v3")
     device = settings.get("device", "cpu")
     compute_type = settings.get("compute_type", "int8")
+    filename = download_audio(link)
+    audio = open(filename, "rb")
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        temp_file.write(audio.read())
+        temp_file_path = temp_file.name
+        model = WhisperModel(
+            model_size_or_path, device=device, compute_type=compute_type
+        )
+        result = model.transcribe(temp_file_path, language=language)
+        os.remove(temp_file_path)
+        os.remove(filename)
+        return result
+
+
+def download_audio(link, output_folder="media/youtube"):  # pragma: no cover
     try:
-        yt = YouTube(link)
-        filename = yt.streams.filter(only_audio=True).first().download("media/youtube")
+
+        inner_output_folder = os.path.join(os.path.dirname(__file__), output_folder)
+        os.makedirs(inner_output_folder, exist_ok=True)
+
+        ydl_opts = {
+            "format": "bestaudio/best",
+            "outtmpl": f"{inner_output_folder}/%(title)s.%(ext)s",
+            "noprogress": False,
+        }
+
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info_dict = ydl.extract_info(link, download=True)
+            filename = ydl.prepare_filename(info_dict)
+
+        return filename
+
     except Exception as e:
-        raise e
-    else:
-        audio = open(filename, "rb")
-        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-            temp_file.write(audio.read())
-            temp_file_path = temp_file.name
-            model = WhisperModel(
-                model_size_or_path, device=device, compute_type=compute_type
-            )
-            result = model.transcribe(temp_file_path, language=language)
-            os.remove(temp_file_path)
-            os.remove(filename)
-            return result
+        raise e

youtube_to_text.py

@@ -56,14 +56,11 @@ class TranscriptionForm(CatForm):
     def submit(self, form_data):
         """Submit the form."""
         settings = self.cat.mad_hatter.plugins["youtube2text"].load_settings()
-        try:
-            segments, _ = transcribe(
-                form_data["youtube_link"], form_data["language"], settings
-            )
-            result = "".join([s.text for s in segments])
-            prompt = f"Summerize the following text: {result}"
-            summary = self.cat.llm(prompt)
-            output = f"The transcription is: \n{result}\n\nSummary: {summary}"
-            return {"output": output}
-        except Exception as e:
-            return {"output": f"An error occurred: {e}"}
+        segments, _ = transcribe(
+            form_data["youtube_link"], form_data["language"], settings
+        )
+        result = "".join([s.text for s in segments])
+        prompt = f"Summerize the following text: {result}"
+        summary = self.cat.llm(prompt)
+        output = f"The transcription is: \n{result}\n\nSummary: {summary}"
+        return {"output": output}