dslighting 1.3.9__py3-none-any.whl

dsat/__init__.py

@@ -0,0 +1,3 @@
+# Core framework exports
+from .models import TaskDefinition, TaskType
+from .workflows import DSATWorkflow

dsat/benchmark/__init__.py

@@ -0,0 +1 @@
+# benchmarks/__init__.py

dsat/benchmark/benchmark.py

@@ -0,0 +1,168 @@
+import asyncio
+import json
+from pathlib import Path
+from typing import Any, Callable, List, Tuple, Optional, Dict
+import pandas as pd
+import logging
+
+logger = logging.getLogger(__name__)
+
+class BaseBenchmark:
+    """
+    Abstract base class for all benchmark tests.
+    """
+    def __init__(self, name: str, file_path: Optional[str], log_path: str, **kwargs):
+        self.name = name
+        self.file_path = file_path
+        self.log_path = log_path
+        self.problems = self._load_problems()
+        # self.results_path = Path(self.log_path) / f"{self.name}_results.csv"
+        from datetime import datetime
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.results_path = Path(self.log_path) / f"{self.name}_results_{timestamp}.csv"
+        self.mismatches_path = Path(self.log_path) / f"{self.name}_mismatches.log"
+
+    def _load_problems(self) -> List[Dict[str, Any]]:
+        """Load problems from jsonl file."""
+        # MODIFICATION: Handle cases where file_path is not provided.
+        if not self.file_path:
+            logger.debug("No file_path provided. Subclass is expected to override _load_problems.")
+            return []
+        
+        with open(self.file_path, "r", encoding="utf-8") as f:
+            return [json.loads(line) for line in f]
+
+    def get_result_columns(self) -> List[str]:
+        raise NotImplementedError
+
+    async def evaluate_problem(self, problem: Dict, eval_fn: Callable, **kwargs) -> Tuple:
+        raise NotImplementedError
+
+    def log_mismatch(self, **kwargs):
+        with open(self.mismatches_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(kwargs) + "\n")
+
+    # REFACTORED: The main evaluation loop now accepts and passes down `eval_fn`.
+    async def run_evaluation(self, eval_fn: Callable, **kwargs):
+        """
+        Run the entire benchmark evaluation.
+
+        Args:
+            eval_fn: The generic evaluation function provided by DSATRunner.get_eval_function().
+        """
+        if not self.problems:
+            logger.error(f"Evaluation for '{self.name}' aborted: No problems were loaded.")
+            return
+
+        logger.info(f"Starting evaluation for benchmark '{self.name}' with {len(self.problems)} problems.")
+
+        results = []
+        tasks = [self.evaluate_problem(problem, eval_fn=eval_fn) for problem in self.problems]
+
+        # Use native asyncio for task completion
+        for future in asyncio.as_completed(tasks):
+            try:
+                # evaluate_problem returns (csv_tuple, report, error_message)
+                result_tuple, report, error_message = await future
+                results.append(result_tuple)
+            except Exception as e:
+                logger.error(f"An unexpected error occurred in evaluate_problem: {e}", exc_info=True)
+
+        # Save results to CSV
+        df = pd.DataFrame(results, columns=self.get_result_columns())
+        df.to_csv(self.results_path, index=False)
+
+        # Add metadata summary
+        self._append_metadata_to_csv(df, **kwargs)
+
+        logger.info(f"Evaluation complete. Results saved to {self.results_path}")
+
+    def _append_metadata_to_csv(self, df: pd.DataFrame, **kwargs):
+        """Append metadata summary to the CSV file."""
+        try:
+            import numpy as np
+
+            # Calculate statistics
+            score_col = 'score' if 'score' in df.columns else None
+            cost_col = 'cost' if 'cost' in df.columns else None
+            running_time_col = 'running_time' if 'running_time' in df.columns else None
+            input_tokens_col = 'input_tokens' if 'input_tokens' in df.columns else None
+            output_tokens_col = 'output_tokens' if 'output_tokens' in df.columns else None
+            total_tokens_col = 'total_tokens' if 'total_tokens' in df.columns else None
+
+            stats = {}
+            if score_col:
+                valid_scores = df[score_col].dropna()
+                if len(valid_scores) > 0:
+                    stats['avg_score'] = valid_scores.mean()
+                    stats['median_score'] = valid_scores.median()
+                    stats['std_score'] = valid_scores.std()
+
+            if cost_col:
+                valid_costs = df[cost_col].dropna()
+                if len(valid_costs) > 0:
+                    stats['avg_cost'] = valid_costs.mean()
+                    stats['total_cost'] = valid_costs.sum()
+
+            if running_time_col:
+                valid_times = df[running_time_col].dropna()
+                if len(valid_times) > 0:
+                    stats['avg_running_time'] = valid_times.mean()
+                    stats['total_running_time'] = valid_times.sum()
+
+            if input_tokens_col:
+                valid_input_tokens = df[input_tokens_col].dropna()
+                if len(valid_input_tokens) > 0:
+                    stats['avg_input_tokens'] = valid_input_tokens.mean()
+                    stats['total_input_tokens'] = valid_input_tokens.sum()
+
+            if output_tokens_col:
+                valid_output_tokens = df[output_tokens_col].dropna()
+                if len(valid_output_tokens) > 0:
+                    stats['avg_output_tokens'] = valid_output_tokens.mean()
+                    stats['total_output_tokens'] = valid_output_tokens.sum()
+
+            if total_tokens_col:
+                valid_total_tokens = df[total_tokens_col].dropna()
+                if len(valid_total_tokens) > 0:
+                    stats['avg_total_tokens'] = valid_total_tokens.mean()
+                    stats['total_total_tokens'] = valid_total_tokens.sum()
+
+            # Get model info from kwargs
+            model_info = kwargs.get('model_name', 'N/A')
+
+            # Create metadata rows
+            meta_rows = [
+                [''] * len(df.columns),  # Empty separator row
+                ['=== METADATA ==='] + [''] * (len(df.columns) - 1),
+            ]
+
+            meta_data = {
+                'Model': model_info,
+                'Total Tasks': len(df),
+                'Average Score': f"{stats.get('avg_score', 0):.4f}" if 'avg_score' in stats else 'N/A',
+                'Median Score': f"{stats.get('median_score', 0):.4f}" if 'median_score' in stats else 'N/A',
+                'Std Score': f"{stats.get('std_score', 0):.4f}" if 'std_score' in stats else 'N/A',
+                'Average Cost': f"${stats.get('avg_cost', 0):.4f}" if 'avg_cost' in stats else 'N/A',
+                'Total Cost': f"${stats.get('total_cost', 0):.4f}" if 'total_cost' in stats else 'N/A',
+                'Average Running Time': f"{stats.get('avg_running_time', 0):.4f}s" if 'avg_running_time' in stats else 'N/A',
+                'Total Running Time': f"{stats.get('total_running_time', 0):.4f}s" if 'total_running_time' in stats else 'N/A',
+                'Average Input Tokens': f"{stats.get('avg_input_tokens', 0):.0f}" if 'avg_input_tokens' in stats else 'N/A',
+                'Total Input Tokens': f"{stats.get('total_input_tokens', 0):.0f}" if 'total_input_tokens' in stats else 'N/A',
+                'Average Output Tokens': f"{stats.get('avg_output_tokens', 0):.0f}" if 'avg_output_tokens' in stats else 'N/A',
+                'Total Output Tokens': f"{stats.get('total_output_tokens', 0):.0f}" if 'total_output_tokens' in stats else 'N/A',
+                'Average Total Tokens': f"{stats.get('avg_total_tokens', 0):.0f}" if 'avg_total_tokens' in stats else 'N/A',
+                'Total Total Tokens': f"{stats.get('total_total_tokens', 0):.0f}" if 'total_total_tokens' in stats else 'N/A',
+            }
+
+            for key, value in meta_data.items():
+                meta_rows.append([key, value] + [''] * (len(df.columns) - 2))
+
+            # Append to CSV
+            with open(self.results_path, 'a', encoding='utf-8') as f:
+                for row in meta_rows:
+                    f.write(','.join(str(x) for x in row) + '\n')
+
+            logger.info(f"Metadata appended to {self.results_path}")
+        except Exception as e:
+            logger.warning(f"Failed to append metadata: {e}")

