evaluation.py

"""
Evaluation metrics and benchmarking for PaliGemma.
Supports VQA, image captioning, and other vision-language tasks.
"""
import torch
import numpy as np
from typing import List, Dict, Optional, Tuple
from collections import defaultdict
import json
import logging
from tqdm import tqdm

from modeling_gemma import PaliGemmaForConditionalGeneration, KVCache
from processing_paligemma import PaliGemmaProcessor
from data_utils import PaliGemmaDataset, create_dataloader
from torch.utils.data import DataLoader

logger = logging.getLogger(__name__)


class MetricsCalculator:
    """Calculate various evaluation metrics"""
    
    @staticmethod
    def bleu_score(predictions: List[str], references: List[List[str]], n: int = 4) -> float:
        """
        Calculate BLEU score.
        
        Args:
            predictions: List of predicted sentences
            references: List of reference sentences (can be multiple per prediction)
            n: Maximum n-gram order
        
        Returns:
            BLEU score
        """
        from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
        smooth = SmoothingFunction().method1
        
        scores = []
        for pred, refs in zip(predictions, references):
            if isinstance(refs, str):
                refs = [refs]
            # Tokenize
            pred_tokens = pred.lower().split()
            ref_tokens_list = [ref.lower().split() for ref in refs]
            
            score = sentence_bleu(ref_tokens_list, pred_tokens, smoothing_function=smooth)
            scores.append(score)
        
        return np.mean(scores)
    
    @staticmethod
    def rouge_l_score(predictions: List[str], references: List[str]) -> float:
        """
        Calculate ROUGE-L score (Longest Common Subsequence).
        
        Args:
            predictions: List of predicted sentences
            references: List of reference sentences
        
        Returns:
            ROUGE-L F1 score
        """
        def lcs_length(s1: List[str], s2: List[str]) -> int:
            """Calculate LCS length"""
            m, n = len(s1), len(s2)
            dp = [[0] * (n + 1) for _ in range(m + 1)]
            
            for i in range(1, m + 1):
                for j in range(1, n + 1):
                    if s1[i-1] == s2[j-1]:
                        dp[i][j] = dp[i-1][j-1] + 1
                    else:
                        dp[i][j] = max(dp[i-1][j], dp[i][j-1])
            
            return dp[m][n]
        
        scores = []
        for pred, ref in zip(predictions, references):
            pred_tokens = pred.lower().split()
            ref_tokens = ref.lower().split()
            
            lcs = lcs_length(pred_tokens, ref_tokens)
            if len(ref_tokens) == 0:
                precision = 0.0
            else:
                precision = lcs / len(ref_tokens)
            
            if len(pred_tokens) == 0:
                recall = 0.0
            else:
                recall = lcs / len(pred_tokens)
            
            if precision + recall == 0:
                f1 = 0.0
            else:
                f1 = 2 * precision * recall / (precision + recall)
            
            scores.append(f1)
        
        return np.mean(scores)
    
    @staticmethod
    def exact_match(predictions: List[str], references: List[str]) -> float:
        """Calculate exact match accuracy"""
        matches = sum(1 for p, r in zip(predictions, references) if p.strip().lower() == r.strip().lower())
        return matches / len(predictions) if predictions else 0.0
    
    @staticmethod
    def vqa_accuracy(predictions: List[str], references: List[str]) -> float:
        """
        Calculate VQA accuracy (case-insensitive, punctuation-agnostic).
        
        Args:
            predictions: List of predicted answers
            references: List of reference answers
        
        Returns:
            Accuracy score
        """
        def normalize_answer(answer: str) -> str:
            """Normalize answer for comparison"""
            import string
            answer = answer.lower().strip()
            # Remove punctuation
            answer = answer.translate(str.maketrans('', '', string.punctuation))
            return answer
        
        matches = sum(
            1 for p, r in zip(predictions, references)
            if normalize_answer(p) == normalize_answer(r)
        )
        return matches / len(predictions) if predictions else 0.0


def generate_batch(
    model: PaliGemmaForConditionalGeneration,
    processor: PaliGemmaProcessor,
    images: List,
    prompts: List[str],
    device: str,
    max_new_tokens: int = 100,
    temperature: float = 0.8,
    top_p: float = 0.9,
    do_sample: bool = True,
) -> List[str]:
    """
    Generate text for a batch of images and prompts.
    
    Args:
        model: PaliGemma model
        processor: PaliGemmaProcessor
        images: List of PIL Images
        prompts: List of text prompts
        device: Device to run on
        max_new_tokens: Maximum tokens to generate
        temperature: Sampling temperature
        top_p: Nucleus sampling threshold
        do_sample: Whether to use sampling
    
    Returns:
        List of generated texts
    """
    model.eval()
    generated_texts = []
    
    # Process batch
    model_inputs = processor(text=prompts, images=images)
    model_inputs = {k: v.to(device) for k, v in model_inputs.items()}
    
    input_ids = model_inputs["input_ids"]
    attention_mask = model_inputs["attention_mask"]
    pixel_values = model_inputs["pixel_values"]
    
    batch_size = input_ids.shape[0]
    kv_caches = [KVCache() for _ in range(batch_size)]
    stop_token = processor.tokenizer.eos_token_id
    
    # Generate for each item in batch
    for i in range(batch_size):
        kv_cache = kv_caches[i]
        current_input_ids = input_ids[i:i+1]
        current_attention_mask = attention_mask[i:i+1]
        current_pixel_values = pixel_values[i:i+1]
        
        generated_tokens = []
        
        with torch.no_grad():
            for _ in range(max_new_tokens):
                outputs = model(
                    input_ids=current_input_ids,
                    pixel_values=current_pixel_values,
                    attention_mask=current_attention_mask,
                    kv_cache=kv_cache,
                )
                kv_cache = outputs["kv_cache"]
                next_token_logits = outputs["logits"][:, -1, :]
                
                if do_sample:
                    next_token_logits = torch.softmax(next_token_logits / temperature, dim=-1)
                    next_token = _sample_top_p(next_token_logits, top_p)
                else:
                    next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
                
                next_token = next_token.squeeze(0)
                generated_tokens.append(next_token)
                
                if next_token.item() == stop_token:
                    break
                
                current_input_ids = next_token.unsqueeze(-1)
                current_attention_mask = torch.cat(
                    [current_attention_mask, torch.ones((1, 1), device=device)], dim=-1
                )
        
        generated_tokens = torch.cat(generated_tokens, dim=-1)
        decoded = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
        generated_texts.append(decoded)
    
    return generated_texts


