unslothai--unsloth/tests/utils/ocr_eval.py

"""Evaluate OCR models on datasets with WER and CER metrics."""

import os
import torch
from tqdm import tqdm
import pandas as pd
from jiwer import wer, cer
from qwen_vl_utils import process_vision_info
import matplotlib.pyplot as plt
from typing import List, Dict, Tuple, Optional, Any
import traceback


class OCRModelEvaluator:
    """OCR model evaluator over multiple models with WER/CER analysis."""

    def __init__(self):
        """Initialize the OCR evaluator."""
        self.model_comparison_results = {}

    def evaluate_model(
        self,
        model: Any,
        processor: Any,
        dataset: List[Dict],
        output_dir: str = "ocr_evaluation_results",
        max_new_tokens: int = 1024,
        temperature: float = 1.5,
        min_p: float = 0.1,
        verbose: bool = True,
    ) -> Tuple[Optional[float], Optional[float]]:
        """Evaluate a model on an OCR dataset."""
        os.makedirs(output_dir, exist_ok = True)

        results = []

        for i, sample in enumerate(
            tqdm(dataset, desc = "Evaluating OCR performance", disable = not verbose)
        ):
            try:
                messages = sample["messages"]

                ground_truth, image, question, input_messages = self._extract_sample_components(
                    messages, i, verbose
                )

                if ground_truth is None or image is None or question is None:
                    continue

                generated_response = self._generate_response(
                    model, processor, input_messages, max_new_tokens, temperature, min_p
                )

                word_error = wer(ground_truth, generated_response)
                char_error = cer(ground_truth, generated_response)

                self._save_individual_result(
                    output_dir,
                    i,
                    question,
                    generated_response,
                    ground_truth,
                    word_error,
                    char_error,
                )

                results.append(
                    {
                        "sample_id": i,
                        "wer": word_error,
                        "cer": char_error,
                        "model_output": generated_response.strip(),
                        "ground_truth": ground_truth,
                        "question": question,
                    }
                )

            except Exception as e:
                if verbose:
                    print(f"Error processing sample {i}: {str(e)}")
                    traceback.print_exc()

        return self._generate_summary_report(results, output_dir, verbose)

    def _extract_sample_components(
        self, messages: List[Dict], sample_idx: int, verbose: bool
    ) -> Tuple[Optional[str], Optional[Any], Optional[str], List[Dict]]:
        """Extract ground truth, image, question, and input messages from sample."""

        system_message = next((msg for msg in messages if msg["role"] == "system"), None)

        user_message = next((msg for msg in messages if msg["role"] == "user"), None)
        if not user_message:
            if verbose:
                print(f"Skipping sample {sample_idx}: No user message found")
            return None, None, None, []

        assistant_message = next((msg for msg in messages if msg["role"] == "assistant"), None)
        if not assistant_message:
            if verbose:
                print(f"Skipping sample {sample_idx}: No assistant message (ground truth) found")
            return None, None, None, []

        ground_truth = None
        for content_item in assistant_message["content"]:
            if content_item["type"] == "text":
                ground_truth = content_item["text"]
                break

        if not ground_truth:
            if verbose:
                print(f"Skipping sample {sample_idx}: No text found in assistant message")
            return None, None, None, []

        # Extract image and question from user message
        image = None
        question = None
        for content_item in user_message["content"]:
            if content_item["type"] == "image":
                image = content_item["image"]
            elif content_item["type"] == "text":
                question = content_item["text"]

        if not image:
            if verbose:
                print(f"Skipping sample {sample_idx}: No image found in user message")
            return None, None, None, []

        if not question:
            if verbose:
                print(f"Skipping sample {sample_idx}: No question found in user message")
            return None, None, None, []

        # Model input excludes the assistant message
        input_messages = []
        if system_message:
            input_messages.append(system_message)
        input_messages.append(user_message)

        return ground_truth, image, question, input_messages

    def _generate_response(
        self,
        model: Any,
        processor: Any,
        input_messages: List[Dict],
        max_new_tokens: int,
        temperature: float,
        min_p: float,
    ) -> str:
        """Generate response from the model."""

        text = processor.apply_chat_template(
            input_messages, tokenize = False, add_generation_prompt = True
        )

        image_inputs, video_inputs = process_vision_info(input_messages)

        inputs = processor(
            text = [text],
            images = image_inputs,
            videos = video_inputs,
            padding = True,
            return_tensors = "pt",
        )
        inputs = inputs.to(model.device)

        with torch.no_grad():
            generated_ids = model.generate(
                **inputs,
                max_new_tokens = max_new_tokens,
                temperature = temperature,
                min_p = min_p,
                use_cache = True,
            )

        # Keep only the generated tokens, not the input
        generated_ids_trimmed = [
            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]

        generated_response = processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens = True,
            clean_up_tokenization_spaces = False,
        )[0]

        return generated_response

    def _save_individual_result(
        self,
        output_dir: str,
        sample_idx: int,
        question: str,
        generated_response: str,
        ground_truth: str,
        word_error: float,
        char_error: float,
    ):
        """Save individual sample result to file."""
        output_file = os.path.join(output_dir, f"sample_{sample_idx}.txt")
        with open(output_file, "w", encoding = "utf-8") as f:
            f.write(f"Sample {sample_idx}\n")
            f.write(f"Question: {question}\n\n")
            f.write(f"Model output:\n{generated_response.strip()}\n\n")
            f.write(f"Ground truth:\n{ground_truth}\n\n")
            f.write(f"WER: {word_error:.4f}, CER: {char_error:.4f}")

    def _generate_summary_report(
        self, results: List[Dict], output_dir: str, verbose: bool
    ) -> Tuple[Optional[float], Optional[float]]:
        """Generate and save summary report."""
        if not results:
            if verbose:
                print("No results to summarize.")
            return None, None

        df = pd.DataFrame(results)

        avg_wer = df["wer"].mean()
        avg_cer = df["cer"].mean()

        with open(os.path.join(output_dir, "avg_metrics.txt"), "w") as f:
            f.write(f"Average WER: {avg_wer:.4f}\n")
            f.write(f"Average CER: {avg_cer:.4f}\n")

        df.to_csv(os.path.join(output_dir, "detailed_results.csv"), index = False)

        if verbose:
            print("\nResults Summary:")
            print(f"Average WER: {avg_wer:.4f}")
            print(f"Average CER: {avg_cer:.4f}")
            print(f"\nDetailed results saved to {output_dir}/")

        return avg_wer, avg_cer

    def add_to_comparison(self, model_name: str, wer: float, cer: float):
        """Add model results to the comparison tracker."""
        self.model_comparison_results[model_name] = {"wer": wer, "cer": cer}

    def print_model_comparison(
        self,
        save_csv: bool = True,
        save_plot: bool = True,
    ) -> Optional[pd.DataFrame]:
        """Print a comparison of all models evaluated so far."""
        if not self.model_comparison_results:
            print("No model results available for comparison")
            return None

        print("\n==== MODEL COMPARISON REPORT ====")

        comparison_df = pd.DataFrame(
            {
                "Model": list(self.model_comparison_results.keys()),
                "WER": [results["wer"] for results in self.model_comparison_results.values()],
                "CER": [results["cer"] for results in self.model_comparison_results.values()],
            }
        )

        # Sort by WER (best first)
        comparison_df = comparison_df.sort_values("WER")

        print("\nComparison Table (sorted by WER):")
        print(comparison_df.to_string(index = False))

        if save_csv:
            comparison_file = "model_comparison_results.csv"
            comparison_df.to_csv(comparison_file, index = False)
            print(f"\nComparison table saved to {comparison_file}")

        if save_plot:
            self._create_comparison_plot(comparison_df)

        return comparison_df

    def _create_comparison_plot(self, comparison_df: pd.DataFrame):
        """Create and save comparison plot."""
        plt.figure(figsize = (12, 6))

        # Plot WER
        plt.subplot(1, 2, 1)
        plt.bar(comparison_df["Model"], comparison_df["WER"], color = "skyblue")
        plt.title("Word Error Rate Comparison")
        plt.ylabel("WER (lower is better)")
        plt.ylim(bottom = 0)
        plt.xticks(rotation = 45, ha = "right")

        # Plot CER
        plt.subplot(1, 2, 2)
        plt.bar(comparison_df["Model"], comparison_df["CER"], color = "lightgreen")
        plt.title("Character Error Rate Comparison")
        plt.ylabel("CER (lower is better)")
        plt.ylim(bottom = 0)
        plt.xticks(rotation = 45, ha = "right")

        plt.tight_layout()
        plt.savefig("ocr_model_comparison.png")
        plt.show()

        print(f"\nVisualization saved to ocr_model_comparison.png")

    def get_comparison_results(self) -> Dict[str, Dict[str, float]]:
        """Get the current comparison results."""
        return self.model_comparison_results.copy()

    def clear_comparison_results(self):
        """Clear all comparison results."""
        self.model_comparison_results.clear()


def evaluate_ocr_model(
    model,
    processor,
    dataset,
    output_dir = "ocr_evaluation_results",
    **kwargs,
):
    """Convenience wrapper kept for backward compatibility."""
    evaluator = OCRModelEvaluator()
    return evaluator.evaluate_model(model, processor, dataset, output_dir, **kwargs)


def create_evaluator():
    """Create a new OCR evaluator instance."""
    return OCRModelEvaluator()