large commit

openbb_terminal/core/config/paths.py

@@ -35,3 +35,4 @@ USER_PRESETS_DIRECTORY = USER_DATA_DIRECTORY / "presets"
 USER_REPORTS_DIRECTORY = USER_DATA_DIRECTORY / "reports"
 USER_CUSTOM_REPORTS_DIRECTORY = USER_DATA_DIRECTORY / "reports" / "custom reports"
 USER_FORECAST_MODELS_DIRECTORY = USER_DATA_DIRECTORY / "exports" / "forecast_models"
+USER_FORECAST_WHISPER_DIRECTORY = USER_DATA_DIRECTORY / "exports" / "whisper"

openbb_terminal/forecast/forecast_controller.py

@@ -26,11 +26,23 @@ except ModuleNotFoundError:
         "Please install the forecast version of the terminal. Instructions can be found "
         "under the python tab: https://docs.openbb.co/terminal/quickstart/installation"
     )
+
+try:
+    import whisper
+    from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE
+    from openbb_terminal.forecast.whisper_utils import str2bool
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "Please install whisper model. Instructions can be found here: "
+        "https://github.com/openai/whisper"
+    )
+
 import pandas as pd
 import psutil
 from openbb_terminal.core.config.paths import (
     USER_EXPORTS_DIRECTORY,
     USER_CUSTOM_IMPORTS_DIRECTORY,
+    USER_FORECAST_WHISPER_DIRECTORY,
 )
 from openbb_terminal import feature_flags as obbff
 from openbb_terminal.custom_prompt_toolkit import NestedCompleter
@@ -3132,3 +3144,98 @@ class ForecastController(BaseController):
                 naive=ns_parser.naive,
                 export_pred_raw=ns_parser.export_pred_raw,
             )
+
+    @log_start_end(log=logger)
+    def call_whisper(self, other_args: List[str]):
+        """Utilize Whisper Model to transcribe a video. Currently only supports Youtube URLS"""
+        parser = argparse.ArgumentParser(
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+            add_help=False,
+            prog="whisper",
+            description="""
+                Utilize Whisper Model to transcribe a video. Currently only supports Youtube URLS:
+                https://github.com/openai/whisper
+            """,
+        )
+        parser.add_argument(
+            "--video",
+            dest="video",
+            type=str,
+            default="",
+            help="video URLs to transcribe",
+        )
+        parser.add_argument(
+            "--model_name",
+            dest="model_name",
+            choices=whisper.available_models(),
+            default="small",
+            help="name of the Whisper model to use",
+        )
+        parser.add_argument(
+            "--subtitles_format",
+            dest="subtitles_format",
+            type=str,
+            choices=["vtt", "srt"],
+            help="the subtitle format to output",
+        )
+        parser.add_argument(
+            "--verbose",
+            dest="verbose",
+            type=str2bool,
+            default=False,
+            help="Whether to print out the progress and debug messages",
+        )
+        parser.add_argument(
+            "--task",
+            dest="task",
+            type=str,
+            choices=["transcribe", "translate"],
+            help="whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')",
+        )
+        parser.add_argument(
+            "--language",
+            dest="language",
+            type=str,
+            default=None,
+            choices=sorted(LANGUAGES.keys())
+            + sorted([k.title() for k in TO_LANGUAGE_CODE.keys()]),
+            help="language spoken in the audio, skip to perform language detection",
+        )
+        parser.add_argument(
+            "--breaklines",
+            dest="breaklines",
+            type=int,
+            default=0,
+            help="Whether to break lines into a bottom-heavy pyramid shape if line length exceeds N characters. 0 disables line breaking.",
+        )
+        parser.add_argument(
+            "--save",
+            dest="save",
+            type=str,
+            default=USER_FORECAST_WHISPER_DIRECTORY,
+            help="Whether to break lines into a bottom-heavy pyramid shape if line length exceeds N characters. 0 disables line breaking.",
+        )
+
+        parser = self.add_standard_args(
+            parser,
+        )
+        ns_parser = self.parse_known_args_and_warn(
+            parser,
+            other_args,
+            export_allowed=EXPORT_ONLY_FIGURES_ALLOWED,
+        )
+
+        if ns_parser:
+            if not helpers.check_parser_input(ns_parser, self.datasets):
+                return
+
+            whisper_view.transcribe_and_summarize(
+                video=ns_parser.video,
+                model_name=ns_parser.model_name,
+                subtitles_format=ns_parser.subtitles_format,
+                verbose=ns_parser.verbose,
+                task=ns_parser.task,
+                language=ns_parser.language,
+                breaklines=ns_parser.breaklines,
+                export=ns_parser.save,
+            )

openbb_terminal/forecast/whisper_utils.py

