Upload folder using huggingface_hub

server/data_cleaning_env_environment.py

@@ -149,6 +149,7 @@ def _grade_hard(rows: list[dict]) -> float:
 
     return round(sum(scores) / len(scores), 4)
 
+
 GRADERS = {
     "easy":   _grade_easy,
     "medium": _grade_medium,
@@ -161,6 +162,7 @@ DATASETS = {
     "hard":   _make_hard_dataset,
 }
 
+
 # ---------------------------------------------------------------------------
 # Environment
 # ---------------------------------------------------------------------------
@@ -171,27 +173,21 @@ MAX_STEPS = {
     "hard": 25,
 }
 
-VALID_OPERATIONS_PER_TASK = {
-    "easy": {
-        "impute_mean",
-        "impute_mode",
-        "drop_missing_rows",
-    },
-    "medium": {
-        "remove_duplicates",
-        "fix_type_errors",
-        "drop_missing_rows",
-    },
-    "hard": {
-        "drop_missing_rows",
-        "fill_quantity_mean",
-        "remove_duplicates",
-        "fix_type_errors",
-        "remove_outliers",
-        "normalize_text",
-    },
+VALID_OPERATIONS = {
+    # Easy task ops
+    "impute_mean",        # fill numeric NaN with column mean
+    "impute_mode",        # fill categorical NaN with column mode
+    "drop_missing_rows",  # drop all rows that contain any None
+    # Medium task ops
+    "remove_duplicates",  # drop exact duplicate rows
+    "fix_type_errors",    # coerce non-numeric 'age' / numeric cols to float (NaN if fails)
+    # Hard task ops
+    "remove_outliers",    # drop rows where price < 0 or price > 500
+    "normalize_text",     # strip + title-case all string columns
+    "fill_quantity_mean", # fill missing quantity with column mean
 }
 
+
 class DataCleaningEnvironment(Environment):
     """
     Data Cleaning RL Environment.
@@ -202,7 +198,9 @@ class DataCleaningEnvironment(Environment):
     Each task presents a different level of difficulty and requires
     different cleaning operations.
     """
+
     SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
     # -----------------------------------------------------------------------
     # Lifecycle
     # -----------------------------------------------------------------------
@@ -268,13 +266,12 @@ class DataCleaningEnvironment(Environment):
         op = action.operation.strip().lower()
         col = getattr(action, "column", None)
 
-        valid_ops = VALID_OPERATIONS_PER_TASK[self._task]
-        if op not in valid_ops:
+        if op not in VALID_OPERATIONS:
             # Invalid operation → small penalty, episode continues
             reward = -0.05
             self._last_reward = reward
             obs = self._make_observation(reward=reward, done=False)
-            obs.metadata["error"] = f"Unknown operation '{op}'. Valid: {sorted(valid_ops)}"
+            obs.metadata["error"] = f"Unknown operation '{op}'. Valid: {sorted(VALID_OPERATIONS)}"
             return obs
 
         before_score = GRADERS[self._task](self._rows)
@@ -283,17 +280,16 @@ class DataCleaningEnvironment(Environment):
 
         # Reward = improvement in quality score (partial progress signal)
         improvement = after_score - before_score
-        # 0.01
         if improvement > 0:
-            reward = round(improvement + 0.02, 4)   # bonus for improvement
+            reward = round(improvement + 0.02, 4)   # bonus for any improvement
         elif improvement < 0:
             reward = round(improvement - 0.02, 4)   # penalty for harming dataset
         else:
-            reward = 0.0                            # neutral if no change
-        
+            reward = 0.0                             # neutral if no change
+
         if after_score >= 1.0:
-            reward += 1.0                       # big bonus for perfect score (episode complete)
-        
+            reward += 1.0                            # big bonus for completing the task
+
         self._last_reward = reward
         self._applied_ops.append(op)
 
@@ -327,7 +323,7 @@ class DataCleaningEnvironment(Environment):
     def grade(self) -> dict:
         """Return final grader score for the current episode."""
         raw = GRADERS[self._task](self._rows)
-        score = max(0.001, min(0.999, raw))
+        score = max(0.001, min(0.999, raw))   # strictly (0, 1) as required
         return {
             "task": self._task,
             "score": score,
@@ -379,7 +375,7 @@ class DataCleaningEnvironment(Environment):
                 "action_schema": {
                     "operation": {
                         "type": "string",
-                        "enum": list(VALID_OPERATIONS_PER_TASK),
+                        "enum": list(VALID_OPERATIONS),
                     },
                     "column": {"type": "string", "description": "Target column (optional)"},
                 },
@@ -403,7 +399,7 @@ class DataCleaningEnvironment(Environment):
                 if not self._done:
                     self.step(DataCleaningAction(operation=op))
             raw = GRADERS[task_id](self._rows)
-            results[task_id] = max(0.001, min(0.999, raw))
+            results[task_id] = max(0.001, min(0.999, raw))  # strictly (0, 1)
         return {"baseline_scores": results}
 
     # -----------------------------------------------------------------------
@@ -415,7 +411,7 @@ class DataCleaningEnvironment(Environment):
         has_dupes = self._has_duplicates()
         has_outliers = self._has_outliers()
         score = GRADERS[self._task](self._rows)
-        valid_ops = VALID_OPERATIONS_PER_TASK[self._task]
+
         return DataCleaningObservation(
             current_text=self._rows_to_text(),
             is_normalized=not has_outliers and missing_count == 0 and not has_dupes,
@@ -431,11 +427,13 @@ class DataCleaningEnvironment(Environment):
                 "has_duplicates": has_dupes,
                 "has_outliers": has_outliers,
                 "quality_score": score,
-                "valid_operations": sorted(valid_ops),
+                "valid_operations": sorted(VALID_OPERATIONS),
                 "ops_already_applied": list(self._applied_ops),
+                "recommended_next": _recommend_next(
+                    self._task, missing_count, has_dupes, has_outliers, self._applied_ops
+                ),
             },
         )
-
     def _rows_to_text(self) -> str:
         if not self._rows:
             return "[]"
@@ -506,10 +504,10 @@ class DataCleaningEnvironment(Environment):
                 for r in self._rows:
                     val = r.get(c)
                     if val is not None and not isinstance(val, (int, float)):
-                        val_str = str(val).strip()
-                        if re.fullmatch(r"-?\d+(\.\d+)?", val_str):
-                            r[c] = float(val_str)
-                        else:
+                        # Try to coerce to float; set None if it fails
+                        try:
+                            r[c] = float(re.sub(r"[^\d.\-]", "", str(val)))
+                        except ValueError:
                             r[c] = None
 
         elif op == "remove_outliers":
@@ -533,6 +531,7 @@ class DataCleaningEnvironment(Environment):
                     if r.get("quantity") is None:
                         r["quantity"] = mean
 
+
 # ---------------------------------------------------------------------------
 # Baseline policies (deterministic rule-based agents)
 # ---------------------------------------------------------------------------
@@ -550,47 +549,49 @@ _BASELINE_POLICIES: dict[str, list[str]] = {
     ],
 }
 
+
 # ---------------------------------------------------------------------------
 # Recommendation helper (guides the LLM toward correct next op)
 # ---------------------------------------------------------------------------
 
-# def _recommend_next(
-#     task: str,
-#     missing_count: int,
-#     has_dupes: bool,
-#     has_outliers: bool,
-#     applied_ops: list,
-# ) -> str:
-#     """Return a plain-English hint for the LLM about the best next operation."""
-#     applied = set(applied_ops)
-
-#     if task == "easy":
-#         if missing_count > 0:
-#             return "There are missing values. Use impute_mean (numeric) or impute_mode (text)."
-#         return "No issues remain. Episode should be complete."
-
-#     if task == "medium":
-#         if has_dupes and "remove_duplicates" not in applied:
-#             return "Duplicate rows exist. Use remove_duplicates."
-#         if "fix_type_errors" not in applied:
-#             return "Non-numeric values in numeric columns. Use fix_type_errors."
-#         if missing_count > 0 and "drop_missing_rows" not in applied:
-#             return "Some values still missing after type fix. Use drop_missing_rows."
-#         return "No issues remain. Episode should be complete."
-
-#     # hard
-#     if missing_count > 0 and "fill_quantity_mean" not in applied:
-#         return "Missing quantity values. Use fill_quantity_mean first."
-#     if missing_count > 0 and "drop_missing_rows" not in applied:
-#         return "Missing product/category values. Use drop_missing_rows."
-#     if has_outliers and "remove_outliers" not in applied:
-#         return "Price outliers detected (price<=0 or price>=500). Use remove_outliers."
-#     if "normalize_text" not in applied:
-#         return "String columns have inconsistent casing/whitespace. Use normalize_text."
-#     if has_dupes and "remove_duplicates" not in applied:
-#         return "Duplicate rows remain. Use remove_duplicates."
-#     return "All issues fixed. Episode should be complete."
-
+def _recommend_next(
+    task: str,
+    missing_count: int,
+    has_dupes: bool,
+    has_outliers: bool,
+    applied_ops: list,
+) -> str:
+    """Return a plain-English hint for the LLM about the best next operation."""
+    applied = set(applied_ops)
+
+    if task == "easy":
+        if missing_count > 0:
+            return "There are missing values. Use impute_mean (numeric) or impute_mode (text)."
+        return "No issues remain. Episode should be complete."
+
+    if task == "medium":
+        if has_dupes and "remove_duplicates" not in applied:
+            return "Duplicate rows exist. Use remove_duplicates."
+        if "fix_type_errors" not in applied:
+            return "Non-numeric values in numeric columns. Use fix_type_errors."
+        if missing_count > 0 and "drop_missing_rows" not in applied:
+            return "Some values still missing after type fix. Use drop_missing_rows."
+        return "No issues remain. Episode should be complete."
+
+    # hard
+    if missing_count > 0 and "fill_quantity_mean" not in applied:
+        return "Missing quantity values. Use fill_quantity_mean first."
+    if missing_count > 0 and "drop_missing_rows" not in applied:
+        return "Missing product/category values. Use drop_missing_rows."
+    if has_outliers and "remove_outliers" not in applied:
+        return "Price outliers detected (price<=0 or price>=500). Use remove_outliers."
+    if "normalize_text" not in applied:
+        return "String columns have inconsistent casing/whitespace. Use normalize_text."
+    if has_dupes and "remove_duplicates" not in applied:
+        return "Duplicate rows remain. Use remove_duplicates."
+    return "All issues fixed. Episode should be complete."
+
+\
 # ---------------------------------------------------------------------------
 # Column utility helpers
 # ---------------------------------------------------------------------------
@@ -603,6 +604,7 @@ def _numeric_cols(rows: list[dict]) -> list[str]:
         if any(isinstance(r.get(c), (int, float)) for r in rows)
     ]
 
+
 def _string_cols(rows: list[dict]) -> list[str]:
     if not rows:
         return []