sravaniamere commited on
Commit
fd851a9
·
1 Parent(s): c7d40b5

fix state as method not property, fix abstract class error

Browse files
Files changed (1) hide show
  1. sql_env/server.py +44 -1
sql_env/server.py CHANGED
@@ -29,7 +29,6 @@ class SQLCorrectionEnvironment(Environment):
29
  self._rewards_history = []
30
 
31
  def reset(self, difficulty: str = "easy", task_id: str = None, **kwargs) -> SQLObservation:
32
- # task_id and difficulty are the same thing in our env
33
  actual_difficulty = task_id or difficulty or "easy"
34
  self._difficulty = actual_difficulty
35
  tasks = TASK_SETS.get(actual_difficulty, TASK_SETS["easy"])
@@ -49,6 +48,50 @@ class SQLCorrectionEnvironment(Environment):
49
  reward=0.001,
50
  done=False,
51
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
52
 
53
  def step(self, action: SQLAction) -> SQLObservation:
54
  # Auto-reset if no task loaded (create_app may use fresh instances)
 
29
  self._rewards_history = []
30
 
31
  def reset(self, difficulty: str = "easy", task_id: str = None, **kwargs) -> SQLObservation:
 
32
  actual_difficulty = task_id or difficulty or "easy"
33
  self._difficulty = actual_difficulty
34
  tasks = TASK_SETS.get(actual_difficulty, TASK_SETS["easy"])
 
48
  reward=0.001,
49
  done=False,
50
  )
51
+
52
+ def step(self, action: SQLAction) -> SQLObservation:
53
+ if self._current_task is None:
54
+ self.reset()
55
+ self._step_count += 1
56
+ reward_obj = grade(action, self._current_task)
57
+ reward = reward_obj.value
58
+ self._last_reward = reward
59
+ self._rewards_history.append(reward)
60
+ done = (reward >= 0.95) or (self._step_count >= self._current_task.max_steps)
61
+ self._done = done
62
+ feedback = generate_feedback(action, self._current_task, reward_obj)
63
+ return SQLObservation(
64
+ task_id=self._current_task.task_id,
65
+ broken_query=self._current_task.broken_query,
66
+ schema_context=self._current_task.schema_context,
67
+ error_hint=self._current_task.error_hint,
68
+ step_number=self._step_count,
69
+ previous_attempt=action.corrected_query,
70
+ feedback=feedback,
71
+ reward=reward,
72
+ done=done,
73
+ )
74
+
75
+ def state(self) -> SQLState:
76
+ if self._current_task is None:
77
+ return SQLState(
78
+ task_id="none",
79
+ difficulty="none",
80
+ step_count=0,
81
+ max_steps=0,
82
+ done=False,
83
+ last_reward=0.0,
84
+ rewards_history=[],
85
+ )
86
+ return SQLState(
87
+ task_id=self._current_task.task_id,
88
+ difficulty=self._difficulty,
89
+ step_count=self._step_count,
90
+ max_steps=self._current_task.max_steps,
91
+ done=self._done,
92
+ last_reward=self._last_reward,
93
+ rewards_history=self._rewards_history,
94
+ )
95
 
96
  def step(self, action: SQLAction) -> SQLObservation:
97
  # Auto-reset if no task loaded (create_app may use fresh instances)