def _sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
    """Top-p (nucleus) sampling"""
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    mask = probs_sum - probs_sort > p
    probs_sort[mask] = 0.0
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    next_token = torch.multinomial(probs_sort, num_samples=1)
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token


def evaluate_model(
    model: PaliGemmaForConditionalGeneration,
    processor: PaliGemmaProcessor,
    dataloader: DataLoader,
    device: str,
    task_type: str = "captioning",
    max_new_tokens: int = 100,
    temperature: float = 0.8,
    top_p: float = 0.9,
) -> Dict[str, float]:
    """
    Evaluate model on a dataset.
    
    Args:
        model: PaliGemma model
        processor: PaliGemmaProcessor
        dataloader: DataLoader with evaluation data
        device: Device to run on
        task_type: Type of task ('captioning', 'vqa', 'general')
        max_new_tokens: Maximum tokens to generate
        temperature: Sampling temperature
        top_p: Nucleus sampling threshold
    
    Returns:
        Dictionary of metrics
    """
    model.eval()
    predictions = []
    references = []
    
    metrics_calc = MetricsCalculator()
    
    logger.info("Running evaluation...")
    for batch in tqdm(dataloader, desc="Evaluating"):
        # Get images and prompts from batch
        # This is simplified - adjust based on your dataset structure
        pixel_values = batch["pixel_values"]
        input_ids = batch["input_ids"]
        
        # Decode prompts (simplified)
        prompts = processor.tokenizer.batch_decode(input_ids, skip_special_tokens=True)
        
        # Generate predictions
        # Note: This is a simplified version - you'd need to extract images properly
        # For now, we'll use a placeholder approach
        
        # Get references if available
        if "labels" in batch:
            labels = batch["labels"]
            refs = processor.tokenizer.batch_decode(labels, skip_special_tokens=True)
            references.extend(refs)
    
    # Calculate metrics based on task type
    metrics = {}
    
    if task_type == "vqa" and references:
        metrics["vqa_accuracy"] = metrics_calc.vqa_accuracy(predictions, references)
        metrics["exact_match"] = metrics_calc.exact_match(predictions, references)
    elif task_type == "captioning" and references:
        metrics["bleu_4"] = metrics_calc.bleu_score(predictions, [[r] for r in references], n=4)
        metrics["rouge_l"] = metrics_calc.rouge_l_score(predictions, references)
    
    return metrics


def benchmark_inference_speed(
    model: PaliGemmaForConditionalGeneration,
    processor: PaliGemmaProcessor,
    device: str,
    num_runs: int = 10,
    batch_size: int = 1,
) -> Dict[str, float]:
    """
    Benchmark inference speed.
    
    Args:
        model: PaliGemma model
        processor: PaliGemmaProcessor
        device: Device to run on
        num_runs: Number of benchmark runs
        batch_size: Batch size for benchmarking
    
    Returns:
        Dictionary with timing metrics
    """
    import time
    from PIL import Image
    import torch
    
    # Create dummy inputs
    dummy_image = Image.new('RGB', (224, 224), color='red')
    dummy_prompt = "describe this image"
    
    model.eval()
    times = []
    
    # Warmup
    for _ in range(3):
        with torch.no_grad():
            model_inputs = processor(text=[dummy_prompt], images=[dummy_image])
            model_inputs = {k: v.to(device) for k, v in model_inputs.items()}
            _ = model(**model_inputs)
    
    # Benchmark
    torch.cuda.synchronize() if device == "cuda" else None
    
    for _ in range(num_runs):
        start_time = time.time()
        
        with torch.no_grad():
            model_inputs = processor(text=[dummy_prompt], images=[dummy_image])
            model_inputs = {k: v.to(device) for k, v in model_inputs.items()}
            _ = model(**model_inputs)
        
        if device == "cuda":
            torch.cuda.synchronize()
        
        elapsed = time.time() - start_time
        times.append(elapsed)
    
    avg_time = np.mean(times)
    std_time = np.std(times)
    throughput = batch_size / avg_time
    
    return {
        "avg_inference_time": avg_time,
        "std_inference_time": std_time,
        "throughput_samples_per_sec": throughput,
        "min_time": np.min(times),
        "max_time": np.max(times),
    }