"""Grading system for AutoClean-Ai data cleaning environment. Implements deterministic reward calculation, quality scoring, and task grading. All scores are normalized between 0.0 and 1.0. """ import pandas as pd import numpy as np from typing import Dict, Any, Tuple import re from models import CleaningActionType, DataCleaningAction def calculate_dataset_quality_score(df: pd.DataFrame, task_id: str) -> float: """ Calculate overall dataset quality score from 0.0 (worst) to 1.0 (perfect). Combines multiple quality metrics: - Null value percentage - Duplicate row percentage - Outlier percentage - Email validity (if applicable) - Data type correctness """ if df.empty: return 0.0 total_rows = len(df) quality_components = [] # 1. Null value handling (0.0 = all nulls, 1.0 = no nulls) null_percentage = df.isnull().sum().sum() / (df.shape[0] * df.shape[1]) null_score = 1.0 - null_percentage quality_components.append(null_score) # 2. Duplicate handling (0.0 = all duplicates, 1.0 = no duplicates) duplicate_count = df.duplicated().sum() duplicate_percentage = duplicate_count / total_rows if total_rows > 0 else 0.0 duplicate_score = 1.0 - duplicate_percentage quality_components.append(duplicate_score) # 3. Outlier detection for numeric columns numeric_columns = df.select_dtypes(include=[np.number]).columns outlier_scores = [] for col in numeric_columns: if len(df[col].dropna()) < 4: continue Q1 = df[col].quantile(0.25) Q3 = df[col].quantile(0.75) IQR = Q3 - Q1 if IQR == 0: continue outlier_mask = (df[col] < (Q1 - 1.5 * IQR)) | (df[col] > (Q3 + 1.5 * IQR)) outlier_percentage = outlier_mask.sum() / len(df[col]) outlier_scores.append(1.0 - outlier_percentage) if outlier_scores: outlier_score = sum(outlier_scores) / len(outlier_scores) quality_components.append(outlier_score) # 4. Email validation if email column exists if 'email' in df.columns: email_pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$' valid_emails = df['email'].astype(str).str.match(email_pattern, na=False).sum() email_validity = valid_emails / len(df['email']) if len(df['email']) > 0 else 0.0 quality_components.append(email_validity) # 5. Task specific quality checks if task_id == "task_1_basic_cleaning": # Basic task only requires null and duplicate handling weights = [0.5, 0.5] final_score = (null_score * weights[0] + duplicate_score * weights[1]) elif task_id == "task_2_intermediate_cleaning": # Intermediate adds outliers and email validation components = [null_score, duplicate_score] if outlier_scores: components.append(outlier_score) if 'email' in df.columns: components.append(email_validity) final_score = sum(components) / len(components) elif task_id == "task_3_full_pipeline": # Advanced task uses all components final_score = sum(quality_components) / len(quality_components) if quality_components else 0.0 else: # Default average final_score = sum(quality_components) / len(quality_components) if quality_components else 0.0 return max(0.0, min(1.0, final_score)) def calculate_reward( df: pd.DataFrame, initial_df: pd.DataFrame, previous_quality: float, current_quality: float, action: DataCleaningAction, task_id: str, step_count: int ) -> Tuple[float, Dict[str, Any]]: """ Calculate reward for an action. Returns (reward, reward_info). Reward components: - Quality improvement delta - Action validity bonus - Efficiency bonus - Progress bonus - Penalties for invalid/ineffective actions """ reward_info = { "quality_improvement": 0.0, "action_validity": 0.0, "efficiency_bonus": 0.0, "progress_bonus": 0.0, "penalty": 0.0, "total": 0.0 } # Base reward is quality improvement quality_delta = current_quality - previous_quality reward = quality_delta * 2.0 # Amplify delta for stronger signal reward_info["quality_improvement"] = quality_delta # Bonus for valid action execution reward += 0.05 reward_info["action_validity"] = 0.05 # Bonus for progress if quality_delta > 0: progress_bonus = 0.03 reward += progress_bonus reward_info["progress_bonus"] = progress_bonus # Efficiency bonus (earlier steps give higher reward) if step_count < 5 and quality_delta > 0: efficiency_bonus = 0.02 * (5 - step_count) reward += efficiency_bonus reward_info["efficiency_bonus"] = efficiency_bonus # Penalties if quality_delta < -0.01: # Negative quality change penalty penalty = abs(quality_delta) * 0.5 reward -= penalty reward_info["penalty"] += penalty # Penalty for dropping too many rows if len(df) < len(initial_df) * 0.5: row_loss_penalty = 0.1 reward -= row_loss_penalty reward_info["penalty"] += row_loss_penalty # Clamp final reward between -0.2 and 0.5 per step final_reward = max(-0.2, min(0.5, reward)) reward_info["total"] = final_reward return final_reward, reward_info def grade_task_result( initial_df: pd.DataFrame, final_df: pd.DataFrame, task_id: str, step_count: int ) -> float: """ Grade final task result. Returns score from 0.0 to 1.0. Grading criteria per task: - task_1_basic_cleaning: 40% null handling, 40% duplicate handling, 20% efficiency - task_2_intermediate_cleaning: 25% null, 30% email, 25% outliers, 20% efficiency - task_3_full_pipeline: Full weighted criteria """ if final_df.empty: return 0.0 final_quality = calculate_dataset_quality_score(final_df, task_id) initial_quality = calculate_dataset_quality_score(initial_df, task_id) quality_improvement = final_quality - initial_quality # Calculate efficiency score (better score for fewer steps) max_steps = 15 efficiency_score = max(0.0, 1.0 - (step_count / max_steps)) # Task specific grading if task_id == "task_1_basic_cleaning": # Basic: 40% null, 40% duplicates, 20% efficiency null_percentage = final_df.isnull().sum().sum() / (final_df.shape[0] * final_df.shape[1]) null_score = 1.0 - null_percentage duplicate_percentage = final_df.duplicated().sum() / len(final_df) duplicate_score = 1.0 - duplicate_percentage score = (null_score * 0.4) + (duplicate_score * 0.4) + (efficiency_score * 0.2) elif task_id == "task_2_intermediate_cleaning": # Intermediate: 25% null, 30% email, 25% outliers, 20% efficiency null_percentage = final_df.isnull().sum().sum() / (final_df.shape[0] * final_df.shape[1]) null_score = 1.0 - null_percentage email_score = 0.5 if 'email' in final_df.columns: email_pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$' valid_emails = final_df['email'].astype(str).str.match(email_pattern, na=False).sum() email_score = valid_emails / len(final_df['email']) outlier_score = 0.5 numeric_columns = final_df.select_dtypes(include=[np.number]).columns for col in numeric_columns: if len(final_df[col].dropna()) >= 4: Q1 = final_df[col].quantile(0.25) Q3 = final_df[col].quantile(0.75) IQR = Q3 - Q1 if IQR > 0: outlier_mask = (final_df[col] < (Q1 - 1.5 * IQR)) | (final_df[col] > (Q3 + 1.5 * IQR)) outlier_percentage = outlier_mask.sum() / len(final_df[col]) outlier_score = 1.0 - outlier_percentage break score = (null_score * 0.25) + (email_score * 0.30) + (outlier_score * 0.25) + (efficiency_score * 0.20) elif task_id == "task_3_full_pipeline": # Advanced: quality * 0.8 + efficiency * 0.2 score = (final_quality * 0.8) + (efficiency_score * 0.2) else: # Default grading score = final_quality # Ensure score is in valid range final_score = max(0.0, min(1.0, score)) return final_score