import torch from PIL import Image from accelerate import Accelerator from datasets import load_dataset from torch.distributed import all_gather_object from transformers import AutoProcessor, AutoConfig, AutoTokenizer, LlavaOnevisionForConditionalGeneration from trl.models import unwrap_model_for_generation from data_utils.aokvqa.evaluator import eval_aokvqa_direct from reward_utils.compute_rewards import split_initial_context accelerator = Accelerator() from tqdm import tqdm import numpy as np DEVICE = accelerator.device # Model and Processor Configuration model_args = {} # Use {"torch_dtype":torch.bfloat16} if desired and supported model_id = '/path/to/dyme-aok-local/final_checkpoint' config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained(model_id, config=config, trust_remote_code=True) model = LlavaOnevisionForConditionalGeneration.from_pretrained( model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ).to(DEVICE) model.eval() processor = AutoProcessor.from_pretrained(model_id) # Configure image processor size if hasattr(processor.image_processor, 'size') and isinstance(processor.image_processor.size, dict): processor.tokenizer.padding_side = 'left' else: print( f"Warning: Could not directly set 'longest_edge' via dict. Current image processor size config: {processor.image_processor.size}") PROMPT_TEMPLATE = ( "{question} Answer the question with a single word (or phrase)." ) def run_kh_batch(batch_data_list): # Renamed from run_kh, takes a batch batch_images = [] batch_formatted_prompts_for_chat_template = [] for item in batch_data_list: image_path = item['image_path'] item_model_input_text = item['model_input_text'].strip() question_with_tags = PROMPT_TEMPLATE.format(question=item_model_input_text) if isinstance(image_path, str): image = Image.open(image_path).convert("RGB") else: image = image_path.convert("RGB") # Assuming image_path is already a PIL Image object batch_images.append(image) messages = [ { "role": "user", "content": [ {"type": "image"}, {"type": "text", "text": question_with_tags}, ] }, ] try: templated_prompt_str = processor.apply_chat_template(messages, add_generation_prompt=True) templated_prompt_str = templated_prompt_str.strip() except: templated_prompt_str = f"USER: \n{question_with_tags}\nASSISTANT:" batch_formatted_prompts_for_chat_template.append(templated_prompt_str) inputs = processor( text=batch_formatted_prompts_for_chat_template, images=batch_images, return_tensors="pt", padding=True, truncation=True ) inputs = { k: v.to(DEVICE).to(torch.bfloat16) if v.is_floating_point() else v.to( DEVICE) for k, v in inputs.items() } with unwrap_model_for_generation(model, accelerator) as unwrapped_model_instance: generated_ids = unwrapped_model_instance.generate(**inputs, max_new_tokens=1024, do_sample=False, ) input_ids_length = inputs['input_ids'].shape[1] newly_generated_ids = generated_ids[:, input_ids_length:] generated_texts = processor.batch_decode( newly_generated_ids, skip_special_tokens=True, # Special tokens like are removed. might be too. ) return [text.strip('.').strip() for text in generated_texts] task = 'aokvqa' dt_record_local = {} if task == 'aokvqa': if accelerator.is_main_process: print("Loading A-OKVQA dataset...") try: full_dataset = load_dataset("HuggingFaceM4/A-OKVQA", trust_remote_code=True)['validation'] except Exception as e: if accelerator.is_main_process: print(f"Failed to load dataset directly. Error: {e}") print("Attempting to load with specific revision if applicable, or check path/connection.") raise eval_datasets_all_prepared = [] for d_item in tqdm(full_dataset, desc="Preparing dataset", disable=not accelerator.is_main_process): image_path = d_item['image'] # --- 修改：'query' -> 'question' --- raw_question = d_item['question'] # --- 修改：'label' -> 'direct_answers' --- ground_truth_answers = d_item.get('direct_answers') if not ground_truth_answers: if accelerator.is_main_process: tqdm.write( f"Warning: Item missing 'direct_answers' or 'direct_answers' is empty. Question: {raw_question[:50]}...") continue model_input_text_for_template = raw_question eval_datasets_all_prepared.append({ 'image_path': image_path, 'model_input_text': model_input_text_for_template, 'direct_answers_list': ground_truth_answers, 'original_question': raw_question }) num_processes = accelerator.num_processes process_index = accelerator.process_index total_items = len(eval_datasets_all_prepared) if total_items == 0: if accelerator.is_main_process: print("No data prepared for evaluation after filtering. Exiting A-OKVQA evaluation.") else: items_per_proc = total_items // num_processes extra_items = total_items % num_processes local_start_index = process_index * items_per_proc + min(process_index, extra_items) num_local_items = items_per_proc + (1 if process_index < extra_items else 0) local_end_index = local_start_index + num_local_items eval_datasets_local = eval_datasets_all_prepared[local_start_index:local_end_index] BATCH_SIZE = 32 REPORT_INTERVAL_BATCHES = 1 pbar = None if accelerator.is_main_process and len(eval_datasets_local) > 0: pbar = tqdm(total=len(eval_datasets_local), desc=f"Eval Proc {process_index}", dynamic_ncols=True) dt_record_local['res'] = [] num_local_batches = (len(eval_datasets_local) + BATCH_SIZE - 1) // BATCH_SIZE for batch_idx_local in range(num_local_batches): start_idx = batch_idx_local * BATCH_SIZE end_idx = min((batch_idx_local + 1) * BATCH_SIZE, len(eval_datasets_local)) current_batch_list = eval_datasets_local[start_idx:end_idx] if not current_batch_list: continue batch_predictions_texts = run_kh_batch(current_batch_list) for item_idx_in_batch, full_pred_text in enumerate(batch_predictions_texts): original_item = current_batch_list[item_idx_in_batch] ground_truth_answers_list = eval(original_item['direct_answers_list']) _, parsed_pred_answer = split_initial_context(full_pred_text) if not parsed_pred_answer.strip(): parsed_pred_answer = full_pred_text score = eval_aokvqa_direct(parsed_pred_answer, ground_truth_answers_list) dt_record_local['res'].append(score) if accelerator.is_main_process: print(parsed_pred_answer, "######", ground_truth_answers_list, "######", score) if pbar: pbar.update(len(current_batch_list)) is_last_local_batch = (batch_idx_local == num_local_batches - 1) should_sync_and_report = ((batch_idx_local + 1) % REPORT_INTERVAL_BATCHES == 0) or is_last_local_batch if len(eval_datasets_local) == 0: should_sync_and_report = False if num_local_batches == 0 and is_last_local_batch: should_sync_and_report = True if should_sync_and_report: accelerator.wait_for_everyone() gathered_all_processes_data = [None] * num_processes all_gather_object(gathered_all_processes_data, dt_record_local) if accelerator.is_main_process: current_global_scores_list = [] for process_data_dict in gathered_all_processes_data: if process_data_dict and 'res' in process_data_dict: current_global_scores_list.extend(process_data_dict['res']) total_samples_processed_globally = len(current_global_scores_list) report_title = "--- Intermediate Report ---" if is_last_local_batch and total_samples_processed_globally == total_items: report_title = "--- Final Report ---" elif is_last_local_batch: report_title = f"--- Report (Main Proc Last Batch, {batch_idx_local + 1}/{num_local_batches}) ---" tqdm.write(f"\n{report_title}") if current_global_scores_list: mean_acc_global = np.array(current_global_scores_list).mean() if accelerator.is_main_process: print(f"Global samples processed: {total_samples_processed_globally} / {total_items}") print(f"Current Global Mean Accuracy (VQA Acc): {mean_acc_global:.4f}") # 标签更新为 VQA Acc if pbar: pbar.set_description( f"Global Acc: {mean_acc_global:.4f} ({total_samples_processed_globally}/{total_items})") else: if accelerator.is_main_process: print( f"No scores to report globally yet (Total processed: {total_samples_processed_globally}).") accelerator.wait_for_everyone() if pbar: pbar.close() if accelerator.is_main_process and len(eval_datasets_local) == 0 and total_items > 0: print( f"Main process had no data, but other processes might have. Final global metrics are printed by the last reporting sync.") elif accelerator.is_main_process and total_items == 0: print("No data was prepared for evaluation. Nothing to report.") else: if accelerator.is_main_process: print(f"Task '{task}' is not configured for batched evaluation in this script.")