Upload with huggingface_hub

dataset/hf_dataset.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+from datasets import load_dataset, Dataset
+from datasets.data_files import EmptyDatasetError
+
+class HFDataset(ABC):
+  """
+  Create a dataset to save the transcripts from Youtube.
+  """
+  def __init__(self, name) -> None:
+    self.name = name
+    if name != "":
+      self._init_dataset()
+    else:
+      self.dataset = Dataset.from_dict({})
+      self.exist = False
+      self.is_empty = True
+
+  @abstractmethod
+  def generate_dataset():
+    pass
+
+  def _init_dataset(self):
+    try:
+      self.dataset = load_dataset(self.name)
+      self.exist = True
+      self.is_empty = False
+    except EmptyDatasetError:
+      self.dataset = Dataset.from_dict({})
+      self.exist = True
+      self.is_empty = True
+      pass
+    except FileNotFoundError:
+      self.dataset = Dataset.from_dict({})
+      self.exist = False
+      self.is_empty = True
+      pass
+  
+  def upload(self):
+    self.dataset.push_to_hub(self.name)

dataset/transcript_dataset.py

@@ -0,0 +1,63 @@
+import glob
+import os
+import validators
+import pandas as pd
+from downloader import WhisperPP, YoutubeDownloader
+from interpreter import WhisperInterpreter
+from datasets import load_dataset, concatenate_datasets, Dataset
+from dataset.hf_dataset import HFDataset
+
+class TranscriptDataset(HFDataset):
+
+  def __init__(self, name) -> None:
+    super().__init__(name)
+
+  def generate_dataset(self, input, download_path, overwrite, whisper_config):
+    if validators.url(input):
+      self.from_url(input, download_path, overwrite, **whisper_config)
+    else:
+      self.from_files(input, overwrite,  **whisper_config)
+
+  def from_url(self, url: str, download_path: str = "tmp/", overwrite: bool = False, **whisper_config: dict) -> None:
+    if self.is_empty:
+      emptyDataset = self.dataset
+    else:
+      #emptyDataset=self.dataset["train"].filter(lambda e: e["id"] is None)
+      emptyDataset=self.dataset["train"]
+    whisper_config["number_videos"] = 5
+    whisperPP = WhisperPP(emptyDataset, **whisper_config)
+    downloader = YoutubeDownloader(download_path)
+    if not overwrite:
+      downloader.config["download_archive"] = os.path.join(download_path,"video_record.txt")
+      self._fill_archive(downloader.config["download_archive"])
+    downloader.download(url, whisperPP)
+    self._concatenate_datasets(whisperPP.get_data())
+
+  def from_files(self, input:str, overwrite: bool = False, **whisper_config):
+    if (whisper_config.get("mode", None) is not None):
+      interpreter = WhisperInterpreter(whisper_config.pop("model_size"))
+      process = getattr(interpreter, whisper_config.pop("mode"))
+      result = process(input, **whisper_config)
+      if type(result) == list:
+        dataset = Dataset.from_list(result)
+      else:
+        dataset = Dataset.from_list([result])
+    else:
+      fileName = "tmp/*.json" if os.path.isdir(input) else input
+      dataset=load_dataset("json", data_files=glob.glob(fileName), split="train")
+    
+    self._concatenate_datasets(dataset)
+
+  def _fill_archive(self, archive_file):
+    if not self.is_empty:
+      with open(archive_file, "w") as f:
+        for id in self.dataset["train"]["id"]:
+          f.write(f"youtube {id}\n")
+
+  def _concatenate_datasets(self, dataset):
+    if not self.is_empty:
+      selectedIDs = list(set(dataset["id"])-set(self.dataset["train"]["id"]))
+      filteredDataset = dataset.filter(lambda element: element["id"] in selectedIDs)
+      self.dataset["train"] = concatenate_datasets([self.dataset["train"],filteredDataset])
+    else:
+      self.dataset = dataset

downloader/downloader.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+ 
+class Downloader(ABC):
+  """
+  A video downloader from online platforms to a specified format
+  """
+
+  @abstractmethod
+  def __init__(self, download_path):
+    self.download_path = download_path
+
+  @abstractmethod
+  def download(self):
+    pass

downloader/whisper_post_processor.py

