(self, action)
| 191 | return 0 |
| 192 | |
| 193 | def step(self, action): |
| 194 | # TODO: first step obs does not have question. |
| 195 | obs, _, done, info = self.env.step(action) |
| 196 | reward = self.get_reward(info) |
| 197 | if done: |
| 198 | obs = f"Episode finished, reward = {reward}\n" |
| 199 | info.update({"gt_answer": self.data[self.data_idx][1], "question_idx": self.data_idx}) |
| 200 | info.update({'em': reward, 'reward': reward, 'f1': reward}) |
| 201 | return obs, reward, done, info |
| 202 | |
| 203 | def __len__(self): |
| 204 | return len(self.data) |