#!/usr/bin/env python3
# Copyright 2025 NVIDIA CORPORATION & AFFILIATES
# Modified from chengyue_script/eval_dfm_qwen.py for Qwen2.5-VL evaluation

"""
LM Eval Harness wrapper for Qwen2.5-VL-3B-Instruct (text-only evaluation).

Usage:
accelerate launch shiyi_scripts/eval_qwen2_5_vl_3B_lm_eval.py \
    --num_fewshot 0 \
    --tasks gsm8k \
    --batch_size 16 \
    --model qwen2_5_vl \
    --model_args model_path=Qwen/Qwen2.5-VL-3B-Instruct
"""

import accelerate
import torch
import random
import numpy as np
from lm_eval.__main__ import cli_evaluate
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
from tqdm import tqdm
import os
from transformers import AutoTokenizer, AutoModel
import json
import time


def set_seed(seed):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


@register_model("qwen2_5_vl")
class Qwen2_5_VLEvalHarness(LM):
    """
    LM Eval Harness wrapper for Qwen2.5-VL (text-only generation).
    """
    def __init__(
        self,
        model_path='Qwen/Qwen2.5-VL-3B-Instruct',
        device="cuda",
        show_speed=False,
        save_dir=None,
        max_new_tokens=512,
        temperature=0.0,
        do_sample=False,
        **kwargs,
    ):
        super().__init__()

        accelerator = accelerate.Accelerator()
        if accelerator.num_processes > 1:
            self.accelerator = accelerator
        else:
            self.accelerator = None
        
        model_kwargs = {}
        if self.accelerator is not None:
            model_kwargs.update({'device_map': {'': f'{self.accelerator.device}'}})
        
        # Use AutoModel (handles VL models correctly)
        print(f"Loading Qwen2.5-VL from: {model_path}")
        self.model = AutoModel.from_pretrained(
            model_path, 
            trust_remote_code=True, 
            torch_dtype=torch.bfloat16, 
            **model_kwargs
        )
        self.model.eval()
        
        self.device = torch.device(device)
        if self.accelerator is not None:
            self.model = self.accelerator.prepare(self.model)
            self.device = torch.device(f'{self.accelerator.device}')
            self._rank = self.accelerator.local_process_index
            self._world_size = self.accelerator.num_processes
        else: 
            self.model = self.model.to(device)
            self._rank = 0
            self._world_size = 1

        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self._tokenizer_name = model_path  # Required by lm_eval for chat templates
        
        self.show_speed = show_speed
        self.save_dir = save_dir
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.do_sample = do_sample
        
        print(f"✓ Model loaded: {type(self.model).__name__}")
        print(f"  Rank: {self._rank}/{self._world_size}")
        
    @property
    def rank(self):
        return self._rank
    
    @property
    def world_size(self):
        return self._world_size
    
    @property
    def tokenizer_name(self):
        """Required by lm_eval for chat template support."""
        return self._tokenizer_name

    def loglikelihood(self, requests):
        """Not implemented - use generate_until for accuracy evaluation."""
        raise NotImplementedError("Use generate_until for generation-based eval")
    
    def loglikelihood_rolling(self, requests):
        """Not implemented - use generate_until for accuracy evaluation."""
        raise NotImplementedError("Use generate_until for generation-based eval")
    
    def generate_until(self, requests):
        """
        Generate answers for evaluation (text-only).
        """
        output = []
        num_tokens = 0
        processed_count = 0
        
        if self.save_dir is not None:
            os.makedirs(self.save_dir, exist_ok=True)
            rank = self.rank
            save_path = os.path.join(self.save_dir, f'rank_{rank}.jsonl')
            print(f"Save path: {save_path}")
            if os.path.exists(save_path):
                print(f"Loading from {save_path}")
                with open(save_path, 'r', encoding='utf-8') as f:
                    output = [json.loads(line) for line in f]
                    processed_count = len(output)
                print(f"Processed count: {processed_count}")
        
        start_time = time.time()
        
        for i, req in enumerate(tqdm(requests, desc="Generating...", disable=self.rank != 0)):
            if i < processed_count:
                continue
                
            question = req.args[0]
            stop_tokens = req.args[1]['until'] if len(req.args) > 1 else []
            
            # Apply chat template
            messages = [{"role": "user", "content": question}]
            user_input = self.tokenizer.apply_chat_template(
                messages, 
                add_generation_prompt=True, 
                tokenize=False
            )
            
            # Tokenize input
            model_inputs = self.tokenizer([user_input], return_tensors="pt").to(self.device)
            
            # Generate (text-only, no vision inputs)
            with torch.no_grad():
                if self.accelerator is not None:
                    generated_ids = self.accelerator.unwrap_model(self.model).generate(
                        **model_inputs,
                        max_new_tokens=self.max_new_tokens,
                        temperature=self.temperature,
                        do_sample=self.do_sample,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        use_cache=True
                    )
                else:
                    generated_ids = self.model.generate(
                        **model_inputs,
                        max_new_tokens=self.max_new_tokens,
                        temperature=self.temperature,
                        do_sample=self.do_sample,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        use_cache=True
                    )
            
            # Extract generated tokens only
            generated_ids = [
                output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
            ]
            
            # Decode
            generated_answer = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
            
            # Count tokens
            if self.show_speed:
                num_tokens += len(generated_ids[0])
            
            output.append(generated_answer)

            if self.save_dir is not None:
                with open(save_path, 'a', encoding='utf-8') as f:
                    f.write(json.dumps(generated_answer, ensure_ascii=False) + '\n')

            if self.rank == 0 and i % 10 == 0:
                print('=' * 20)
                print('Question:', question[:100])
                print('Answer:', generated_answer[:100])
                print('=' * 20)
            
        end_time = time.time()
        if self.show_speed and self.rank == 0:
            print(f"Total tokens generated: {num_tokens}")
            print(f"Total time: {end_time - start_time:.2f}s")
            print(f"Tokens/sec: {num_tokens / (end_time - start_time):.2f}")
            
        return output


if __name__ == "__main__":
    cli_evaluate()

