From 737cb895ac8b0559da2bf22d73fc346c8eebe31d Mon Sep 17 00:00:00 2001 From: martinb-bb <105685594+martinb-bb@users.noreply.github.com> Date: Mon, 16 Jan 2023 11:59:02 -0500 Subject: [PATCH] large commit --- openbb_terminal/core/config/paths.py | 1 + .../forecast/forecast_controller.py | 107 ++++++++++++++++++ openbb_terminal/forecast/whisper_utils.py | 87 ++++++++++++++ openbb_terminal/forecast/whisper_view.py | 91 ++++++++++++++- 4 files changed, 284 insertions(+), 2 deletions(-) create mode 100644 openbb_terminal/forecast/whisper_utils.py diff --git a/openbb_terminal/core/config/paths.py b/openbb_terminal/core/config/paths.py index 3aea1930ef3..058d9ad36bc 100644 --- a/openbb_terminal/core/config/paths.py +++ b/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" diff --git a/openbb_terminal/forecast/forecast_controller.py b/openbb_terminal/forecast/forecast_controller.py index 368c880f78e..526385c0d68 100644 --- a/openbb_terminal/forecast/forecast_controller.py +++ b/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, + ) diff --git a/openbb_terminal/forecast/whisper_utils.py b/openbb_terminal/forecast/whisper_utils.py new file mode 100644 index 00000000000..b3ea7bfae27 --- /dev/null +++ b/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("_") diff --git a/openbb_terminal/forecast/whisper_view.py b/openbb_terminal/forecast/whisper_view.py index 97d8d02c233..95cb74dff00 100644 --- a/openbb_terminal/forecast/whisper_view.py +++ b/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))