code/OpenTSLM/evaluation/baseline/common_evaluator.py

# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT

import json
import os
import io
import re
import sys
import base64
from typing import Type, Callable, Dict, List, Any, Optional

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
from tqdm import tqdm
from transformers.pipelines import pipeline
import matplotlib.pyplot as plt
from time import sleep

from opentslm.logger import get_logger

from .openai_pipeline import OpenAIPipeline


class CommonEvaluator:
    """
    A common evaluation framework for testing LLMs on time series datasets.
    """

    def __init__(self, device: Optional[str] = None):
        """
        Initialize the evaluator.

        Args:
            device: Device to use for inference ('cuda', 'mps', 'cpu', or None for auto)
        """
        self.device = device or self._get_best_device()

        if self.device == "mps":
            print(
                "⚠️ Warning!! MPS is available but not recommended for evaluation. Many LLMs do not produce reasonable output!"
            )
            print(
                "⚠️ Warning!! MPS is available but not recommended for evaluation. Many LLMs do not produce reasonable output!"
            )
            print("⚠️ Better use CPU or CUDA for evaluation.")
            sleep(10)

    def _get_best_device(self) -> str:
        """Get the best available device."""
        if torch.cuda.is_available():
            return "cuda"
        elif torch.backends.mps.is_available():
            return "mps"
        else:
            return "cpu"

    def load_model(self, model_name: str, **pipeline_kwargs) -> pipeline:
        """
        Load a model using transformers pipeline or OpenAI API.
        """
        self.current_model_name = (
            model_name  # Track the current model name for formatter selection
        )
        if model_name.startswith("openai-"):
            # Use OpenAI API
            openai_model = model_name.replace("openai-", "")
            return OpenAIPipeline(model_name=openai_model, **pipeline_kwargs)
        print(f"Loading model: {model_name}")
        print(f"Using device: {self.device}")
        # Default pipeline arguments
        default_kwargs = {
            "task": "text-generation",
            "device": self.device,
            "temperature": 0.1,
        }
        default_kwargs.update(pipeline_kwargs)
        pipe = pipeline(model=model_name, **default_kwargs)
        print(f"Model loaded successfully: {model_name}")
        return pipe

    def load_dataset(
        self,
        dataset_class: Type[Dataset],
        split: str = "test",
        format_sample_str: bool = True,
        max_samples: Optional[int] = None,
        **dataset_kwargs,
    ) -> Dataset:
        """
        Load a dataset with proper formatting.
        """
        print(f"Loading dataset: {dataset_class.__name__}")

        # Import the gruver formatters
        from .gruver_llmtime_tokenizer import gpt_formatter, llama_formatter

        # Choose formatter based on model type
        model_name = getattr(self, "current_model_name", None)
        if model_name is None and "model_name" in dataset_kwargs:
            model_name = dataset_kwargs["model_name"]
        if model_name is not None:
            if model_name.startswith("openai-") or "gpt" in model_name.lower():
                formatter = gpt_formatter
                print(f"Using GPT formatter for model: {model_name}")
            elif "llama" in model_name.lower():
                formatter = llama_formatter
                print(f"Using Llama formatter for model: {model_name}")
            else:
                print(f"Defaulting to Llama formatter for model: {model_name}")

                formatter = llama_formatter
        else:
            formatter = llama_formatter

        # Default dataset arguments
        default_kwargs = {
            "split": split,
            "EOS_TOKEN": "",
            "format_sample_str": format_sample_str,
            "time_series_format_function": formatter,
        }

        # Add max_samples if provided
        if max_samples is not None:
            default_kwargs["max_samples"] = max_samples

        # Update with provided kwargs
        default_kwargs.update(dataset_kwargs)

        dataset = dataset_class(**default_kwargs)
        print(f"Loaded {len(dataset)} {split} samples")

        return dataset

    def evaluate_model_on_dataset(
        self,
        model_name: str,
        dataset_class: Type[Dataset],
        evaluation_function: Callable[[str, str], Dict[str, Any]],
        max_samples: Optional[int] = None,
        use_plot: bool = False,
        **pipeline_kwargs,
    ) -> Dict[str, Any]:
        """
        Evaluate a model on a dataset using a custom evaluation function.

        Args:
            model_name: Name of the model to evaluate
            dataset_class: Dataset class to use
            evaluation_function: Function that takes (ground_truth, prediction) and returns metrics
            max_samples: Maximum number of samples to evaluate (None for all)
            **pipeline_kwargs: Additional arguments for model pipeline

        Returns:
            Dictionary containing evaluation results
        """
        print(
            f"Starting evaluation with model {model_name} on dataset {dataset_class.__name__}"
        )
        print("=" * 60)

        # Load model
        pipe = self.load_model(model_name, **pipeline_kwargs)

        # Load dataset
        dataset = self.load_dataset(dataset_class, max_samples=max_samples)

        # Check for existing results to resume from
        existing_count = self._get_existing_results_count(
            model_name, dataset_class.__name__
        )
        start_idx = existing_count

        # Limit samples if specified
        if max_samples is not None:
            dataset_size = min(len(dataset), max_samples)
            print(f"Processing samples {start_idx} to {dataset_size}...")
        else:
            dataset_size = len(dataset)
            print(f"Processing samples {start_idx} to {dataset_size}...")

        if start_idx >= dataset_size:
            print(f"✅ All {dataset_size} samples already processed!")
            return self._consolidate_jsonl_results(model_name, dataset_class.__name__)

        # Initialize tracking
        total_samples = dataset_size
        successful_inferences = existing_count  # Start with existing count
        all_metrics = []
        results = []
        first_error_printed = False  # Track if we've printed the first error

        print("\nRunning inference...")
        print("=" * 80)

        # Get max_new_tokens for generation (default 1000)
        max_new_tokens = pipeline_kwargs.pop("max_new_tokens", 1000)

        # Load existing metrics if resuming
        if start_idx > 0:
            print(f"📂 Loading existing results from {start_idx} completed samples...")
            jsonl_file = self._get_jsonl_file_path(model_name, dataset_class.__name__)
            if os.path.exists(jsonl_file):
                with open(jsonl_file, "r") as f:
                    for line in f:
                        if line.strip():
                            result = json.loads(line.strip())
                            all_metrics.append(result["metrics"])
                            results.append(result)

        # Process each sample starting from where we left off
        for idx in tqdm(range(start_idx, dataset_size), desc="Processing samples"):
            try:
                sample = dataset[idx]
                plot_data = None

                # Clean up prompt for TSQADataset (if needed)
                if use_plot and hasattr(sample, "get") and sample.get("prompt"):
                    plot_data = self.get_plot_from_prompt(sample["prompt"])
                    pattern = r"The following is the accelerometer data on the [xyz]-axis\n([\-0-9, ]+)"
                    sample["prompt"] = re.sub(pattern, "", sample["prompt"])

                # Clean up prompt for TSQADataset (if needed)
                if hasattr(sample, "get") and sample.get("prompt"):
                    pattern = r"This is the time series, it has mean (-?\d+\.\d{4}) and std (-?\d+\.\d{4})\."
                    replacement = "This is the time series:"
                    sample["prompt"] = re.sub(pattern, replacement, sample["prompt"])

                # Create input text
                input_text = sample["prompt"]
                target_answer = sample["answer"]

                # Generate prediction
                outputs = pipe(
                    input_text,
                    max_new_tokens=max_new_tokens,
                    return_full_text=False,
                    plot_data=plot_data,
                )

                # Extract generated text
                if outputs and len(outputs) > 0:
                    generated_text = outputs[0]["generated_text"].strip()
                    successful_inferences += 1

                    # Evaluate using custom function (optionally with sample)
                    try:
                        import inspect

                        sig = inspect.signature(evaluation_function)
                        if len(sig.parameters) >= 3:
                            metrics = evaluation_function(
                                target_answer, generated_text, sample
                            )
                        else:
                            metrics = evaluation_function(target_answer, generated_text)
                    except Exception:
                        # Fallback to 2-arg call
                        metrics = evaluation_function(target_answer, generated_text)
                    all_metrics.append(metrics)

                    # Store detailed results
                    result = {
                        "sample_idx": idx,
                        "input_text": input_text,
                        "target_answer": target_answer,
                        "generated_answer": generated_text,
                        "metrics": metrics,
                    }
                    # Include template_id if present in sample for downstream analysis
                    if isinstance(sample, dict) and "template_id" in sample:
                        result["template_id"] = sample["template_id"]
                    results.append(result)

                    # Save individual result immediately to prevent data loss
                    self._save_individual_result(
                        result, model_name, dataset_class.__name__
                    )

                    # Print progress for first few samples
                    if idx < 10:
                        print(f"\nSAMPLE {idx + 1}:")
                        print(f"PROMPT: {input_text}...")
                        print(f"TARGET: {target_answer}")
                        print(f"PREDICTION: {generated_text}")
                        print(f"METRICS: {metrics}")
                        print("=" * 80)

                    # Print first error for debugging
                    if not first_error_printed and metrics.get("accuracy", 1) == 0:
                        print(f"\n❌ FIRST ERROR (Sample {idx + 1}):")
                        print(f"TARGET: {target_answer}")
                        print(f"PREDICTION: {generated_text}")
                        print("=" * 80)
                        first_error_printed = True
            except Exception as e:
                print(f"Error processing sample {idx}: {e}")
                continue

        # Calculate aggregate metrics
        if successful_inferences > 0:
            # Aggregate metrics across all samples
            aggregate_metrics = self._aggregate_metrics(all_metrics)

            # Calculate success rate
            success_rate = successful_inferences / total_samples

            # Prepare final results
            final_results = {
                "model_name": model_name,
                "dataset_name": dataset_class.__name__,
                "total_samples": total_samples,
                "successful_inferences": successful_inferences,
                "success_rate": success_rate,
                "metrics": aggregate_metrics,
                "detailed_results": results,
            }

            # Print summary
            self._print_summary(final_results)

            # Consolidate JSONL results into final JSON file
            consolidated_file = self._consolidate_jsonl_results(
                model_name, dataset_class.__name__
            )
            if consolidated_file:
                # Update the consolidated file with correct total_samples
                with open(consolidated_file, "r") as f:
                    consolidated_data = json.load(f)
                consolidated_data["total_samples"] = total_samples
                consolidated_data["success_rate"] = success_rate
                with open(consolidated_file, "w") as f:
                    json.dump(consolidated_data, f, indent=2)

            return final_results
        else:
            print("❌ No successful inferences completed!")
            return {
                "model_name": model_name,
                "dataset_name": dataset_class.__name__,
                "total_samples": total_samples,
                "successful_inferences": 0,
                "success_rate": 0.0,
                "metrics": {},
                "detailed_results": [],
            }

    def _aggregate_metrics(self, metrics_list: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Aggregate metrics across all samples.

        Args:
            metrics_list: List of metric dictionaries

        Returns:
            Aggregated metrics
        """
        if not metrics_list:
            return {}

        # Get all unique metric keys
        all_keys = set()
        for metrics in metrics_list:
            all_keys.update(metrics.keys())

        aggregated = {}
        for key in all_keys:
            values = [metrics.get(key, 0) for metrics in metrics_list]
            if all(isinstance(v, (int, float)) for v in values):
                # Calculate overall accuracy/percentage
                accuracy = np.mean(values) * 100
                aggregated[key] = accuracy
            else:
                # For non-numeric metrics, just count occurrences
                aggregated[key] = {
                    "values": values,
                    "count": len(values),
                }

        return aggregated

    def _print_summary(self, results: Dict[str, Any]):
        """Print evaluation summary."""
        print("\n" + "=" * 80)
        print("EVALUATION RESULTS")
        print("=" * 80)
        print(f"Model: {results['model_name']}")
        print(f"Dataset: {results['dataset_name']}")
        print(f"Total samples processed: {results['total_samples']}")
        print(f"Successful inferences: {results['successful_inferences']}")
        print(f"Success rate: {results['success_rate']:.2%}")

        if results["metrics"]:
            print("\nAggregated Metrics:")
            for metric_name, metric_values in results["metrics"].items():
                if isinstance(metric_values, (int, float)):
                    print(f"  {metric_name}: {metric_values:.1f}%")
                else:
                    print(f"  {metric_name}: {metric_values}")

    def _save_results(self, results: Dict[str, Any]):
        """Save detailed results to file."""
        import os

        current_dir = os.path.dirname(os.path.abspath(__file__))
        detailed_dir = os.path.join(
            current_dir, "..", "results", "baseline", "detailed"
        )
        os.makedirs(detailed_dir, exist_ok=True)
        normalized_model_id = re.sub(r"[^a-z0-9]", "-", results["model_name"].lower())
        normalized_dataset_name = re.sub(
            r"[^a-z0-9]", "-", results["dataset_name"].lower()
        )
        results_file = os.path.join(
            detailed_dir,
            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.json",
        )

        with open(results_file, "w") as f:
            json.dump(results, f, indent=2)

        print(f"\nDetailed results saved to: {results_file}")

    def _save_individual_result(
        self, result: Dict[str, Any], model_name: str, dataset_name: str
    ):
        """Save individual result incrementally to prevent data loss."""
        import os

        current_dir = os.path.dirname(os.path.abspath(__file__))
        detailed_dir = os.path.join(
            current_dir, "..", "results", "baseline", "detailed"
        )
        os.makedirs(detailed_dir, exist_ok=True)
        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
        results_file = os.path.join(
            detailed_dir,
            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
        )

        # Append individual result as JSONL
        with open(results_file, "a") as f:
            json.dump(result, f)
            f.write("\n")

    def _consolidate_jsonl_results(self, model_name: str, dataset_name: str) -> str:
        """Consolidate JSONL results into final JSON file."""
        import os

        current_dir = os.path.dirname(os.path.abspath(__file__))
        detailed_dir = os.path.join(
            current_dir, "..", "results", "baseline", "detailed"
        )
        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())

        jsonl_file = os.path.join(
            detailed_dir,
            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
        )
        json_file = os.path.join(
            detailed_dir,
            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.json",
        )

        # Read all JSONL results
        individual_results = []
        if os.path.exists(jsonl_file):
            with open(jsonl_file, "r") as f:
                for line in f:
                    if line.strip():
                        individual_results.append(json.loads(line.strip()))

        # Create consolidated results structure
        if individual_results:
            # Calculate aggregate metrics
            all_metrics = [result["metrics"] for result in individual_results]
            aggregate_metrics = self._aggregate_metrics(all_metrics)

            # Calculate success rate
            successful_inferences = len(individual_results)
            total_samples = len(individual_results)  # This will be updated by caller

            consolidated_results = {
                "model_name": model_name,
                "dataset_name": dataset_name,
                "total_samples": total_samples,
                "successful_inferences": successful_inferences,
                "success_rate": (
                    successful_inferences / total_samples if total_samples > 0 else 0.0
                ),
                "metrics": aggregate_metrics,
                "detailed_results": individual_results,
            }

            # Save consolidated results
            with open(json_file, "w") as f:
                json.dump(consolidated_results, f, indent=2)

            print(f"\nConsolidated results saved to: {json_file}")
            return json_file

        return None

    def _get_existing_results_count(self, model_name: str, dataset_name: str) -> int:
        """Get count of existing results from JSONL file for resuming interrupted evaluations."""
        import os

        current_dir = os.path.dirname(os.path.abspath(__file__))
        detailed_dir = os.path.join(
            current_dir, "..", "results", "baseline", "detailed"
        )
        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
        jsonl_file = os.path.join(
            detailed_dir,
            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
        )

        if os.path.exists(jsonl_file):
                    if line.strip():
                        count += 1
            return count
        return 0

    def _get_jsonl_file_path(self, model_name: str, dataset_name: str) -> str:
        """Get the JSONL file path for a model-dataset combination."""
        import os

        current_dir = os.path.dirname(os.path.abspath(__file__))
        detailed_dir = os.path.join(
            current_dir, "..", "results", "baseline", "detailed"
        )
        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
        return os.path.join(
            detailed_dir,
            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
        )

    def evaluate_multiple_models(
        self,
        model_names: List[str],
        dataset_classes: List[Type[Dataset]],
        evaluation_functions: Dict[str, Callable[[str, str], Dict[str, Any]]],
        max_samples: Optional[int] = None,
        **pipeline_kwargs,
    ) -> pd.DataFrame:
        """
        Evaluate multiple models on multiple datasets.

        Args:
            model_names: List of model names to evaluate
            dataset_classes: List of dataset classes to evaluate on
            evaluation_functions: Dictionary mapping dataset class names to evaluation functions
            max_samples: Maximum number of samples per evaluation
            **pipeline_kwargs: Additional arguments for model pipeline

        Returns:
            DataFrame with results for all model-dataset combinations
        """
        all_results = []

        # Generate filename once at the beginning
        import os

        current_dir = os.path.dirname(os.path.abspath(__file__))
        results_dir = os.path.join(current_dir, "..", "results", "baseline")
        os.makedirs(results_dir, exist_ok=True)
        df_filename = os.path.join(results_dir, "evaluation_results.csv")
        print(f"Results will be saved to: {df_filename}")

        # Load existing results if file exists
        existing_df = None
        if os.path.exists(df_filename):
            try:
                existing_df = pd.read_csv(df_filename)
                print(f"Found existing results file with {len(existing_df)} entries")
            except Exception as e:
                print(f"Warning: Could not read existing results file: {e}")

        for model_name in model_names:
            for dataset_class in dataset_classes:
                dataset_name = dataset_class.__name__

                if dataset_name not in evaluation_functions:
                    print(f"Warning: No evaluation function found for {dataset_name}")
                    continue

                # Check if this model-dataset combination already exists in results
                if existing_df is not None:
                    existing_result = existing_df[
                        (existing_df["model"] == model_name)
                        & (existing_df["dataset"] == dataset_name)
                    ]
                    if not existing_result.empty:
                        print(
                            f"⏭️  Skipping {model_name} on {dataset_name} (already evaluated)"
                        )
                        continue

                evaluation_function = evaluation_functions[dataset_name]

                print(f"\n{'=' * 80}")
                print(f"Evaluating {model_name} on {dataset_name}")
                print(f"{'=' * 80}")

                try:
                    results = self.evaluate_model_on_dataset(
                        model_name=model_name,
                        dataset_class=dataset_class,
                        evaluation_function=evaluation_function,
                        max_samples=max_samples,
                        use_plot=False,
                        **pipeline_kwargs,
                    )

                    # Extract key metrics for DataFrame
                    row = {
                        "model": model_name,
                        "dataset": dataset_name,
                        "total_samples": results["total_samples"],
                        "successful_inferences": results["successful_inferences"],
                        "success_rate": results["success_rate"],
                    }

                    # Add specific metrics
                    if results["metrics"]:
                        for metric_name, metric_values in results["metrics"].items():
                            if isinstance(metric_values, (int, float)):
                                row[metric_name] = metric_values
                            else:
                                row[metric_name] = str(metric_values)

                    all_results.append(row)

                    # Combine with existing results and save
                    current_df = pd.DataFrame(all_results)
                    if existing_df is not None:
                        # Append new results
                        final_df = pd.concat(
                            [existing_df, current_df], ignore_index=True
                        )
                    else:
                        final_df = current_df

                    final_df.to_csv(df_filename, index=False)
                    print(f"✅ Results updated: {df_filename}")

                except Exception as e:
                    print(f"Error evaluating {model_name} on {dataset_name}: {e}")
                    all_results.append(
                        {
                            "model": model_name,
                            "dataset": dataset_name,
                            "status": "Failed",
                        }
                    )

                    # Save DataFrame even after errors
                    current_df = pd.DataFrame(all_results)
                    if existing_df is not None:
                        final_df = pd.concat(
                            [existing_df, current_df], ignore_index=True
                        )
                    else:
                        final_df = current_df
                    final_df.to_csv(df_filename, index=False)
                    print(f"⚠️  Results updated (with error): {df_filename}")

        print(f"\nFinal results saved to: {df_filename}")
        return final_df

    def get_plot_from_prompt(self, prompt: str):
        """
        Parse time series data from the prompt and return a base64 image.
        """
        # Parse the time series data from the prompt
        time_series_data = []

        # Extract data for each axis using regex
        axes = ["x-axis", "y-axis", "z-axis"]
        for axis in axes:
            pattern = f"accelerometer data on the {axis}\\n([\\-0-9, ]+)"
            match = re.search(pattern, prompt.lower())
            if match:
                # Extract the data and convert to a list of integers
                data_str = match.group(1).strip()
                data_str = data_str.replace(" ", "")
                data = [int(val.strip()) for val in data_str.split(",") if val.strip()]
                time_series_data.append(data)

        # Create the plot
        num_series = len(time_series_data)
        fig, axes = plt.subplots(
            num_series, 1, figsize=(10, 4 * num_series), sharex=True
        )
        # If there's only one series, axes won't be an array
        if num_series == 1:
            axes = [axes]

        # Plot each time series in its own subplot
        axis_names = {0: "X-axis", 1: "Y-axis", 2: "Z-axis"}
        for i, series in enumerate(time_series_data):
            axes[i].plot(series, marker="o", linestyle="-", markersize=0)
            axes[i].grid(True, alpha=0.3)
            axes[i].set_title(f"Accelerometer - {axis_names.get(i)}")

        plt.tight_layout()

        # Convert plot to base64 image
        img_buffer = io.BytesIO()
        plt.savefig(img_buffer, format="png", bbox_inches="tight", dpi=100)
        plt.close()
        img_buffer.seek(0)
        image_data = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
        return image_data