Spaces:
Running
Running
| import sys | |
| import os | |
| project_dir = os.getcwd() | |
| sys.path.append(project_dir) | |
| import json | |
| from tqdm import tqdm | |
| from goldfish_lv import GoldFish_LV,split_subtitles,time_to_seconds | |
| import argparse | |
| import json | |
| import argparse | |
| import torch | |
| import re | |
| from tqdm import tqdm | |
| from PIL import Image | |
| # from openai import OpenAI | |
| from index import MemoryIndex | |
| import pysrt | |
| import chardet | |
| import torch | |
| import random | |
| import numpy as np | |
| import torch.backends.cudnn as cudnn | |
| def str2bool(v): | |
| if isinstance(v, bool): | |
| return v | |
| if v.lower() in ('yes', 'true', 't', 'y', '1'): | |
| return True | |
| elif v.lower() in ('no', 'false', 'f', 'n', '0'): | |
| return False | |
| else: | |
| raise argparse.ArgumentTypeError('Boolean value expected.') | |
| def get_arguments(): | |
| parser = argparse.ArgumentParser(description="Inference parameters") | |
| parser.add_argument("--neighbours", type=int, default=-1) | |
| parser.add_argument("--name", type=str,default="ckpt_92",help="name of the experiment") | |
| parser.add_argument("--exp_name", type=str,default="",help="name of the experiment") | |
| parser.add_argument("--add_unknown", action='store_true') | |
| parser.add_argument("--use_chatgpt", action='store_true') | |
| parser.add_argument("--use_choices_for_info", action='store_true') | |
| parser.add_argument("--use_gt_information", action='store_true') | |
| parser.add_argument("--inference_text", action='store_true') | |
| parser.add_argument("--use_gt_information_with_distraction", action='store_true') | |
| parser.add_argument("--num_distraction", type=int, default=2) | |
| parser.add_argument("--add_confidance_score", action='store_true') | |
| parser.add_argument("--use_original_video", action='store_true') | |
| parser.add_argument("--use_video_embedding", action='store_true') | |
| parser.add_argument("--use_clips_for_info", action='store_true') | |
| parser.add_argument("--use_GT_video", action='store_true') | |
| parser.add_argument("--use_gt_summary", action='store_true') | |
| parser.add_argument("--index_subtitles_together", action='store_true') | |
| parser.add_argument("--ask_the_question_early", action='store_true') | |
| parser.add_argument("--clip_in_ask_early", action='store_true') | |
| parser.add_argument("--use_coherent_description", action='store_true') | |
| parser.add_argument("--start", default=0, type=int) | |
| parser.add_argument("--end", default=100000, type=int) | |
| parser.add_argument("--vision_only", action='store_true') | |
| parser.add_argument("--model_summary_only", action='store_true') | |
| parser.add_argument("--subtitles_only", action='store_true') | |
| parser.add_argument("--subtitles_only_after_retrieval", action='store_true') | |
| parser.add_argument("--info_only", action='store_true') | |
| parser.add_argument("--cfg-path", default="test_configs/llama2_test_config.yaml") | |
| parser.add_argument("--ckpt", type=str, default="checkpoints/video_llama_checkpoint_last.pth") | |
| parser.add_argument("--add_subtitles", action='store_true') | |
| parser.add_argument("--eval_opt", type=str, default='all') | |
| parser.add_argument("--max_new_tokens", type=int, default=300) | |
| parser.add_argument("--batch_size", type=int, default=8) | |
| parser.add_argument("--lora_r", type=int, default=64) | |
| parser.add_argument("--lora_alpha", type=int, default=16) | |
| parser.add_argument("--video_path", type=str, help="path to the video") | |
| parser.add_argument("--use_openai_embedding",type=str2bool, default=False) | |
| parser.add_argument("--annotation_path", type=str, help="path to the annotation file") | |
| parser.add_argument("--videos_frames", type=str, help="path to the dataset extracted frames") | |
| parser.add_argument("--tvqa_json_subtitles", type=str, help="path to the tvqa json subtitles") | |
| parser.add_argument("--tvqa_clips_subtitles", type=str, help="path to the tvqa json") | |
| parser.add_argument("--options", nargs="+") | |
| return parser.parse_args() | |
| def clean_text(subtitles_text): | |
| # Remove unwanted characters except for letters, digits, and single quotes | |
| subtitles_text = re.sub(r'[^a-zA-Z0-9\s\']', '', subtitles_text) | |
| # Replace multiple spaces with a single space | |
| subtitles_text = re.sub(r'\s+', ' ', subtitles_text) | |
| return subtitles_text.strip() | |
| class TVQAEVAL (GoldFish_LV): | |
| def __init__(self, args: argparse.Namespace) -> None: | |
| super().__init__(args) | |
| self.tv_shows_mapping={"Grey's Anatomy":"grey_frames", 'How I Met You Mother':"met_frames", 'Friends':"friends_frames", 'The Big Bang Theory':"bbt_frames", 'House M.D.':"house_frames", 'Castle':"castle_frames"} | |
| self.save_long_videos_path = f"new_workspace/clips_summary/tvqa" | |
| if args.use_openai_embedding: | |
| self.save_embedding_path = f"new_workspace/open_ai_embedding/tvqa" | |
| else: | |
| self.save_embedding_path = f"new_workspace/embedding/tvqa" | |
| os.makedirs(self.save_long_videos_path, exist_ok=True) | |
| self.max_sub_len=400 | |
| self.max_num_images=45 | |
| self.fps=3 | |
| with open(args.tvqa_json_subtitles) as f: | |
| self.subtitles_list=json.load(f) | |
| self.subtitles={} | |
| for sub in self.subtitles_list: | |
| self.subtitles[sub["vid_name"]]=sub["sub"] | |
| def _get_TVs_data(self): | |
| json_file_path=args.annotation_path | |
| frames_path=args.videos_frames | |
| subtitle_path=args.tvqa_clips_subtitles | |
| with open (json_file_path) as f: | |
| tv_shows_data=json.load(f) | |
| return tv_shows_data,frames_path,subtitle_path | |
| def _get_shows_subtitles(self,clip_subtitles_path): | |
| try : | |
| with open(clip_subtitles_path, 'rb') as f: | |
| result = chardet.detect(f.read()) | |
| clip_subtitles = pysrt.open(clip_subtitles_path, encoding=result['encoding']) | |
| return clip_subtitles | |
| except: | |
| print("No subtitles found") | |
| return [] | |
| def episode_inference(self,clips,folder_name,use_subtitles): | |
| max_caption_index = 0 | |
| max_subtitle_index = 0 | |
| preds={} | |
| important_data = {} | |
| videos_summaries=[] | |
| batch_size=args.batch_size | |
| batch_images=[] | |
| batch_instructions=[] | |
| conversations=[] | |
| clips_names=[] | |
| for clip_name in tqdm(clips,desc="Inference Episode clips"): | |
| conversation="" | |
| try: | |
| for subtitle in self.subtitles[clip_name]: | |
| conversation+=subtitle['text']+" " | |
| except: | |
| pass | |
| conversations.append(clean_text(conversation)) | |
| images,img_placeholder=self.prepare_input_images(clip_name,folder_name,use_subtitles) | |
| instruction = img_placeholder + '\n' + self.summary_instruction | |
| batch_images.append(images) | |
| batch_instructions.append(instruction) | |
| clips_names.append(clip_name) | |
| if len(batch_images) < batch_size: | |
| continue | |
| batch_images = torch.stack(batch_images) | |
| batch_pred=self.run_images(batch_images,batch_instructions) | |
| for i,pred in enumerate(batch_pred): | |
| max_caption_index += 1 | |
| videos_summaries.append(pred) | |
| if args.use_coherent_description: | |
| preds[f'caption_{max_caption_index}__{clips_names[i]}'] = f"model_summary :{pred}\nVideo conversation :{conversations[i]}" | |
| else: | |
| if args.index_subtitles_together and use_subtitles: | |
| if conversations[i] != "": | |
| max_subtitle_index+=1 | |
| important_data.update({f"subtitle_{max_subtitle_index}__{clips_names[i]}": conversations[i]}) | |
| preds[f'caption_{max_caption_index}__{clips_names[i]}'] = pred | |
| batch_images=[] | |
| batch_instructions=[] | |
| clips_names=[] | |
| conversations=[] | |
| # run inference for the last batch | |
| if len(batch_images)>0: | |
| batch_images = torch.stack(batch_images) | |
| batch_pred=self.run_images(batch_images,batch_instructions) | |
| for i,pred in enumerate(batch_pred): | |
| max_caption_index += 1 | |
| videos_summaries.append(pred) | |
| if args.use_coherent_description: | |
| preds[f'caption_{max_caption_index}__{clips_names[i]}'] = f"model_summary :{pred}\nVideo conversation :{conversations[i]}" | |
| else: | |
| if args.index_subtitles_together and use_subtitles: | |
| if conversations[i] != "": | |
| max_subtitle_index+=1 | |
| important_data.update({f"subtitle_{max_subtitle_index}__{clips_names[i]}": conversations[i]}) | |
| preds[f'caption_{max_caption_index}__{clips_names[i]}'] = pred | |
| batch_images=[] | |
| batch_instructions=[] | |
| clips_names=[] | |
| return preds,important_data | |
| def episode_inference_only_subtitles(self,clips,tv_images_path,subtitle_path): | |
| max_subtitle_index = 0 | |
| important_data = {} | |
| for c_name in tqdm(clips,desc="Inference Episode clips"): | |
| clip_subtitles_path=os.path.join(subtitle_path,c_name+".srt") | |
| clip_subtitles=self._get_shows_subtitles(clip_subtitles_path) | |
| conversation="" | |
| if args.index_subtitles_together: | |
| if self.subtitles.get(c_name,False): | |
| for subtitle in self.subtitles[c_name]: | |
| conversation+=subtitle['text']+" " | |
| conversation=clean_text(conversation) | |
| if conversation != "": | |
| max_subtitle_index+=1 | |
| important_data.update({f"subtitle_{max_subtitle_index}__{c_name}": conversation}) | |
| return important_data | |
| def prepare_input_images(self,clip_name,folder_name,use_subtitles): | |
| tv_shows_data,frames_path,subtitle_path=self._get_TVs_data() | |
| tv_images_path =os.path.join(frames_path,folder_name) | |
| clip_path=os.path.join(tv_images_path,clip_name) | |
| total_frames=len(os.listdir(clip_path)) | |
| sampling_interval=int(total_frames//self.max_num_images) | |
| if sampling_interval==0: | |
| sampling_interval=1 | |
| images=[] | |
| img_placeholder = "" | |
| video_frames_path = os.path.join(frames_path,folder_name,clip_name) | |
| total_num_frames=len(os.listdir(video_frames_path)) | |
| sampling_interval = round(total_num_frames / self.max_num_images) | |
| if sampling_interval == 0: | |
| sampling_interval = 1 | |
| subtitle_text_in_interval = "" | |
| history_subtitles = {} | |
| number_of_sub_words=0 | |
| for i,frame in enumerate(sorted(os.listdir(video_frames_path))): | |
| # Find the corresponding subtitle for the frame and combine the interval subtitles into one subtitle | |
| # we choose 1 frame for every 2 seconds,so we need to combine the subtitles in the interval of 2 seconds | |
| if self.subtitles.get(clip_name,False) and use_subtitles: | |
| for subtitle in self.subtitles[clip_name]: | |
| if (subtitle['start'] <= (i / self.fps) <= subtitle['end']) and subtitle['text'] not in subtitle_text_in_interval: | |
| if not history_subtitles.get(subtitle['text'],False): | |
| subtitle_text_in_interval+=subtitle['text']+" " | |
| history_subtitles[subtitle['text']]=True | |
| break | |
| if i % sampling_interval == 0: | |
| frame = Image.open(os.path.join(video_frames_path,frame)).convert("RGB") | |
| frame = self.vis_processor(frame) | |
| images.append(frame) | |
| img_placeholder += '<Img><ImageHere>' | |
| if number_of_sub_words<self.max_sub_len and use_subtitles: | |
| if subtitle_text_in_interval != "": | |
| subtitle_text_in_interval=clean_text(subtitle_text_in_interval) | |
| img_placeholder+=f'<Cap>{subtitle_text_in_interval}' | |
| number_of_sub_words+=len(subtitle_text_in_interval.split(' ')) | |
| subtitle_text_in_interval = "" | |
| if len(images) >= self.max_num_images: | |
| break | |
| if len(images) ==0: | |
| print("Video not found",video_frames_path) | |
| if 0 <len(images) < self.max_num_images: | |
| last_item = images[-1] | |
| while len(images) < self.max_num_images: | |
| images.append(last_item) | |
| img_placeholder += '<Img><ImageHere>' | |
| images = torch.stack(images) | |
| return images,img_placeholder | |
| def clip_inference(self,clips_name,folders_name,prompts): | |
| setup_seeds(seed) | |
| images_batch, instructions_batch = [], [] | |
| for clip_name,folder_name, prompt in zip(clips_name,folders_name, prompts): | |
| images,img_placeholder=self.prepare_input_images(clip_name,folder_name,use_subtitles=not args.vision_only) | |
| instruction = img_placeholder + '\n' + prompt | |
| images_batch.append(images) | |
| instructions_batch.append(instruction) | |
| # run inference for the batch | |
| images_batch=torch.stack(images_batch) | |
| batch_pred=self.run_images(images_batch,instructions_batch) | |
| return batch_pred | |
| def prepare_prompt(self,qa): | |
| prompt=qa["q"]+" \n\n As you watched in this video Choose ONE suitable answer from these mutiple choices \n" | |
| for i,choice in enumerate(["a0","a1","a2","a3","a4"]): | |
| prompt+=f"option {i}: {qa[choice]} \n" | |
| if args.add_unknown and args.add_confidance_score: | |
| # Add unknown option | |
| prompt+=f"option 5: Can't answer based on the provided information\n" | |
| prompt+="\n Your output should be THE NUMBER OF THE CORRECT ANSWER FROM THE CHOICES FROM 0 TO 5 INCLUSIVE and aslo output a CONFIDANCE SCORE FROM 0 TO 5 representing how confident you are with your answer where 0 is the least confident and 5 is the most confident" | |
| elif args.add_unknown: | |
| prompt+=f"option 5: Can't answer based on the provided information\n" | |
| prompt+="\n Your output should be THE NUMBER OF THE CORRECT ANSWER FROM THE CHOICES FROM 0 TO 5 INCLUSIVE" | |
| elif args.add_confidance_score: | |
| prompt+="\n Your output should be THE NUMBER OF THE CORRECT ANSWER FROM THE CHOICES FROM 0 TO 4 INCLUSIVE and aslo output a CONFIDANCE SCORE FROM 0 TO 5 representing how confident you are with your answer where 0 is the least confident and 5 is the most confident" | |
| else: | |
| prompt+="\n Your output should be THE NUMBER OF THE CORRECT ANSWER FROM THE CHOICES FROM 0 TO 4 INCLUSIVE" | |
| return prompt | |
| def get_most_related_clips(self,qa,related_context_keys): | |
| if args.use_gt_information: | |
| most_related_clips=[qa['vid_name']] | |
| elif args.use_gt_information_with_distraction: | |
| most_related_clips=[qa['vid_name']] | |
| for context_key in related_context_keys: | |
| if len(context_key.split('__'))>1: | |
| most_related_clips.append(context_key.split('__')[1]) | |
| if len(most_related_clips)==args.num_distraction+1: | |
| break | |
| else: | |
| most_related_clips=[] | |
| for context_key in related_context_keys: | |
| if len(context_key.split('__'))>1: | |
| most_related_clips.append(context_key.split('__')[1]) | |
| if len(most_related_clips)==args.neighbours: | |
| break | |
| assert len(most_related_clips)!=0, f"No related clips found {related_context_keys}" | |
| return most_related_clips | |
| def use_clips_for_info(self,qa_list,related_context_keys_list,external_memory): | |
| total_batch_pred=[] | |
| questions=[] | |
| related_information_list=[] | |
| related_context_keys_list_new=[] | |
| for qa,related_context_keys in zip(qa_list,related_context_keys_list): | |
| most_related_clips=self.get_most_related_clips(qa,related_context_keys) | |
| folder_name=self.tv_shows_mapping[qa['show_name']] | |
| question=qa['q']+ "\nand these are the choices :\n" | |
| for i,choice in enumerate(["a0","a1","a2","a3","a4"]): | |
| question+=f"option {i}: {qa[choice]} \n" | |
| if args.add_unknown: | |
| question+= "option 5: Can't answer based on the provided information\n" | |
| question+="\n Your output should be THE NUMBER OF THE CORRECT ANSWER FROM THE CHOICES FROM 0 TO 5 INCLUSIVE" | |
| else: | |
| question+="\n Your output should be THE NUMBER OF THE CORRECT ANSWER FROM THE CHOICES FROM 0 TO 4 INCLUSIVE" | |
| if args.use_choices_for_info: | |
| # prompt=self.prepare_prompt(qa) | |
| # prompt+=" and also provide an EXPLAINATION for your answer and If you don't know the answer, say that you don't know.\n\n" | |
| prompt=f"From this video extract the related information to This multichioce question and provide an explaination for your answer and If you don't know the answer, say 'I DON'T KNOW' as option 5 because maybe the questoin is not related to the video content.\n the question is :\n {question}\n your answer :" | |
| else: | |
| prompt=f"As you watched in this video answer this {qa['q']}\n\n and also provide an EXPLAINATION for your answer and If you don't know the answer, say that you don't know.\n\n" | |
| all_info=self.clip_inference(most_related_clips,[folder_name]*len(most_related_clips),[prompt]*len(most_related_clips)) | |
| # concatinate all the information together | |
| related_information="" | |
| for info,clip_name in zip(all_info,most_related_clips): | |
| clip_conversation="" | |
| general_sum="" | |
| for key in external_memory.documents.keys(): | |
| if clip_name in key and 'caption' in key: | |
| general_sum="Clip Summary: "+external_memory.documents[key] | |
| if clip_name in key and 'subtitle' in key: | |
| clip_conversation="Clip Subtitles: "+external_memory.documents[key] | |
| if args.use_coherent_description: | |
| related_information+=f"question_related_information: {info},{general_sum}\n" | |
| else: | |
| # related_information+=f"{general_sum},{clip_conversation},question_related_information: {info}\n" | |
| # related_information+=f"question_related_information: {info},{clip_conversation}\n" | |
| if args.model_summary_only: | |
| related_information+=f"{general_sum},question_related_information: {info}\n" | |
| elif args.info_only: | |
| related_information+=f"question_related_information: {info}\n" | |
| elif args.subtitles_only: | |
| related_information+=f"{clip_conversation},question_related_information: {info}\n" | |
| elif args.subtitles_only_after_retrieval: | |
| related_information+=f"{clip_conversation},question_related_information: {info}\n" | |
| else: | |
| related_information+=f"{general_sum},{clip_conversation},question_related_information: {info}\n" | |
| questions.append(question) | |
| related_information_list.append(related_information) | |
| related_context_keys.append(related_information) | |
| related_context_keys_list_new.append(related_context_keys) | |
| if len(questions)< args.batch_size: | |
| continue | |
| setup_seeds(seed) | |
| if args.use_chatgpt : | |
| batch_pred=self.inference_RAG_chatGPT(questions, related_information_list) | |
| else: | |
| batch_pred=self.inference_RAG(questions, related_information_list) | |
| for pred in batch_pred: | |
| total_batch_pred.append(pred) | |
| questions=[] | |
| related_information_list=[] | |
| if len(questions)>0: | |
| setup_seeds(seed) | |
| if args.use_chatgpt : | |
| batch_pred=self.inference_RAG_chatGPT(questions, related_information_list) | |
| else: | |
| batch_pred=self.inference_RAG(questions, related_information_list) | |
| for pred in batch_pred: | |
| total_batch_pred.append(pred) | |
| return total_batch_pred,related_context_keys_list_new | |
| def answer_TV_questions_RAG(self,qa_list,external_memory,episode_clips,episode_name): | |
| related_context_keys_list,related_context_documents_list=[],[] | |
| setup_seeds(seed) | |
| for qa in qa_list: | |
| question_choices=qa['q']+ "\n and these are the options for the question\n\n" | |
| for i,choice in enumerate(["a0","a1","a2","a3","a4"]): | |
| question_choices+=f"option {i}: {qa[choice]} \n\n" | |
| related_context_documents,related_context_keys = external_memory.search_by_similarity(question_choices) | |
| related_context_documents_list.append(related_context_documents) | |
| related_context_keys_list.append(related_context_keys) | |
| if args.use_clips_for_info: | |
| batch_pred,related_context_keys_list=self.use_clips_for_info(qa_list,related_context_keys_list,external_memory) | |
| else: | |
| prompts=[] | |
| related_context_documents_text_list=[] | |
| for qa,related_context_documents,related_context_keys in zip(qa_list,related_context_documents_list,related_context_keys_list): | |
| related_information="" | |
| most_related_clips=self.get_most_related_clips(qa,related_context_keys) | |
| for clip_name in most_related_clips: | |
| clip_conversation="" | |
| general_sum="" | |
| for key in external_memory.documents.keys(): | |
| if clip_name in key and 'caption' in key: | |
| general_sum="Clip Summary: "+external_memory.documents[key] | |
| if clip_name in key and 'subtitle' in key: | |
| clip_conversation="Clip Subtitles: "+external_memory.documents[key] | |
| # related_information+=f"{general_sum},{clip_conversation}\n" | |
| if args.use_coherent_description: | |
| related_information+=f"{general_sum}\n" | |
| else: | |
| if args.model_summary_only: | |
| related_information+=f"{general_sum}\n" | |
| elif args.subtitles_only: | |
| related_information+=f"{clip_conversation}\n" | |
| else: | |
| related_information+=f"{general_sum},{clip_conversation}\n" | |
| prompt=self.prepare_prompt(qa) | |
| prompts.append(prompt) | |
| related_context_documents_text_list.append(related_information) | |
| setup_seeds(seed) | |
| if args.use_chatgpt: | |
| batch_pred=self.inference_RAG_chatGPT(prompts, related_context_documents_text_list) | |
| else: | |
| batch_pred=self.inference_RAG(prompts, related_context_documents_text_list) | |
| return batch_pred ,related_context_keys_list | |
| def answer_episode_questions(self,questions,information_RAG_path,embedding_path,episode_clips): | |
| external_memory=MemoryIndex(args.neighbours, use_openai=args.use_openai_embedding) | |
| if os.path.exists(embedding_path): | |
| external_memory.load_embeddings_from_pkl(embedding_path) | |
| else: | |
| external_memory.load_documents_from_json(information_RAG_path,embedding_path) | |
| episode_name=information_RAG_path.split('/')[-1].split('.')[0] | |
| pred_json=[] | |
| batch_questions=[] | |
| for qa in tqdm(questions,desc="Answering questions"): | |
| batch_questions.append(qa) | |
| if len(batch_questions)<args.batch_size: | |
| continue | |
| batch_pred,batch_related_context_keys = self.answer_TV_questions_RAG(batch_questions,external_memory,episode_clips,episode_name) | |
| for pred,related_context_keys,qa in zip(batch_pred,batch_related_context_keys,batch_questions): | |
| qa['pred']=pred | |
| qa['related_context_keys']=related_context_keys | |
| pred_json.append(qa) | |
| batch_questions=[] | |
| if len(batch_questions)>0: | |
| batch_pred,batch_related_context_keys = self.answer_TV_questions_RAG(batch_questions,external_memory,episode_clips,episode_name) | |
| for pred,related_context_keys,qa in zip(batch_pred,batch_related_context_keys,batch_questions): | |
| qa['pred']=pred | |
| qa['related_context_keys']=related_context_keys | |
| pred_json.append(qa) | |
| return pred_json | |
| def eval_tv_shows(self,): | |
| tv_shows_data,frames_path,subtitle_path=self._get_TVs_data() | |
| full_questions_result=[] | |
| number_of_episodes=0 | |
| start=args.start | |
| end=args.end | |
| for show in tqdm(tv_shows_data,desc="Inference TV shows"): | |
| for season in tqdm(tv_shows_data[show],desc=f"Inference {show} seasons"): | |
| for episode in tqdm(tv_shows_data[show][season],desc=f"Inference {show} {season} episodes"): | |
| # Generate clips summary and store the important data (summary and subtitles) in json file | |
| if start<=number_of_episodes<end: | |
| folder_name=self.tv_shows_mapping[show] | |
| tv_images_path =os.path.join(frames_path,folder_name) | |
| os.makedirs(self.save_long_videos_path, exist_ok=True) | |
| save_name="" if args.index_subtitles_together else "no_subtitles_" | |
| save_name="subtitles_only" if args.subtitles_only else save_name | |
| save_name="use_coherent_description" if args.use_coherent_description else save_name | |
| file_path=os.path.join(self.save_long_videos_path,save_name+folder_name+"_"+season+"_"+episode+".json") | |
| embedding_path=os.path.join(self.save_embedding_path,save_name+folder_name+"_"+season+"_"+episode+".pkl") | |
| # options don't require rerunning the inference | |
| save_name+="_unknown" if args.add_unknown else "" | |
| save_name+="_clips_for_info" if args.use_clips_for_info else "" | |
| save_name+="_chatgpt" if args.use_chatgpt else "" | |
| save_name+="_choices_for_info" if args.use_choices_for_info else "" | |
| save_name+="_info_only" if args.info_only else "" | |
| save_name+="_subtitles_only" if args.subtitles_only else "" | |
| save_name+="_subtitles_only_after_retrieval" if args.subtitles_only_after_retrieval else "" | |
| if os.path.exists(file_path): | |
| with open(file_path, 'r') as file: | |
| important_data = json.load(file) | |
| print("Already processed") | |
| else: | |
| episode_clips=tv_shows_data[show][season][episode]['clips'] | |
| if args.subtitles_only : | |
| important_data=self.episode_inference_only_subtitles(episode_clips,tv_images_path,subtitle_path) | |
| else: | |
| preds,important_data=self.episode_inference(episode_clips,folder_name,use_subtitles=not args.vision_only) | |
| important_data.update(preds) | |
| # if not args.subtitles_only : | |
| # summary = self.compine_summaries(important_data) | |
| # preds['summary'] = summary | |
| # important_data["summary"]=summary | |
| with open(file_path, 'w') as file: | |
| json.dump(important_data, file, indent=4) | |
| # Answer questions | |
| questions=tv_shows_data[show][season][episode]['questions'] | |
| episode_clips=tv_shows_data[show][season][episode]['clips'] | |
| episode_name=file_path.split('/')[-1].split('.')[0] | |
| pred_json=self.answer_episode_questions(questions,file_path,embedding_path,episode_clips) | |
| full_questions_result.extend(pred_json) | |
| save_dir=f"new_workspace/results/tvqa/{args.exp_name}/{save_name}_{args.neighbours}_neighbours" | |
| os.makedirs(save_dir, exist_ok=True) | |
| with open(f"{save_dir}/{episode_name}.json", 'w') as fp: | |
| json.dump(pred_json, fp) | |
| print(f"Episode {episode_name} prediction saved to {save_dir}/{episode_name}_pred_{args.neighbours}.json") | |
| number_of_episodes+=1 | |
| with open(f"{save_dir}/full_pred_{start}_{end}.json", 'w') as fp: | |
| json.dump(full_questions_result, fp) | |
| print(f"TV shows prediction saved to {save_dir}/full_pred_{start}{end}.json") | |
| args=get_arguments() | |
| def setup_seeds(seed): | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed(seed) | |
| cudnn.benchmark = False | |
| cudnn.deterministic = True | |
| import yaml | |
| with open('test_configs/llama2_test_config.yaml') as file: | |
| config = yaml.load(file, Loader=yaml.FullLoader) | |
| seed=config['run']['seed'] | |
| print("seed",seed) | |
| if __name__ == "__main__": | |
| setup_seeds(seed) | |
| tvqa_eval=TVQAEVAL(args) | |
| tvqa_eval.eval_tv_shows() |