Upload folder using huggingface_hub

server/sql_env_environment.py

@@ -95,60 +95,40 @@ class SqlEnvironment(Environment):
             }
 
     def _reward(self, result: dict, query: str) -> float:
-        """
-        Reward based on query quality:
-        - Error:                        0.0
-        - Runs successfully:            +0.3
-        - Returns rows:                 +0.2
-        - Uses JOIN:                    +0.2
-        - Uses aggregation:             +0.2
-        - Uses WHERE:                   +0.1
-        - Columns match task expected:  +0.2 bonus
-        Max possible:                   1.0
-        """
         if result["error"]:
             return 0.0
 
         score = 0.0
         q = query.lower()
 
-        # Base score for running successfully
         score += 0.3
 
-        # Returned actual data
         if result["rowcount"] > 0:
             score += 0.2
 
-        # Bonus for JOIN
         if "join" in q:
             score += 0.2
 
-        # Bonus for aggregation
         if any(fn in q for fn in ["count(", "sum(", "avg(", "max(", "min("]):
             score += 0.2
 
-        # Bonus for WHERE filter
         if "where" in q:
             score += 0.1
 
-        # Bonus for matching expected columns from task
         if self._task and result["columns"]:
             matched = sum(
                 1 for c in self._task.expected_columns
-                if any(
-                    c in col.lower() or col.lower() in c  # ← bidirectional match
-                    for col in result["columns"]
-                )
+                if any(c in col.lower() or col.lower() in c
+                    for col in result["columns"])
             )
             score += 0.2 * (matched / len(self._task.expected_columns))
 
-        return min(score, 1.0)
+        # ← run grader if available
+        if self._task and self._task.grader and result["rows"]:
+            grader_score = self._task.grader(result["rows"])
+            score = min(score + (grader_score * 0.2), 1.0)
 
-    def _pick_task(self) -> Task:
-        """Pick a task based on TASK_DIFFICULTY env var, or random if unset."""
-        if TASK_DIFFICULTY in ("easy", "medium", "hard"):
-            return get_task_by_difficulty(TASK_DIFFICULTY)
-        return get_random_task()
+        return min(score, 1.0)
 
     # ------------------------------------------------------------------
     # Environment interface

tasks.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
-from typing import List
+from dataclasses import dataclass, field
+from typing import List, Callable, Any
 import random
 
 @dataclass
@@ -9,6 +9,35 @@ class Task:
     expected_columns: List[str]
     validation_query: str
     difficulty: str
+    grader: Callable[[List[Any]], float] = None  # ← add this
+
+
+def grader_salary_above_50000(rows: List[Any]) -> float:
+    """All returned rows must have salary > 50000."""
+    if not rows:
+        return 0.0
+    correct = all(row[1] > 50000 for row in rows)
+    return 1.0 if correct else 0.0
+
+
+def grader_employee_department_join(rows: List[Any]) -> float:
+    """Must return rows with at least 2 columns (name + department)."""
+    if not rows:
+        return 0.0
+    correct = all(len(row) >= 2 for row in rows)
+    return 1.0 if correct else 0.0
+
+
+def grader_avg_salary_per_department(rows: List[Any]) -> float:
+    """Must return rows with department name and a numeric avg salary."""
+    if not rows:
+        return 0.0
+    try:
+        correct = all(isinstance(row[1], (int, float)) for row in rows)
+        return 1.0 if correct else 0.0
+    except (IndexError, TypeError):
+        return 0.0
+
 
 TASKS = [
     Task(
@@ -17,6 +46,7 @@ TASKS = [
         expected_columns=["name", "salary"],
         validation_query="SELECT name, salary FROM employees WHERE salary > 50000",
         difficulty="easy",
+        grader=grader_salary_above_50000,
     ),
     Task(
         id="task_02",
@@ -24,6 +54,7 @@ TASKS = [
         expected_columns=["name", "department"],
         validation_query="SELECT e.name, d.name AS department FROM employees e JOIN departments d ON e.department = d.id",
         difficulty="medium",
+        grader=grader_employee_department_join,
     ),
     Task(
         id="task_03",
@@ -31,6 +62,7 @@ TASKS = [
         expected_columns=["name", "avg_salary"],
         validation_query="SELECT d.name, AVG(e.salary) AS avg_salary FROM employees e JOIN departments d ON e.department = d.id GROUP BY d.id",
         difficulty="hard",
+        grader=grader_avg_salary_per_department,
     ),
 ]