Spaces:
Runtime error
Runtime error
File size: 2,171 Bytes
91bd31b 4a4836d 91bd31b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import unittest
from models import LlmRewardLabAction
from server.environment import LlmRewardLabEnvironment
class TestEnvLoop(unittest.TestCase):
def test_full_episode(self):
env = LlmRewardLabEnvironment()
obs = env.reset(task_id="task_detect_localize", seed=42)
self.assertEqual(obs.task_id, "task_detect_localize")
self.assertGreater(obs.budget_remaining, 0)
self.assertEqual(len(obs.quality_stats), 0) # stats hidden until inspect
self.assertFalse(obs.done)
obs = env.step(
LlmRewardLabAction(
action_type="inspect_samples",
parameters={"task_type": "summarization", "limit": 10},
)
)
self.assertLessEqual(len(obs.samples), 10)
self.assertTrue(all(s.task_type == "summarization" for s in obs.samples))
obs = env.step(
LlmRewardLabAction(
action_type="submit_diagnosis",
parameters={
"drift_events": ["data_contamination"],
"remediations": ["rollback_finetune_checkpoint"],
},
)
)
self.assertTrue(obs.done)
self.assertGreaterEqual(float(obs.reward or 0.0), 0.0)
self.assertLessEqual(float(obs.reward or 0.0), 1.0)
obs = env.step(
LlmRewardLabAction(action_type="inspect_samples", parameters={}),
)
self.assertTrue(obs.done)
def test_all_three_tasks(self):
env = LlmRewardLabEnvironment()
for task_id in ["task_detect_localize", "task_diagnose", "task_multi_drift"]:
obs = env.reset(task_id=task_id, seed=42)
self.assertEqual(obs.task_id, task_id)
obs = env.step(
LlmRewardLabAction(
action_type="submit_diagnosis",
parameters={"drift_events": [], "remediations": []},
)
)
self.assertTrue(obs.done)
self.assertGreaterEqual(float(obs.reward or 0.0), 0.0)
self.assertLessEqual(float(obs.reward or 0.0), 1.0)
if __name__ == "__main__":
unittest.main()
|