import argparse import re import os from tqdm import tqdm import pandas as pd import torch from natsort import index_natsorted from vllm import LLM, SamplingParams from transformers import AutoTokenizer from utils.logger import logger def extract_output(s, prefix='"rewritten description": '): """Customize the function according to the prompt.""" # Since some LLMs struggles to output strictly formatted JSON strings as specified by the prompt, # thus manually parse the output string `{"rewritten description": "your rewritten description here"}`. match = re.search(r"{(.+?)}", s, re.DOTALL) if not match: logger.warning(f"{s} is not in the json format. Return None.") return None output = match.group(1).strip() if output.startswith(prefix): output = output[len(prefix) :] if output[0] == '"' and output[-1] == '"': return output[1:-1] else: logger.warning(f"{output} does not start and end with the double quote. Return None.") return None else: logger.warning(f"{output} does not start with {prefix}. Return None.") return None def parse_args(): parser = argparse.ArgumentParser(description="Rewrite the video caption by LLMs.") parser.add_argument( "--video_metadata_path", type=str, required=True, help="The path to the video dataset metadata (csv/jsonl)." ) parser.add_argument( "--video_path_column", type=str, default=None, help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).", ) parser.add_argument( "--caption_column", type=str, default="caption", help="The column contains the video caption.", ) parser.add_argument( "--batch_size", type=int, default=128, required=False, help="The batch size for vllm inference. Adjust according to the number of GPUs to maximize inference throughput.", ) parser.add_argument( "--model_name", type=str, default="NousResearch/Meta-Llama-3-8B-Instruct", ) parser.add_argument( "--prompt", type=str, required=True, help="A string or a txt file contains the prompt.", ) parser.add_argument( "--prefix", type=str, required=True, help="The prefix to extract the output from LLMs.", ) parser.add_argument("--saved_path", type=str, required=True, help="The save path to the output results (csv/jsonl).") parser.add_argument("--saved_freq", type=int, default=1, help="The frequency to save the output results.") args = parser.parse_args() return args def main(): args = parse_args() if args.video_metadata_path.endswith(".csv"): video_metadata_df = pd.read_csv(args.video_metadata_path) elif args.video_metadata_path.endswith(".jsonl"): video_metadata_df = pd.read_json(args.video_metadata_path, lines=True) elif args.video_metadata_path.endswith(".json"): video_metadata_df = pd.read_json(args.video_metadata_path) else: raise ValueError(f"The {args.video_metadata_path} must end with .csv, .jsonl or .json.") saved_suffix = os.path.splitext(args.saved_path)[1] if saved_suffix not in set([".csv", ".jsonl", ".json"]): raise ValueError(f"The saved_path must end with .csv, .jsonl or .json.") if os.path.exists(args.saved_path) and args.video_path_column is not None: if args.saved_path.endswith(".csv"): saved_metadata_df = pd.read_csv(args.saved_path) elif args.saved_path.endswith(".jsonl"): saved_metadata_df = pd.read_json(args.saved_path, lines=True) # Filter out the unprocessed video-caption pairs by setting the indicator=True. merged_df = video_metadata_df.merge(saved_metadata_df, on=args.video_path_column, how="outer", indicator=True) video_metadata_df = merged_df[merged_df["_merge"] == "left_only"] # Sorting to guarantee the same result for each process. video_metadata_df = video_metadata_df.iloc[index_natsorted(video_metadata_df[args.video_path_column])].reset_index( drop=True ) logger.info( f"Resume from {args.saved_path}: {len(saved_metadata_df)} processed and {len(video_metadata_df)} to be processed." ) if args.prompt.endswith(".txt") and os.path.exists(args.prompt): with open(args.prompt, "r") as f: args.prompt = "".join(f.readlines()) logger.info(f"Prompt: {args.prompt}") if args.video_path_column is not None: video_path_list = video_metadata_df[args.video_path_column].tolist() if args.caption_column in video_metadata_df.columns: sampled_frame_caption_list = video_metadata_df[args.caption_column].tolist() else: # When two columns with the same name, the dataframe merge operation on will distinguish them by adding 'x' and 'y'. sampled_frame_caption_list = video_metadata_df[args.caption_column + "_x"].tolist() CUDA_VISIBLE_DEVICES = os.getenv("CUDA_VISIBLE_DEVICES", None) tensor_parallel_size = torch.cuda.device_count() if CUDA_VISIBLE_DEVICES is None else len(CUDA_VISIBLE_DEVICES.split(",")) logger.info(f"Automatically set tensor_parallel_size={tensor_parallel_size} based on the available devices.") llm = LLM(model=args.model_name, trust_remote_code=True, tensor_parallel_size=tensor_parallel_size) if "Meta-Llama-3" in args.model_name: if "Meta-Llama-3-70B" in args.model_name: # Llama-3-70B should use the tokenizer from Llama-3-8B # https://github.com/vllm-project/vllm/issues/4180#issuecomment-2068292942 tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct") else: tokenizer = AutoTokenizer.from_pretrained(args.model_name) stop_token_ids = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")] sampling_params = SamplingParams(temperature=0.7, top_p=1, max_tokens=1024, stop_token_ids=stop_token_ids) else: tokenizer = AutoTokenizer.from_pretrained(args.model_name) sampling_params = SamplingParams(temperature=0.7, top_p=1, max_tokens=1024) result_dict = {args.caption_column: []} if args.video_path_column is not None: result_dict = {args.video_path_column: [], args.caption_column: []} for i in tqdm(range(0, len(sampled_frame_caption_list), args.batch_size)): if args.video_path_column is not None: batch_video_path = video_path_list[i : i + args.batch_size] batch_caption = sampled_frame_caption_list[i : i + args.batch_size] batch_prompt = [] for caption in batch_caption: # batch_prompt.append("user:" + args.prompt + str(caption) + "\n assistant:") messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": args.prompt + "\n" + str(caption)}, ] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) batch_prompt.append(text) batch_output = llm.generate(batch_prompt, sampling_params) batch_output = [output.outputs[0].text.rstrip() for output in batch_output] batch_output = [extract_output(output, prefix=args.prefix) for output in batch_output] # Filter out data that does not meet the output format. batch_result = [] if args.video_path_column is not None: for video_path, output in zip(batch_video_path, batch_output): if output is not None: batch_result.append((video_path, output)) batch_video_path, batch_output = zip(*batch_result) result_dict[args.video_path_column].extend(batch_video_path) else: for output in batch_output: if output is not None: batch_result.append(output) result_dict[args.caption_column].extend(batch_result) # Save the metadata every args.saved_freq. if i != 0 and ((i // args.batch_size) % args.saved_freq) == 0: if len(result_dict[args.caption_column]) > 0: result_df = pd.DataFrame(result_dict) if args.saved_path.endswith(".csv"): header = True if not os.path.exists(args.saved_path) else False result_df.to_csv(args.saved_path, header=header, index=False, mode="a") elif args.saved_path.endswith(".jsonl"): result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False) elif args.saved_path.endswith(".json"): # Append is not supported. if os.path.exists(args.saved_path): saved_df = pd.read_json(args.saved_path, orient="records") result_df = pd.concat([saved_df, result_df], ignore_index=True) result_df.to_json(args.saved_path, orient="records", indent=4, force_ascii=False) logger.info(f"Save result to {args.saved_path}.") result_dict = {args.caption_column: []} if args.video_path_column is not None: result_dict = {args.video_path_column: [], args.caption_column: []} if len(result_dict[args.caption_column]) > 0: result_df = pd.DataFrame(result_dict) if args.saved_path.endswith(".csv"): header = True if not os.path.exists(args.saved_path) else False result_df.to_csv(args.saved_path, header=header, index=False, mode="a") elif args.saved_path.endswith(".jsonl"): result_df.to_json(args.saved_path, orient="records", lines=True, mode="a") elif args.saved_path.endswith(".json"): # Append is not supported. if os.path.exists(args.saved_path): saved_df = pd.read_json(args.saved_path, orient="records") result_df = pd.concat([saved_df, result_df], ignore_index=True) result_df.to_json(args.saved_path, orient="records", indent=4, force_ascii=False) logger.info(f"Save the final result to {args.saved_path}.") if __name__ == "__main__": main()