ManojINaik/TheoremExplainAgent
0
1import os2import json3import argparse4import tempfile5from typing import Dict, List, Union6from datetime import datetime7 8from dotenv import load_dotenv9from moviepy import VideoFileClip10 11from mllm_tools.litellm import LiteLLMWrapper12from mllm_tools.gemini import GeminiWrapper13from eval_suite.utils import calculate_geometric_mean14from eval_suite.text_utils import parse_srt_to_text, fix_transcript, evaluate_text15from eval_suite.video_utils import evaluate_video_chunk_new16from eval_suite.image_utils import evaluate_sampled_images17 18load_dotenv()19 20with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "src", "utils", "allowed_models.json")) as f:21 ALLOWED_MODELS = json.load(f)["allowed_models"]22 23 24def combine_results(output_folder: str, combined_file: str, results: Dict[str, Dict]) -> None:25 """26 Combine all evaluation results into a single file.27 28 Args:29 output_folder (str): Directory to store the combined file.30 combined_file (str): Name of the combined file.31 results (Dict[str, Dict]): Dictionary of evaluation results with file names as keys.32 33 Returns:34 None35 """36 combined_path = os.path.join(output_folder, combined_file)37 with open(combined_path, 'w') as output_file:38 json.dump(results, output_file, indent=4)39 40 41def save_individual_result(output_folder: str, file_name: str, result: Dict) -> None:42 """43 Save individual evaluation result to a file.44 45 Args:46 output_folder (str): Directory to store the evaluation file.47 file_name (str): Name of the file.48 result (Dict): Evaluation result.49 50 Returns:51 None52 """53 current_time = datetime.now().strftime("%Y%m%d_%H%M%S")54 result_file = f"evaluation_{file_name}_{current_time}.json"55 os.makedirs(output_folder, exist_ok=True)56 result_path = os.path.join(output_folder, result_file)57 with open(result_path, 'w') as output_file:58 json.dump(result, output_file, indent=4)59 60 61def evaluate_text_file(model, transcript_path, retry_limit):62 """63 Evaluate a text file using the provided model.64 65 Args:66 model: The model to use for evaluation.67 transcript_path (str): Path to the transcript file (.srt or .txt).68 retry_limit (int): Number of retry attempts for evaluation.69 70 Returns:71 Dict or None: Evaluation results if successful, None if file format unsupported.72 """73 if not transcript_path.endswith(('.srt', '.txt')):74 print(f"Skipping {transcript_path}: Unsupported file format for text evaluation.")75 return None76 77 if transcript_path.endswith(".srt"):78 transcript = parse_srt_to_text(transcript_path)79 elif transcript_path.endswith(".txt"):80 with open(transcript_path) as f:81 transcript = f.read().strip()82 else:83 raise ValueError("Unrecognized transcript file format.")84 85 capital_letter_proportion = sum(1 for c in transcript if c.isupper()) / sum(1 for c in transcript if c.isalpha())86 if capital_letter_proportion < 0.01:87 transcript = fix_transcript(model, transcript)88 89 print(f"Performing text evaluation: {os.path.basename(transcript_path)}")90 result = evaluate_text(model, transcript, retry_limit)91 return result92 93 94def evaluate_video_file(model, video_path, transcript_path, description_path, target_fps=None, output_folder=None):95 """96 Evaluate a video file using the provided model.97 98 Args:99 model: The model to use for evaluation.100 video_path (str): Path to the video file.101 transcript_path (str): Path to the transcript file.102 description_path (str): Path to the description file.103 target_fps (int, optional): Target frames per second for video processing.104 output_folder (str, optional): Directory to store output files.105 106 Returns:107 Dict or None: Evaluation results if successful, None if file format unsupported.108 """109 if not video_path.endswith(('.mp4', '.mkv')):110 print(f"Skipping {video_path}: Unsupported file format for video evaluation.")111 return None112 113 moviepy_temp_dir = os.path.join(output_folder, "moviepy_temp")114 115 # Chunking116 num_chunks = 10117 with VideoFileClip(video_path) as clip:118 duration = clip.duration119 chunk_duration = duration / num_chunks120 results = []121 122 # Create a temporary directory in the output_folder123 temp_dir_parent = output_folder or os.getcwd()124 with tempfile.TemporaryDirectory(dir=temp_dir_parent) as temp_dir:125 for i in range(10):126 start = i * chunk_duration127 end = min(start + chunk_duration, duration)128 chunk = clip.subclipped(start, end)129 chunk_path = os.path.join(temp_dir, f"chunk_{i+1}.mp4")130 # Explicitly set the temp_audiofile path with matching codec131 temp_audiofile = os.path.join(moviepy_temp_dir, f"temp_audio_chunk_{i+1}.m4a")132 chunk.write_videofile(133 chunk_path,134 codec="libx264",135 audio_codec="aac",136 temp_audiofile=temp_audiofile,137 audio_bitrate="192k",138 preset="ultrafast", # Speed up encoding139 logger=None140 )141 # Create processed videos folder inside output_folder142 processed_videos_dir = os.path.join(output_folder, "processed_videos")143 save_path = os.path.join(processed_videos_dir, f"processed_chunk_{i+1}.mp4")144 result = evaluate_video_chunk_new(145 model,146 chunk_path,147 transcript_path,148 description_path,149 target_fps=target_fps,150 save_processed_video=save_path151 )152 results.append(result)153 154 score_dict = {}155 for key in results[0]["evaluation"].keys():156 score_dict[key] = []157 for result in results:158 score_dict[key].append(result["evaluation"][key]["score"])159 160 evaluation = {}161 for key, scores in score_dict.items():162 evaluation[key] = {"score": calculate_geometric_mean(scores)}163 164 result_json = {165 "evaluation": evaluation,166 "video_chunks": results167 }168 return result_json169 170 171def extract_scores(data: Union[Dict, List]) -> List[int]:172 """173 Extract all score values from a nested dictionary or list structure.174 175 Args:176 data (Union[Dict, List]): The data structure to extract scores from.177 178 Returns:179 List[int]: List of extracted score values.180 """181 scores = []182 if isinstance(data, dict):183 for key, value in data.items():184 if "chunks" in key:185 continue186 elif isinstance(value, dict) or isinstance(value, list):187 scores.extend(extract_scores(value))188 elif key == 'score':189 scores.append(value)190 elif isinstance(data, list):191 for item in data:192 scores.extend(extract_scores(item))193 return scores194 195 196def calculate_overall_score(result: Dict) -> float:197 """198 Calculate the overall score from evaluation results.199 200 Args:201 result (Dict): Dictionary containing evaluation results.202 203 Returns:204 float: The calculated overall score.205 """206 scores = extract_scores(result)207 overall_score = calculate_geometric_mean(scores)208 return overall_score209 210 211def process_topic_name(topic_name: str) -> str:212 """213 Process a topic name by capitalizing words and handling special characters.214 215 Args:216 topic_name (str): The topic name to process.217 218 Returns:219 str: The processed topic name.220 """221 words = topic_name.replace("_s_", "'s_").split("_")222 return " ".join([word.capitalize() for word in words])223 224 225def merge_dicts(dict1: dict, dict2: dict) -> dict:226 """227 Recursively merge two dictionaries.228 229 Args:230 dict1 (dict): First dictionary.231 dict2 (dict): Second dictionary.232 233 Returns:234 dict: Merged dictionary.235 """236 merged = dict1.copy()237 for key, value in dict2.items():238 if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):239 merged[key] = merge_dicts(merged[key], value)240 else:241 merged[key] = value242 return merged243 244 245def process_theorem(models, file_path: str, eval_type: str, retry_limit: int,246 target_fps: int = None, use_parent_folder_as_topic: bool = False,247 output_folder: str = None) -> tuple[str, dict]:248 """249 Process a theorem file or directory for evaluation.250 251 Args:252 models: Dictionary of models for different evaluation types.253 file_path (str): Path to the file or directory to evaluate.254 eval_type (str): Type of evaluation to perform.255 retry_limit (int): Number of retry attempts.256 target_fps (int, optional): Target frames per second for video processing.257 use_parent_folder_as_topic (bool, optional): Use parent folder name as topic.258 output_folder (str, optional): Directory to store output files.259 260 Returns:261 tuple[str, dict]: Tuple of file name and evaluation results.262 """263 ext_map = {264 'text': ('.txt', '.srt'),265 'video': ('.mp4', '.mkv')266 }267 268 # Handle single file evaluation269 if os.path.isfile(file_path):270 file_ext = os.path.splitext(file_path)[1].lower()271 file_name = os.path.basename(file_path)272 273 if eval_type == "text" and file_ext in ext_map['text']:274 return file_name, evaluate_text_file(models['text'], file_path, retry_limit)275 elif eval_type == "video" and file_ext in ext_map['video']:276 if use_parent_folder_as_topic:277 topic_name = os.path.basename(os.path.dirname(file_path))278 else:279 topic_name = None280 topic_name = process_topic_name(topic_name)281 return file_name, evaluate_video_file(models['video'], file_path, None, topic_name, target_fps, output_folder)282 elif eval_type == "image" and file_ext in ext_map['video']:283 if use_parent_folder_as_topic:284 topic_name = os.path.basename(os.path.dirname(file_path))285 else:286 topic_name = None287 topic_name = process_topic_name(topic_name)288 return file_name, evaluate_sampled_images(models['image'], file_path, topic_name, num_chunks=10, output_folder=output_folder)289 elif eval_type == "all":290 raise ValueError("Evaluation type 'all' is not supported for a single file. Try passing a folder with both a video and a subtitle file.")291 else:292 raise ValueError(f"File type of {file_path} does not match evaluation type {eval_type!r}")293 294 # Handle directory evaluation295 theorem_dir = file_path296 all_files = os.listdir(theorem_dir)297 298 # Look for transcript files, prioritizing .srt over .txt if both exist299 transcript_file_candidates = [f for f in all_files if f.endswith(ext_map['text']) and not f.endswith('_scene_outline.txt')]300 srt_files = [f for f in transcript_file_candidates if f.endswith('.srt')]301 txt_files = [f for f in transcript_file_candidates if f.endswith('.txt')]302 303 transcript_path = None304 if srt_files:305 transcript_path = os.path.join(theorem_dir, srt_files[0])306 elif txt_files:307 transcript_path = os.path.join(theorem_dir, txt_files[0])308 309 video_file_candidates = [f for f in all_files if f.endswith(ext_map['video'])]310 video_path = os.path.join(theorem_dir, video_file_candidates[0]) if len(video_file_candidates) == 1 else None311 312 topic_name = os.path.basename(theorem_dir)313 topic_name = process_topic_name(topic_name)314 315 if not video_path:316 print(f"Skipping {theorem_dir}: No video file found")317 return None, None318 319 text_result = video_result = image_result = None320 if eval_type == "text" or eval_type == "all":321 if transcript_path is None:322 print(f"Warning: No suitable transcript file found in {theorem_dir}")323 else:324 text_result = evaluate_text_file(models['text'], transcript_path, retry_limit)325 if eval_type == "video" or eval_type == "all":326 assert video_path is not None, f"Expected 1 video file, got {len(video_file_candidates)} for {theorem_dir}"327 video_result = evaluate_video_file(models['video'], video_path, transcript_path, topic_name, target_fps, output_folder)328 if eval_type == "image" or eval_type == "all":329 assert video_path is not None, f"Expected 1 video file, got {len(video_file_candidates)} for {theorem_dir}"330 image_result = evaluate_sampled_images(models['image'], video_path, topic_name, num_chunks=10, output_folder=output_folder)331 332 if eval_type == "all":333 result = {}334 if text_result:335 result = merge_dicts(result, text_result)336 if video_result:337 result = merge_dicts(result, video_result)338 if image_result:339 result = merge_dicts(result, image_result)340 if result:341 result["evaluation"]["overall_score"] = calculate_overall_score(result)342 else:343 result = text_result if eval_type == "text" else video_result if eval_type == "video" else image_result if eval_type == "image" else None344 345 file_name = os.path.basename(theorem_dir)346 return file_name, result347 348 349def main():350 """351 Main function to run the evaluation script.352 353 Parses command line arguments and orchestrates the evaluation process354 for text, video, and image content using specified AI models.355 """356 parser = argparse.ArgumentParser(description='Automatic evaluation of theorem explanation videos with LLMs')357 parser.add_argument('--model_text', type=str, 358 choices=ALLOWED_MODELS,359 default='azure/gpt-4o',360 help='Select the AI model to use for text evaluation')361 parser.add_argument('--model_video', type=str,362 choices=['gemini/gemini-1.5-pro-002',363 'gemini/gemini-2.0-flash-exp',364 'gemini/gemini-2.0-pro-exp-02-05'],365 default='gemini/gemini-1.5-pro-002',366 help='Select the AI model to use for video evaluation')367 parser.add_argument('--model_image', type=str,368 choices=ALLOWED_MODELS,369 default='azure/gpt-4o',370 help='Select the AI model to use for image evaluation')371 parser.add_argument('--eval_type', type=str, choices=['text', 'video', 'image', 'all'], default='all', help='Type of evaluation to perform')372 parser.add_argument('--file_path', type=str, help='Path to a file or a theorem folder', required=True)373 parser.add_argument('--output_folder', type=str, help='Directory to store the evaluation files', required=True)374 parser.add_argument('--retry_limit', type=int, default=3, help='Number of retry attempts for each inference')375 parser.add_argument('--combine', action='store_true', help='Combine all results into a single JSON file')376 parser.add_argument('--bulk_evaluate', action='store_true', help='Evaluate a folder of theorems together', default=False)377 parser.add_argument('--target_fps', type=int, help='Target FPS for video processing. If not set, original video FPS will be used', required=False)378 parser.add_argument('--use_parent_folder_as_topic', action='store_true', help='Use parent folder name as topic name for single file evaluation', default=True)379 parser.add_argument('--max_workers', type=int, default=4, help='Maximum number of concurrent workers for parallel processing')380 381 args = parser.parse_args()382 383 # Initialize separate models384 text_model = LiteLLMWrapper(385 model_name=args.model_text,386 temperature=0.0,387 )388 video_model = GeminiWrapper(389 model_name=args.model_video,390 temperature=0.0,391 )392 image_model = LiteLLMWrapper(393 model_name=args.model_image,394 temperature=0.0,395 )396 397 models = {398 'text': text_model,399 'video': video_model,400 'image': image_model401 }402 403 theorem_dirs = []404 if args.bulk_evaluate:405 assert os.path.isdir(args.file_path), "File path must be a folder for --bulk_evaluate"406 for root, dirnames, _ in os.walk(args.file_path):407 if not any(f.endswith(".mp4") for f in os.listdir(root)):408 continue409 410 theorem_dirs.append(root)411 elif os.path.isdir(args.file_path):412 assert any(f.endswith(".mp4") for f in os.listdir(args.file_path)), "The provided folder must contain a video file"413 414 theorem_dirs.append(args.file_path)415 416 # Create output directory and its temp subdirectories if it doesn't exist417 os.makedirs(args.output_folder, exist_ok=True)418 moviepy_temp_dir = os.path.join(args.output_folder, "moviepy_temp")419 os.makedirs(moviepy_temp_dir, exist_ok=True)420 VideoFileClip.DEFAULT_TEMP_DIR = moviepy_temp_dir421 422 processed_videos_dir = os.path.join(args.output_folder, "processed_videos")423 os.makedirs(processed_videos_dir, exist_ok=True)424 425 results = {}426 if theorem_dirs:427 for theorem_dir in theorem_dirs:428 file_name, result = process_theorem(429 models,430 theorem_dir,431 args.eval_type,432 args.retry_limit,433 args.target_fps,434 args.use_parent_folder_as_topic,435 args.output_folder436 )437 438 if result is not None:439 results[file_name] = result440 441 if not args.combine:442 save_individual_result(args.output_folder, file_name, result)443 else:444 file_name, result = process_theorem(445 models, 446 args.file_path, 447 args.eval_type, 448 args.retry_limit,449 args.target_fps,450 args.use_parent_folder_as_topic,451 args.output_folder452 )453 454 if result is not None:455 results[file_name] = result456 457 if not args.combine:458 save_individual_result(args.output_folder, file_name, result)459 460 if args.combine:461 if len(results) > 1:462 current_time = datetime.now().strftime("%Y%m%d_%H%M%S")463 combined_file = f"evaluation_{current_time}.json"464 combine_results(args.output_folder, combined_file, results)465 print("Combining results completed.")466 else:467 for file_name, result in results.items():468 save_individual_result(args.output_folder, file_name, result)469 470 os.rmdir(moviepy_temp_dir)471 472 473if __name__ == "__main__":474 main()475