Upload environment.py with huggingface_hub

environment.py

@@ -0,0 +1,404 @@
+
+import json
+import pandas as pd
+import numpy as np
+from typing import Any, Dict, Optional, Tuple
+from models import Action, Observation, Reward, StepResult, TaskInfo
+
+AVAILABLE_OPERATIONS = [
+    "remove_duplicates",
+    "fill_missing", 
+    "fix_dtype",
+    "remove_outliers",
+    "rename_columns",
+    "validate_schema",
+    "finish"
+]
+
+class DataCleaningEnv:
+    """
+    OpenEnv-compliant Data Cleaning Environment.
+    Agent must clean dirty datasets step by step.
+    """
+
+    def __init__(self, task_id: str = "easy_dedup_rename"):
+        self.task_id = task_id
+        self.current_df: Optional[pd.DataFrame] = None
+        self.gold_df: Optional[pd.DataFrame] = None
+        self.task_meta: Dict = {}
+        self.step_count: int = 0
+        self.done: bool = False
+        self.max_steps: int = 10
+        self.reward_history = []
+        self._load_task_metadata()
+
+    # ─────────────────────────────────────────
+    # INTERNAL HELPERS
+    # ─────────────────────────────────────────
+
+    def _load_task_metadata(self):
+        with open("datasets/task_metadata.json", "r") as f:
+            all_meta = json.load(f)
+
+        mapping = {
+            "easy_dedup_rename":   "easy",
+            "medium_missing_dtype": "medium",
+            "hard_full_pipeline":   "hard"
+        }
+        key = mapping.get(self.task_id, "easy")
+        self.task_meta = all_meta[key]
+        self.max_steps = self.task_meta["max_steps"]
+
+    def _load_datasets(self):
+        mapping = {
+            "easy_dedup_rename":    "easy",
+            "medium_missing_dtype": "medium",
+            "hard_full_pipeline":   "hard"
+        }
+        folder = mapping.get(self.task_id, "easy")
+        self.current_df = pd.read_csv(f"datasets/{folder}/dirty.csv")
+        self.gold_df    = pd.read_csv(f"datasets/{folder}/gold.csv")
+
+    def _get_observation(self, message: str = "") -> Observation:
+        df = self.current_df
+        missing = {col: int(df[col].isnull().sum()) for col in df.columns}
+        dtypes  = {col: str(df[col].dtype) for col in df.columns}
+        sample  = df.head(3).fillna("NULL").to_dict(orient="records")
+
+        return Observation(
+            task_id=self.task_id,
+            step=self.step_count,
+            dataset_info={
+                "total_rows": len(df),
+                "total_columns": len(df.columns),
+                "has_duplicates": bool(df.duplicated().any()),
+                "has_missing": bool(df.isnull().any().any()),
+            },
+            columns=list(df.columns),
+            shape=[len(df), len(df.columns)],
+            missing_values=missing,
+            dtypes=dtypes,
+            duplicate_count=int(df.duplicated().sum()),
+            sample_rows=sample,
+            available_operations=self.task_meta.get(
+                "operations_allowed", AVAILABLE_OPERATIONS
+            ),
+            task_description=self.task_meta.get("description", ""),
+            message=message
+        )
+
+    def _compute_reward(self) -> Reward:
+        df   = self.current_df
+        gold = self.gold_df
+        scoring = self.task_meta.get("scoring", {})
+
+        dup_score     = 0.0
+        missing_score = 0.0
+        dtype_score   = 0.0
+        outlier_score = 0.0
+        schema_score  = 0.0
+        penalty       = 0.0
+
+        # ── Duplicate score ──────────────────────────────────────────
+        if "duplicate_score" in scoring:
+            gold_rows = len(gold)
+            curr_rows = len(df)
+            if curr_rows == gold_rows:
+                dup_score = 1.0
+            elif curr_rows < gold_rows:
+                dup_score = max(0.0, curr_rows / gold_rows)
+            else:
+                excess = curr_rows - gold_rows
+                dup_score = max(0.0, 1.0 - (excess / gold_rows))
+
+        # ── Missing value score ──────────────────────────────────────
+        if "missing_score" in scoring:
+            total_cells  = df.shape[0] * df.shape[1]
+            missing_curr = int(df.isnull().sum().sum())
+            missing_gold = int(gold.isnull().sum().sum())
+            if total_cells > 0:
+                filled = max(0, missing_curr - missing_gold)
+                missing_score = 1.0 - (filled / total_cells)
+                missing_score = max(0.0, min(1.0, missing_score))
+
+        # ── Dtype score ──────────────────────────────────────────────
+        if "dtype_score" in scoring:
+            common_cols = [c for c in gold.columns if c in df.columns]
+            if common_cols:
+                matches = sum(
+                    1 for c in common_cols
+                    if str(df[c].dtype) == str(gold[c].dtype)
+                )
+                dtype_score = matches / len(common_cols)
+
+        # ── Outlier score ────────────────────────────────────────────
+        if "outlier_score" in scoring:
+            num_cols = gold.select_dtypes(include=[np.number]).columns
+            scores = []
+            for col in num_cols:
+                if col not in df.columns:
+                    continue
+                gold_mean = gold[col].mean()
+                gold_std  = gold[col].std() + 1e-9
+                curr_col  = pd.to_numeric(df[col], errors="coerce").dropna()
+                outliers  = ((curr_col - gold_mean).abs() > 3 * gold_std).sum()
+                col_score = max(0.0, 1.0 - outliers / (len(curr_col) + 1e-9))
+                scores.append(col_score)
+            outlier_score = float(np.mean(scores)) if scores else 0.0
+
+        # ── Schema score ─────────────────────────────────────────────
+        if "schema_score" in scoring:
+            gold_cols = list(gold.columns)
+            curr_cols = list(df.columns)
+            matched   = sum(1 for c in gold_cols if c in curr_cols)
+            schema_score = matched / len(gold_cols) if gold_cols else 0.0
+
+        # ── Penalty for too many steps ───────────────────────────────
+        step_ratio = self.step_count / self.max_steps
+        if step_ratio > 0.8:
+            penalty = 0.05
+
+        # ── Weighted total ───────────────────────────────────────────
+        weights = {
+            "duplicate_score": scoring.get("duplicate_score", 0.0),
+            "missing_score":   scoring.get("missing_score",   0.0),
+            "dtype_score":     scoring.get("dtype_score",     0.0),
+            "outlier_score":   scoring.get("outlier_score",   0.0),
+            "schema_score":    scoring.get("schema_score",    0.0),
+        }
+        component_scores = {
+            "duplicate_score": dup_score,
+            "missing_score":   missing_score,
+            "dtype_score":     dtype_score,
+            "outlier_score":   outlier_score,
+            "schema_score":    schema_score,
+        }
+        total = sum(
+            component_scores[k] * w
+            for k, w in weights.items()
+        ) - penalty
+        total = max(0.0, min(1.0, total))
+
+        return Reward(
+            total=round(total, 4),
+            duplicate_score=round(dup_score, 4),
+            missing_score=round(missing_score, 4),
+            dtype_score=round(dtype_score, 4),
+            outlier_score=round(outlier_score, 4),
+            schema_score=round(schema_score, 4),
+            penalty=round(penalty, 4)
+        )
+
+    # ─────────────────────────────────────────
+    # OPERATIONS
+    # ─────────────────────────────────────────
+
+    def _op_remove_duplicates(self, params: Dict) -> str:
+        before = len(self.current_df)
+        subset = params.get("subset", None)
+        self.current_df = self.current_df.drop_duplicates(subset=subset)
+        self.current_df = self.current_df.reset_index(drop=True)
+        removed = before - len(self.current_df)
+        return f"Removed {removed} duplicate rows. Rows: {before} → {len(self.current_df)}"
+
+    def _op_fill_missing(self, params: Dict) -> str:
+        col      = params.get("column")
+        strategy = params.get("strategy", "mean")
+        messages = []
+
+        cols_to_fill = [col] if col else list(self.current_df.columns)
+        for c in cols_to_fill:
+            if self.current_df[c].isnull().sum() == 0:
+                continue
+            if strategy == "mean":
+                numeric = pd.to_numeric(self.current_df[c], errors="coerce")
+                fill_val = numeric.mean()
+                self.current_df[c] = numeric.fillna(round(fill_val, 2))
+            elif strategy == "median":
+                numeric = pd.to_numeric(self.current_df[c], errors="coerce")
+                fill_val = numeric.median()
+                self.current_df[c] = numeric.fillna(fill_val)
+            elif strategy == "mode":
+                fill_val = self.current_df[c].mode()[0]
+                self.current_df[c] = self.current_df[c].fillna(fill_val)
+            elif strategy == "ffill":
+                self.current_df[c] = self.current_df[c].ffill()
+            else:
+                fill_val = strategy
+                self.current_df[c] = self.current_df[c].fillna(fill_val)
+            messages.append(f"{c}→{strategy}")
+
+        return f"Filled missing values: {', '.join(messages)}"
+
+    def _op_fix_dtype(self, params: Dict) -> str:
+        col    = params.get("column")
+        dtype  = params.get("dtype", "auto")
+        messages = []
+
+        cols_to_fix = [col] if col else list(self.current_df.columns)
+        for c in cols_to_fix:
+            try:
+                if dtype == "int" or dtype == "auto":
+                    converted = pd.to_numeric(self.current_df[c], errors="coerce")
+                    if converted.notna().all():
+                        self.current_df[c] = converted.astype(int)
+                        messages.append(f"{c}→int")
+                elif dtype == "float":
+                    self.current_df[c] = pd.to_numeric(
+                        self.current_df[c], errors="coerce"
+                    )
+                    messages.append(f"{c}→float")
+                elif dtype == "str":
+                    self.current_df[c] = self.current_df[c].astype(str)
+                    messages.append(f"{c}→str")
+            except Exception as e:
+                messages.append(f"{c}→failed({e})")
+
+        return f"Fixed dtypes: {', '.join(messages)}"
+
+    def _op_remove_outliers(self, params: Dict) -> str:
+        col    = params.get("column")
+        method = params.get("method", "iqr")
+        before = len(self.current_df)
+        messages = []
+
+        cols = [col] if col else list(
+            self.current_df.select_dtypes(include=[np.number]).columns
+        )
+        for c in cols:
+            series = pd.to_numeric(self.current_df[c], errors="coerce")
+            if method == "iqr":
+                Q1  = series.quantile(0.25)
+                Q3  = series.quantile(0.75)
+                IQR = Q3 - Q1
+                mask = (series >= Q1 - 1.5 * IQR) & (series <= Q3 + 1.5 * IQR)
+                self.current_df = self.current_df[mask | series.isna()]
+            elif method == "zscore":
+                mean = series.mean()
+                std  = series.std() + 1e-9
+                mask = ((series - mean).abs() <= 3 * std)
+                self.current_df = self.current_df[mask | series.isna()]
+            self.current_df = self.current_df.reset_index(drop=True)
+            messages.append(c)
+
+        removed = before - len(self.current_df)
+        return f"Removed {removed} outlier rows from: {', '.join(messages)}"
+
+    def _op_rename_columns(self, params: Dict) -> str:
+        mapping = params.get("mapping", {})
+        if not mapping:
+            # Auto snake_case
+            new_names = {
+                col: col.lower().replace(" ", "_")
+                for col in self.current_df.columns
+            }
+            self.current_df = self.current_df.rename(columns=new_names)
+            return f"Auto renamed columns to snake_case: {list(new_names.values())}"
+        self.current_df = self.current_df.rename(columns=mapping)
+        return f"Renamed columns: {mapping}"
+
+    def _op_validate_schema(self, params: Dict) -> str:
+        gold_cols = list(self.gold_df.columns)
+        curr_cols = list(self.current_df.columns)
+        missing   = [c for c in gold_cols if c not in curr_cols]
+        extra     = [c for c in curr_cols if c not in gold_cols]
+        if not missing and not extra:
+            return "Schema valid! All columns match gold standard."
+        msg = []
+        if missing:
+            msg.append(f"Missing columns: {missing}")
+        if extra:
+            msg.append(f"Extra columns: {extra}")
+        return "Schema issues: " + " | ".join(msg)
+
+    # ─────────────────────────────────────────
+    # OPENENV API
+    # ─────────────────────────────────────────
+
+    def reset(self) -> StepResult:
+        self._load_datasets()
+        self.step_count   = 0
+        self.done         = False
+        self.reward_history = []
+
+        obs    = self._get_observation("Environment reset. Start cleaning!")
+        reward = Reward(total=0.0)
+
+        return StepResult(
+            observation=obs,
+            reward=reward,
+            done=False,
+            info={"task_id": self.task_id, "max_steps": self.max_steps}
+        )
+
+    def step(self, action: Action) -> StepResult:
+        if self.done:
+            obs = self._get_observation("Episode already done. Call reset().")
+            return StepResult(
+                observation=obs,
+                reward=Reward(total=0.0),
+                done=True,
+                info={"warning": "Episode already done"}
+            )
+
+        self.step_count += 1
+        op     = action.operation
+        params = action.parameters
+        message = ""
+
+        # ── Route operation ──────────────────────────────────────────
+        try:
+            if op == "remove_duplicates":
+                message = self._op_remove_duplicates(params)
+            elif op == "fill_missing":
+                message = self._op_fill_missing(params)
+            elif op == "fix_dtype":
+                message = self._op_fix_dtype(params)
+            elif op == "remove_outliers":
+                message = self._op_remove_outliers(params)
+            elif op == "rename_columns":
+                message = self._op_rename_columns(params)
+            elif op == "validate_schema":
+                message = self._op_validate_schema(params)
+            elif op == "finish":
+                message = "Agent called finish."
+                self.done = True
+            else:
+                message = f"Unknown operation: {op}. No changes made."
+        except Exception as e:
+            message = f"Operation failed: {str(e)}"
+
+        # ── Check max steps ──────────────────────────────────────────
+        if self.step_count >= self.max_steps:
+            self.done = True
+            message += " | Max steps reached."
+
+        reward = self._compute_reward()
+        self.reward_history.append(reward.total)
+        obs = self._get_observation(message)
+
+        return StepResult(
+            observation=obs,
+            reward=reward,
+            done=self.done,
+            info={
+                "step": self.step_count,
+                "operation": op,
+                "reward_history": self.reward_history
+            }
+        )
+
+    def state(self) -> Dict[str, Any]:
+        if self.current_df is None:
+            return {"status": "not initialized — call reset() first"}
+        return {
+            "task_id":        self.task_id,
+            "step":           self.step_count,
+            "done":           self.done,
+            "shape":          list(self.current_df.shape),
+            "columns":        list(self.current_df.columns),
+            "missing_values": self.current_df.isnull().sum().to_dict(),
+            "duplicate_count": int(self.current_df.duplicated().sum()),
+            "reward_history": self.reward_history,
+            "dtypes":         {c: str(t) for c, t in self.current_df.dtypes.items()}
+        }