mirror of
https://github.com/HKUDS/VideoRAG.git
synced 2025-05-11 03:54:36 +03:00
78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
import os
|
|
import json
|
|
import logging
|
|
import warnings
|
|
import multiprocessing
|
|
import sys
|
|
|
|
warnings.filterwarnings("ignore")
|
|
logging.getLogger("httpx").setLevel(logging.WARNING)
|
|
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description="Set sub-category and CUDA device.")
|
|
parser.add_argument('--collection', type=str, default='4-rag-lecture')
|
|
parser.add_argument('--cuda', type=str, default='0')
|
|
args = parser.parse_args()
|
|
sub_category = args.collection
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda
|
|
os.environ["OPENAI_API_KEY"] = ""
|
|
|
|
from videorag._llm import *
|
|
from videorag.videorag import VideoRAG, QueryParam
|
|
|
|
longervideos_llm_config = LLMConfig(
|
|
embedding_func_raw = openai_embedding,
|
|
embedding_model_name = "text-embedding-3-small",
|
|
embedding_dim = 1536,
|
|
embedding_max_token_size = 8192,
|
|
embedding_batch_num = 32,
|
|
embedding_func_max_async = 16,
|
|
query_better_than_threshold = 0.2,
|
|
|
|
# LLM (we utilize gpt-4o-mini for all experiments)
|
|
best_model_func_raw = gpt_4o_mini_complete,
|
|
best_model_name = "gpt-4o-mini",
|
|
best_model_max_token_size = 32768,
|
|
best_model_max_async = 16,
|
|
|
|
cheap_model_func_raw = gpt_4o_mini_complete,
|
|
cheap_model_name = "gpt-4o-mini",
|
|
cheap_model_max_token_size = 32768,
|
|
cheap_model_max_async = 16
|
|
)
|
|
|
|
if __name__ == '__main__':
|
|
multiprocessing.set_start_method('spawn')
|
|
|
|
## learn
|
|
video_base_path = f'./longervideos/{sub_category}/videos/'
|
|
video_files = sorted(os.listdir(video_base_path))
|
|
video_paths = [os.path.join(video_base_path, f) for f in video_files]
|
|
videorag = VideoRAG(llm=longervideos_llm_config, working_dir=f"./longervideos/videorag-workdir/{sub_category}")
|
|
videorag.insert_video(video_path_list=video_paths)
|
|
|
|
## inference
|
|
with open(f'./longervideos/dataset.json', 'r') as f:
|
|
longervideos = json.load(f)
|
|
|
|
videorag = VideoRAG(llm=longervideos_llm_config, working_dir=f"./longervideos/videorag-workdir/{sub_category}")
|
|
videorag.load_caption_model(debug=False)
|
|
|
|
answer_folder = f'./longervideos/videorag-answers/{sub_category}'
|
|
os.makedirs(answer_folder, exist_ok=True)
|
|
|
|
collection_id = sub_category.split('-')[0]
|
|
querys = longervideos[collection_id][0]['questions']
|
|
for i in range(len(querys)):
|
|
query_id = querys[i]['id']
|
|
query = querys[i]['question']
|
|
param = QueryParam(mode="videorag")
|
|
param.wo_reference = True
|
|
print("Query: ", query)
|
|
|
|
response = videorag.query(query=query, param=param)
|
|
print(response)
|
|
with open(os.path.join(answer_folder, f'answer_{query_id}.md'), 'w') as f:
|
|
f.write(response)
|