add youtube task creation

Former-commit-id: da4a1419b1c085934eba891dbf1079c2da05888b

.gitignore

@@ -10,4 +10,7 @@ test.py
 test.srt
 test.txt
 log_*.csv
-log.csv
+log.csv
+.chroma
+*.ini
+local_dump/

configs/local_launch.yaml

@@ -0,0 +1,5 @@
+# launch config for local environment
+model: "gpt-4"
+local_dump: ./local_dump
+output_type: srt
+environ: local

entries/__init_lib_path.py

@@ -0,0 +1,10 @@
+import sys
+import os
+
+def add_path(custom_path):
+    if custom_path not in sys.path: sys.path.insert(0, custom_path)
+
+this_dir = os.path.dirname(__file__)
+
+lib_path = os.path.join(this_dir, '..')
+add_path(lib_path)

entries/run.py

@@ -0,0 +1,80 @@
+import __init_lib_path
+import logging
+from yaml import Loader, Dumper, load, dump
+from src.task import Task
+import openai
+import argparse
+import os
+from pathlib import Path
+from datetime import datetime
+import shutil
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
+    parser.add_argument("--video_file", help="local video path here", default=None, type=str, required=False)
+    parser.add_argument("--audio_file", help="local audio path here", default=None, type=str, required=False)
+    parser.add_argument("--srt_file", help="srt file input path here", default=None, type=str, required=False)
+    parser.add_argument("--continue", help="task_id that need to continue", default=None, type=str, required=False) # need implement
+    parser.add_argument("--launch_cfg", help="launch config path", default='./configs/local_launch.yaml', type=str, required=False)
+    args = parser.parse_args()
+
+    return args
+
+if __name__ == "__main__":
+    args = parse_args()
+    launch_cfg = load(open(args.launch_cfg), Loader=Loader)
+
+    # initialize dir
+    local_dir = Path(launch_cfg['local_dump'])
+
+    # initialize task queue
+    if not local_dir.exists():
+        local_dir.mkdir(parents=False, exist_ok=False)
+        f = open(local_dir.joinpath("task_queue.yaml"), "w")
+        f.write("Task Queue: []\n")
+        f.close()
+
+    # get task id
+    tasks_queue = load(open(local_dir.joinpath("task_queue.yaml")), Loader = Loader)
+    task_list = tasks_queue['Task Queue']
+    task_id = len(task_list)
+
+    # create locak dir for the task
+    task_dir = local_dir.joinpath(f"task_{task_id}")
+    task_dir.mkdir(parents=False, exist_ok=False)
+    task_dir.joinpath("results").mkdir(parents=False, exist_ok=False)
+    task_dir.joinpath("logs").mkdir(parents=False, exist_ok=False)
+    f = open(task_dir.joinpath("task_info.yaml"), "w")
+    f.write(f"task_id: {task_id}")
+    f.close()
+
+    logging.basicConfig(level=logging.INFO, handlers=[
+        logging.FileHandler(
+            "{}/{}_{}.log".format(task_dir.joinpath("logs"), f"task_{task_id}", datetime.now().strftime("%m%d%Y_%H%M%S")),
+            'w', encoding='utf-8')])
+
+    # task create
+    if args.link is not None:
+        try:
+            task = Task.fromYoutubeLink(args.link, task_id, launch_cfg)
+        except:
+            shutil.rmtree(task_dir)
+            raise RuntimeError("failed to create task from youtube link")
+
+    # add task to the status queue
+    task_list.append({"id": task_id, "status": "created", "resource_status:": "local"})
+    stream = open(local_dir.joinpath("task_queue.yaml"), "w")
+    dump(tasks_queue, stream)
+
+    task.run_pipeline()
+
+
+
+
+    
+
+    
+    
+
+    

src/task.py

@@ -9,6 +9,9 @@ from os import getenv
 from enum import Enum
 from pathlib import Path
 from enum import Enum, auto
+import logging
+import subprocess
+
 
 
 """
@@ -42,6 +45,7 @@ SRT_Script : SrtScript
 """
 
 class TaskStatus(Enum):
+    CREATED = (auto(), None)
     INITIALIZING_ASR = (auto(), None)
     PRE_PROCESSING = (auto(), None)
     TRANSLATING = (auto(), 0.0)
@@ -50,22 +54,43 @@ class TaskStatus(Enum):
 
 
 class Task:
-    def __init__(self, task_id, audio_path, model, output_type):
-        # openai.api_key = getenv("OPENAI_API_KEY")
-        self.audio_path = audio_path
-        self.model = model
+    def __init__(self, task_id, task_local_dir, launch_info):
+        openai.api_key = getenv("OPENAI_API_KEY")
+        self.task_local_dir = task_local_dir
+        self.model = launch_info["model"]
         self.gpu_status = 0
-        self.output_type = output_type
+        self.output_type = launch_info["output_type"]
         self.task_id = task_id
         self.progress = NotImplemented
         self.SRT_Script = None
-        self.local_dump = Path()
+    
 
     @staticmethod
-    def fromYoutubeLink(youtube_url):
+    def fromYoutubeLink(youtube_url, task_id, launch_info):
         # convert to audio
+        local_dump = Path(launch_info['local_dump'])# should get from launch config
+        yt = YouTube(youtube_url)
+        video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
         
-        return Task(...)
+        if video:
+            video.download(str(local_dump.joinpath(f"task_{task_id}")), filename=f"task_{task_id}.mp4")
+            logging.info(f'Video download completed to {local_dump.joinpath(f"task_{task_id}")}!')
+        else:
+            raise FileNotFoundError(f"Video stream not found for link {youtube_url}")
+
+        audio = yt.streams.filter(only_audio=True).first()
+        if audio:
+            audio.download(str(local_dump.joinpath(f"task_{task_id}")), filename=f"task_{task_id}.mp3")
+            logging.info(f'Audio download completed to {local_dump.joinpath(f"task_{task_id}")}!')
+        else:
+            logging.info("download audio failed, using ffmpeg to extract audio")
+            subprocess.run(['ffmpeg', '-i', local_dump.joinpath(f"task_{task_id}").joinpath(f"task_{task_id}.mp4"), '-f', 'mp3', '-ab', '192000', '-vn', local_dump.joinpath(f"task_{task_id}").joinpath(f"task_{task_id}.mp3")])
+            logging.info("audio extraction finished")
+        
+        logging.info("Task Creation Complete.")
+        logging.info("Task Creation method: Youtube Link")
+
+        return Task(task_id, local_dump.joinpath(f"task_{task_id}"), launch_info)
 
     @staticmethod
     def fromAudioFile():