youtube_runner.py

import argparse
import json
import os
import pprint
import time
from os import PathLike
from pathlib import Path
from typing import Any

import sqlalchemy as db
from loguru import logger
from sqlalchemy.orm import sessionmaker
from youtube_transcript_api._errors import TranscriptsDisabled

from generators.youtube_generator import generate
from models import DummyTest
from src.database import YouTubeBase
from src.dataclasses import YouTubeVideo
from src.test_runner import TestRegistry, TestRunner
from src.utils import insert_youtube_result


@TestRegistry.register(DummyTest)
class YouTubeTestRunner(TestRunner):
    """
    Test runner for youtube videos, uses a testplan generated by youtube_generator
    """

    def __init__(
        self,
        testplan_path: PathLike,
        audio_dir: PathLike,
        output_dir: PathLike,
        iterations: int = 1,
        save_transcripts: bool = False,
        save_to_database: bool = False,
        keep_audio: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self._audio_dir = Path(audio_dir)
        self._output_dir = Path(output_dir)
        self._testplan_path = Path(testplan_path)
        self._iterations = iterations
        self._save_transcripts = save_transcripts
        self._save_to_database = save_to_database
        self._keep_audio = keep_audio
        self._session = None

        if self._save_to_database:
            engine = db.create_engine(f"sqlite:///youtube.sqlite")
            YouTubeBase.metadata.create_all(engine)
            self._session = sessionmaker(bind=engine)()

    def run(self) -> None:
        if self.tester.transcriber is None:
            raise ValueError("Transcriber is None")

        if self.tester.normalizer is None:
            logger.warning("Normalizer is None, running without normalizer")

        # check if we can generate more testplans if we need to
        if self._iterations > 1 and "GoogleAPI" not in os.environ:
            raise ValueError(
                "GoogleAPI not in the environment, can not generate more test plans. "
                "Add GoogleAPI or set iterations to 1."
            )

        with open(self._testplan_path, encoding="utf8") as f:
            testplan = json.load(f)

        # run the testplan
        for i in range(self._iterations):
            logger.info(f"Starting {i + 1}/{self._iterations} testplan")
            logger.info(f"Testplan args:\n{pprint.pformat(testplan['args'])}")

            # iterate over the videos in the testplan
            for idx, video_details in enumerate(testplan["items"]):
                logger.info(
                    f"Testplan status: {idx + 1}/{len(testplan['items'])} video, {i + 1}/{self._iterations} testplan"
                )
                video = YouTubeVideo.from_dict(video_details)

                # download the target transcript
                try:
                    target_transcript = video.youtube_transcript(self.tester.language)
                except ValueError as e:
                    logger.warning(
                        f"Skipping the video {video.videoId}, ValueError (youtube transcript): {e}"
                    )
                    video_details["error"] = f"ValueError (youtube transcript): {e}"
                    continue
                except TranscriptsDisabled as e:
                    logger.warning(
                        f"Skipping the video {video.videoId}, TranscriptsDisabled (youtube transcript): {e}"
                    )
                    video_details[
                        "error"
                    ] = f"TranscriptsDisabled (youtube transcript): {e}"
                    continue

                # download the audio
                try:
                    audio = video.download_mp3(self._audio_dir)
                except ValueError as e:
                    logger.warning(
                        f"Skipping the video {video.videoId}, ValueError (download): {e}"
                    )
                    video_details["error"] = f"ValueError (download): {e}"
                    continue

                # transcribe the audio
                try:
                    model_transcript = self.tester.transcribe(audio)
                except TimeoutError as e:
                    logger.warning(
                        f"Skipping the video {video.videoId}, TimeoutError (model transcript): {e}"
                    )
                    video_details["error"] = f"TimeoutError (model transcript): {e}"
                    continue

                if not self._keep_audio:
                    audio.unlink()

                # compare the transcripts
                try:
                    results = self.tester.compare(model_transcript, target_transcript)
                    results.update(self.tester.additional_info())
                    video_details["results"] = results
                except ValueError as e:
                    logger.warning(
                        f"Skipping the video {video.videoId}, ValueError (compare): {e}"
                    )
                    video_details["error"] = f"ValueError (compare): {e}"
                    continue

                # add the transcripts to the video details if we want to save them
                if self._save_transcripts:
                    video_details["modelTranscript"] = model_transcript
                    video_details["targetTranscript"] = target_transcript

            self.save_results(testplan)

            # generate a new testplan if we need to
            if i + 1 < self._iterations:
                args = testplan["args"]
                args["pageToken"] = testplan["nextPageToken"]
                testplan = generate(args)

        logger.info("Testplan finished")

    def save_results(self, results: dict[str, Any]) -> None:
        """
        Save the results to a json file

        Args:
            results: results to save
        """

        time_str = time.strftime("%Y%m%d-%H%M%S")
        filename = f"{results['args']['q']}_{self.__class__.__name__}_{time_str}.json"

        path = Path(__file__).parent.joinpath(self._output_dir, filename)
        path.parent.mkdir(exist_ok=True)

        logger.info(f"Saving results - {path}")

        with open(path, "x", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False)

        if self._save_to_database and not insert_youtube_result(
            self._session, filename, results
        ):
            logger.warning(f"Failed to save results to database")

    def __repr__(self):
        try:
            category = f'_{self._testplan_path.name.split("_")[0]}'
        except IndexError:
            category = ""
        return f"{self.__class__.__name__}{category}"

    @staticmethod
    def runner_args(parser: argparse.ArgumentParser) -> None:
        """
        Add runner specific arguments to the parser

        Args:
            parser: parser to add arguments to
        """

        parser.add_argument(type=str, dest="testplan_path", help="Testplan path")

        parser.add_argument(
            "-st",
            "--save-transcript",
            required=False,
            action="store_true",
            dest="save_transcripts",
            default=False,
        )

        parser.add_argument(
            "-db",
            "--save-to-database",
            required=False,
            action="store_true",
            dest="save_to_database",
            default=False,
        )

        parser.add_argument(
            "--audio-path",
            required=False,
            type=str,
            default="./cache/audio",
            dest="audio_dir",
        )

        parser.add_argument(
            "-k",
            "--keep-audio",
            required=False,
            action="store_true",
            dest="keep_audio",
            default=False,
        )

        parser.add_argument("-it", "--iterations", required=False, type=int, default=1)

        parser.add_argument(
            "-o",
            "--output",
            required=False,
            type=str,
            default="./output",
            dest="output_dir",
        )


if __name__ == "__main__":
    YouTubeTestRunner.from_command_line().run()