chore: import upstream snapshot with attribution

tests/utils/ocr_eval.py

@@ -0,0 +1,326 @@
+"""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()