dsat/benchmark/datasci.py

@@ -0,0 +1,291 @@
+# dsat/benchmark/datasci.py
+
+import json
+import uuid
+import yaml
+import shutil
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Callable, List, Tuple, Optional, Dict
+import pandas as pd
+import logging
+
+from dsat.benchmark.benchmark import BaseBenchmark
+from dsat.models.task import TaskDefinition
+
+logger = logging.getLogger(__name__)
+
+
+class DataSciBenchmark(BaseBenchmark):
+    """
+    Benchmark class for DataSciBench tasks.
+    
+    DataSciBench tasks involve multi-step data science workflows where:
+    - Input: prompt.json (task description) + optional input data files
+    - Output: Generated files that are compared against ground truth
+    - Evaluation: Uses metric.yaml to define evaluation functions
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        file_path: Optional[str],
+        log_path: str,
+        datasci_root_dir: Optional[str] = None,
+        tasks: Optional[List[str]] = None,
+        **kwargs
+    ):
+        # Set up root directory before calling parent constructor
+        if datasci_root_dir:
+            self.root_dir = Path(datasci_root_dir)
+        else:
+            # Default to DataSciBench_Selected in benchmarks directory
+            self.root_dir = Path(__file__).parent.parent.parent / "benchmarks" / "DataSciBench_Selected"
+        
+        self.data_dir = self.root_dir / "data"
+        self.metric_dir = self.root_dir / "metric"
+        self.gt_data_dir = self.root_dir / "gt_data"
+        
+        # Override tasks if provided via CLI
+        self.task_filter = tasks
+        
+        # Call parent constructor (which calls _load_problems)
+        super().__init__(name, file_path, log_path, **kwargs)
+        
+        Path(self.log_path).mkdir(parents=True, exist_ok=True)
+        
+        # Re-initialize problems after setting up directories
+        self.problems = self._load_problems()
+        logger.info(f"DataSciBenchmark initialized with root_dir: {self.root_dir}")
+        logger.info(f"Loaded {len(self.problems)} tasks")
+    
+    def _load_problems(self) -> List[Dict[str, Any]]:
+        """Load DataSciBench tasks from the data directory."""
+        if not self.data_dir.exists():
+            logger.error(f"DataSciBench data directory not found: {self.data_dir}")
+            return []
+        
+        problems = []
+        
+        # Get all task directories
+        task_dirs = sorted([d for d in self.data_dir.iterdir() if d.is_dir()])
+        
+        for task_dir in task_dirs:
+            task_id = task_dir.name
+            
+            # Apply task filter if specified
+            if self.task_filter and task_id not in self.task_filter:
+                continue
+            
+            prompt_file = task_dir / "prompt.json"
+            if not prompt_file.exists():
+                logger.warning(f"Skipping task '{task_id}': no prompt.json found")
+                continue
+            
+            # Load prompt
+            try:
+                with open(prompt_file, 'r', encoding='utf-8') as f:
+                    prompt_data = json.load(f)
+            except Exception as e:
+                logger.warning(f"Skipping task '{task_id}': failed to load prompt.json: {e}")
+                continue
+            
+            # Check for metric file
+            metric_file = self.metric_dir / task_id / "metric.yaml"
+            if not metric_file.exists():
+                logger.warning(f"Skipping task '{task_id}': no metric.yaml found")
+                continue
+            
+            # Get input files (exclude prompt.json)
+            input_files = [f.name for f in task_dir.iterdir() 
+                          if f.is_file() and f.name not in ('prompt.json', 'orig_prompt.json')]
+            
+            problems.append({
+                "task_id": task_id,
+                "prompt": prompt_data.get("prompt", ""),
+                "data_source_type": prompt_data.get("data_source_type", "unknown"),
+                "input_files": input_files,
+                "task_dir": str(task_dir),
+                "metric_file": str(metric_file),
+                "gt_dir": str(self.gt_data_dir / task_id / "gt"),
+            })
+            logger.debug(f"Loaded task: {task_id}")
+        
+        if not problems:
+            logger.error(f"No valid tasks found in {self.data_dir}")
+        
+        return problems
+    
+    def get_result_columns(self) -> List[str]:
+        """Define columns for results CSV."""
+        return [
+            "task_id",
+            "output_dir",
+            "gt_dir", 
+            "score",
+            "cost",
+            "files_generated",
+            "evaluation_passed",
+            "error_message",
+        ]
+    
+    def _evaluate_task_outputs(
+        self,
+        task_id: str,
+        output_dir: Path,
+        metric_file: Path,
+        gt_dir: Path
+    ) -> Tuple[float, bool, Optional[str]]:
+        """
+        Evaluate task outputs using metric.yaml evaluation functions.
+        
+        Returns:
+            Tuple of (score, passed, error_message)
+        """
+        try:
+            # Load metric configuration
+            with open(metric_file, 'r', encoding='utf-8') as f:
+                metric_config = yaml.safe_load(f)
+            
+            if not metric_config or 'TMC-list' not in metric_config:
+                return 0.0, False, "Invalid metric.yaml: missing TMC-list"
+            
+            tmc_list = metric_config['TMC-list']
+            total_score = 0
+            max_score = len(tmc_list) * 2  # Each task can score 0, 1, or 2
+            
+            for metric_item in tmc_list:
+                eval_code = metric_item.get('code', '')
+                gt_file = metric_item.get('ground_truth', '')
+                
+                if not eval_code:
+                    continue
+                
+                # Prepare ground truth path
+                gt_path = gt_dir / gt_file if gt_file else gt_dir
+                
+                try:
+                    # Create evaluation context
+                    import os
+                    original_cwd = os.getcwd()
+                    os.chdir(str(output_dir))
+                    
+                    # Execute evaluation code
+                    local_vars = {'ground_truth': str(gt_path)}
+                    exec(eval_code, {'__builtins__': __builtins__, 'pd': pd}, local_vars)
+                    
+                    # Find the evaluation function and call it
+                    for name, obj in local_vars.items():
+                        if callable(obj) and name != '__builtins__':
+                            try:
+                                result = obj(str(gt_path))
+                                if result is True:
+                                    total_score += 2
+                                elif result is not None:
+                                    total_score += 1
+                            except Exception as func_error:
+                                logger.debug(f"Evaluation function {name} failed: {func_error}")
+                            break
+                    
+                    os.chdir(original_cwd)
+                    
+                except Exception as eval_error:
+                    logger.debug(f"Evaluation error for {task_id}: {eval_error}")
+                    try:
+                        os.chdir(original_cwd)
+                    except:
+                        pass
+            
+            # Calculate normalized score (0-1)
+            score = total_score / max_score if max_score > 0 else 0.0
+            passed = score > 0
+            
+            return score, passed, None
+            
+        except Exception as e:
+            return 0.0, False, str(e)
+    
+    async def evaluate_problem(
+        self,
+        problem: Dict[str, Any],
+        eval_fn: Callable
+    ) -> Tuple[Tuple, Any, Optional[str]]:
+        """
+        Evaluate a single DataSciBench task.
+        """
+        task_id = problem.get("task_id")
+        if not task_id:
+            raise ValueError("Problem data must contain 'task_id'")
+        
+        prompt = problem.get("prompt", "")
+        task_dir = Path(problem.get("task_dir", ""))
+        metric_file = Path(problem.get("metric_file", ""))
+        gt_dir = Path(problem.get("gt_dir", ""))
+        
+        # Create unique output directory
+        unique_id = uuid.uuid4().hex[:6]
+        output_dir = Path(self.log_path) / f"output_{task_id}_{unique_id}"
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        cost = 0.0
+        error_message: Optional[str] = None
+        files_generated = 0
+        evaluation_passed = False
+        score = 0.0
+        
+        try:
+            # Create TaskDefinition for the workflow
+            task = TaskDefinition(
+                task_id=task_id,
+                task_type="datasci",
+                payload={
+                    "prompt": prompt,
+                    "input_dir": str(task_dir),
+                    "output_dir": str(output_dir),
+                    "expected_outputs": [],  # DataSciBench doesn't require specific output files
+                }
+            )
+            
+            # Execute the workflow
+            result, cost = await eval_fn(task)
+            
+            # Check for errors
+            if isinstance(result, str) and result.startswith("[ERROR]"):
+                error_message = result
+                logger.error(f"Task {task_id} failed: {error_message}")
+            else:
+                # Count generated files
+                files_generated = len(list(output_dir.glob("*")))
+                
+                # Evaluate outputs against ground truth
+                if metric_file.exists() and gt_dir.exists():
+                    score, evaluation_passed, eval_error = self._evaluate_task_outputs(
+                        task_id, output_dir, metric_file, gt_dir
+                    )
+                    if eval_error:
+                        error_message = f"Evaluation error: {eval_error}"
+                else:
+                    # If no metric/gt, consider it passed if files were generated
+                    evaluation_passed = files_generated > 0
+                    score = 1.0 if evaluation_passed else 0.0
+                
+                logger.info(f"Task {task_id}: score={score:.2f}, files={files_generated}, passed={evaluation_passed}")
+        
+        except Exception as e:
+            error_message = f"Error during DataSciBenchmark evaluation of {task_id}: {e}"
+            logger.error(error_message, exc_info=True)
+        
+        # Build result tuple
+        csv_tuple = (
+            task_id,
+            str(output_dir),
+            str(gt_dir),
+            score,
+            cost,
+            files_generated,
+            evaluation_passed,
+            error_message,
+        )
+        
+        return csv_tuple, {"score": score, "passed": evaluation_passed}, error_message
+