@@ -0,0 +1,46 @@
+from interpreter import WhisperInterpreter
+from utils import VIDEO_INFO, json_dump
+from yt_dlp.postprocessor import PostProcessor
+from datasets import Dataset
+import re
+
+class WhisperPP(PostProcessor):
+  def __init__(self,data,**whisper_options):
+    super().__init__()
+    self._options = whisper_options
+    interpreter = WhisperInterpreter(self._options.pop("model_size","base"))
+    self.data = data
+    self._process = getattr(interpreter, self._options.pop("mode","transcribe"))
+    self._write = self._options.pop("write")
+    self.videos_to_process = self._options.pop("number_videos",0)
+    self.repoId = self._get_name()
+  
+  def run(self, info):
+    self.to_screen(f"Processing Video {info['id']}")
+    result = {key: info[key] for key in VIDEO_INFO}
+    result.update(self._process(info["filepath"], **self._options))
+    self.to_screen(f"Processed Video {info['id']} and appended results.")
+    self._update_data(result)
+    if self._write:
+      json_dump(result, f"{info['filepath'].split('.')[0]}.json")
+    return [], info
+
+  def _update_data(self, record):
+    dataType = type(self.data)
+    if dataType == list:
+      self.data.append(record)
+    else:
+      self.data = self.data.add_item(record)
+      if self.data.num_rows >= self.videos_to_process and self.videos_to_process != 0:
+        self.data.push_to_hub(self.repoId)
+
+  def get_data(self):
+    return self.data
+
+  def _get_name(self):
+    if self.data.info.download_checksums is not None:
+      regex = r"(?<=datasets\/)(.*?)(?=\/resolve)"
+      repoId = re.compile(regex)
+      url = list(self.data.info.download_checksums.keys())[0]
+      return repoId.findall(url)[0]
+    return ""

downloader/youtube_downloader.py

@@ -0,0 +1,26 @@
+import os
+import yt_dlp
+from downloader import Downloader
+from yt_dlp.postprocessor import PostProcessor
+from utils import YT_OPTIONS
+
+class YoutubeDownloader(Downloader):
+  
+  def __init__(self, download_path:str) -> None:
+    super().__init__(download_path)
+    self._ydl_options = YT_OPTIONS
+    self._ydl_options["outtmpl"] = os.path.join(download_path,"%(id)s.%(ext)s")
+
+
+  def download(self, url: str, CustomPP: PostProcessor, when: str = "post_process") -> None:
+    with yt_dlp.YoutubeDL(self._ydl_options) as ydl:
+      ydl.add_post_processor(CustomPP, when=when)
+      ydl.download(url)
+  
+  @property
+  def config(self):
+    return self._ydl_options
+
+  @config.setter
+  def config(self, key: str, value: str) -> None:
+    self._ydl_options[key] = value

interpreter/interpreter.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+ 
+class Interpreter(ABC):
+  """
+  An interpreter make some audio operations to transcribe or translate 
+  the video content.
+  """
+
+  @abstractmethod
+  def transcribe(self):
+    pass
+  
+  @abstractmethod
+  def translate(self):
+    pass

interpreter/whisper_interpreter.py

@@ -0,0 +1,48 @@
+import glob
+from typing import Any, Optional
+import whisper, os
+from interpreter import Interpreter
+from utils import SEGMENTS_INFO, AUDIO_FILES, json_dump
+
+class WhisperInterpreter(Interpreter):
+  
+  def __init__(self, model_size: str) -> None:
+    self.model = whisper.load_model(model_size)
+  
+  def transcribe(self, file_path: str, **kwargs: Optional[Any]) -> dict:
+    return self._execute_task("transcribe", file_path, **kwargs)
+
+  def translate(self, file_path: str, **kwargs: Optional[Any]) -> dict:
+    return self._execute_task("translate", file_path, **kwargs)
+
+  def _execute_task(self, mode: str, file_path: str, **kwargs: Optional[Any]) -> dict:
+    options = dict(task=mode)
+    options.update(kwargs)
+
+    if os.path.isdir(file_path):
+      result = []
+      files = [x for x in glob.glob(os.path.join(file_path,"*")) if os.path.splitext(x)[1] in AUDIO_FILES]
+      for file in files:
+        file_processed = dict(filename=file)
+        file_processed.update(self._file_extraction(file, **options))
+        result.append(file_processed)
+    else:
+      result = self._file_extraction(file_path, **options)
+
+    return result
+
+  def _formatter_result(self, input: dict) -> dict:
+    output = dict()
+    output["text"] = input["text"]
+    output["segments"] = [{key: segment[key] for key in SEGMENTS_INFO} for segment in input["segments"]]
+    return output
+
+  def _file_extraction(self, file_path: str, **kwargs: Optional[Any]) -> dict:
+    write = kwargs.pop("write",False)
+    result = self._formatter_result(
+            self.model.transcribe(file_path, **kwargs)
+          )
+    if write:
+      json_dump(result, f"{file_path.split('.')[0]}.json")
+    
+    return result