@@ -0,0 +1,87 @@
+from typing import Iterator, TextIO
+
+
+def str2bool(string):
+    str2val = {"True": True, "False": False}
+    if string in str2val:
+        return str2val[string]
+    else:
+        raise ValueError(f"Expected one of {set(str2val.keys())}, got {string}")
+
+
+def format_timestamp(
+    seconds: float, always_include_hours: bool = False, decimal_marker: str = "."
+):
+    assert seconds >= 0, "non-negative timestamp expected"
+    milliseconds = round(seconds * 1000.0)
+
+    hours = milliseconds // 3_600_000
+    milliseconds -= hours * 3_600_000
+
+    minutes = milliseconds // 60_000
+    milliseconds -= minutes * 60_000
+
+    seconds = milliseconds // 1_000
+    milliseconds -= seconds * 1_000
+
+    hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
+    return (
+        f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
+    )
+
+
+def break_line(line: str, length: int):
+    break_index = min(len(line) // 2, length)  # split evenly or at maximum length
+
+    # work backwards from that guess to split between words
+    # if break_index <= 1, we've hit the beginning of the string and can't split
+    while break_index > 1:
+        if line[break_index - 1] == " ":
+            break  # break at space
+        else:
+            break_index -= 1
+    if break_index > 1:
+        # split the line, not including the space at break_index
+        return line[: break_index - 1] + "\n" + line[break_index:]
+    else:
+        return line
+
+
+def process_segment(segment: dict, line_length: int = 0):
+    segment["text"] = segment["text"].strip()
+    if line_length > 0 and len(segment["text"]) > line_length:
+        # break at N characters as per Netflix guidelines
+        segment["text"] = break_line(segment["text"], line_length)
+
+    return segment
+
+
+def write_vtt(transcript: Iterator[dict], file: TextIO, line_length: int = 0):
+    print("WEBVTT\n", file=file)
+    for segment in transcript:
+        segment = process_segment(segment, line_length=line_length)
+
+        print(
+            f"{format_timestamp(segment['start'])} --> {format_timestamp(segment['end'])}\n"
+            f"{segment['text'].strip().replace('-->', '->')}\n",
+            file=file,
+            flush=True,
+        )
+
+
+def write_srt(transcript: Iterator[dict], file: TextIO, line_length: int = 0):
+    for i, segment in enumerate(transcript, start=1):
+        segment = process_segment(segment, line_length=line_length)
+
+        print(
+            f"{i}\n"
+            f"{format_timestamp(segment['start'], always_include_hours=True, decimal_marker=',')} --> "
+            f"{format_timestamp(segment['end'], always_include_hours=True, decimal_marker=',')}\n"
+            f"{segment['text'].strip().replace('-->', '->')}\n",
+            file=file,
+            flush=True,
+        )
+
+
+def slugify(title):
+    return "".join(c if c.isalnum() else "_" for c in title).rstrip("_")

openbb_terminal/forecast/whisper_view.py

@@ -12,9 +12,96 @@ from openbb_terminal.decorators import log_start_end
 from openbb_terminal.forecast import helpers
 from openbb_terminal.rich_config import console
 
+import warnings
+import os
+import whisper
+from openbb_terminal.forecast.whisper_utils import (
+    slugify,
+    str2bool,
+    write_srt,
+    write_vtt,
+)
+import yt_dlp
+
+# from .utils import slugify, str2bool, write_srt, write_vtt
+import tempfile
+
 logger = logging.getLogger(__name__)
 # pylint: disable=too-many-arguments
 
 
-# @log_start_end(log=logger)
-# def transcribe_and_summarize():
+def get_audio(urls):
+    temp_dir = tempfile.gettempdir()
+
+    ydl = yt_dlp.YoutubeDL(
+        {
+            "quiet": True,
+            "verbose": False,
+            "format": "bestaudio",
+            "outtmpl": os.path.join(temp_dir, "%(id)s.%(ext)s"),
+            "postprocessors": [
+                {
+                    "preferredcodec": "mp3",
+                    "preferredquality": "192",
+                    "key": "FFmpegExtractAudio",
+                }
+            ],
+        }
+    )
+
+    paths = {}
+
+    for url in urls:
+        result = ydl.extract_info(url, download=True)
+        print(f"Downloaded video \"{result['title']}\". Generating subtitles...")
+        paths[result["title"]] = os.path.join(temp_dir, f"{result['id']}.mp3")
+
+    return paths
+
+
+@log_start_end(log=logger)
+def transcribe_and_summarize(
+    video: str = "",
+    model_name: str = "small",
+    subtitles_format: str = "vtt",
+    export: str = "",
+    verbose: bool = False,
+    task: str = "transcribe",
+    language: str = None,
+    breaklines: int = 0,
+    output_dir: str = "/Users/martinbufi/OpenBBTerminal/openbb_terminal/forecast/whisper_output",
+):
+    console.print("Transcribing and summarizing...")
+    if model_name.endswith(".en"):
+        warnings.warn(
+            f"{model_name} is an English-only model, forcing English detection."
+        )
+    os.makedirs(output_dir, exist_ok=True)
+
+    model = whisper.load_model(model_name)
+    audios = get_audio(video)
+
+    for title, audio_path in audios.items():
+        warnings.filterwarnings("ignore")
+
+        result = whisper.transcribe(
+            audio_path,
+            model,
+            language=language,
+            verbose=verbose,
+            task=task,
+        )
+        warnings.filterwarnings("default")
+
+        if subtitles_format == "vtt":
+            vtt_path = os.path.join(output_dir, f"{slugify(title)}.vtt")
+            with open(vtt_path, "w", encoding="utf-8") as vtt:
+                write_vtt(result["segments"], file=vtt, line_length=breaklines)
+
+            print("Saved VTT to", os.path.abspath(vtt_path))
+        else:
+            srt_path = os.path.join(output_dir, f"{slugify(title)}.srt")
+            with open(srt_path, "w", encoding="utf-8") as srt:
+                write_srt(result["segments"], file=srt, line_length=breaklines)
+
+            print("Saved SRT to", os.path.abspath(srt_path))