Spaces:
Sleeping
Sleeping
File size: 2,157 Bytes
ee117a1 | 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 63 64 65 | """training_wrapper.py — Minimal training-ready wrapper for the Medical Diagnostic Environment.
This module exposes a small async/sync wrapper that is easy to plug into a training
loop or evaluation script. It is not a full RL algorithm, but it makes the
environment easy to consume for model-based training.
"""
import os
import asyncio
from typing import Optional
from client import DiagnosticEnv
from models import DiagnosticAction, PatientObservation
class TrainingEnv:
"""Minimal wrapper exposing a training-friendly environment interface."""
def __init__(self, base_url: Optional[str] = None):
self.base_url = base_url or os.getenv("ENV_URL", "ws://localhost:8000/ws")
self._env = DiagnosticEnv(base_url=self.base_url)
async def __aenter__(self):
await self._env.__aenter__()
return self
async def __aexit__(self, exc_type, exc, tb):
await self._env.__aexit__(exc_type, exc, tb)
async def reset(self, difficulty: str = "easy") -> PatientObservation:
result = await self._env.reset(difficulty=difficulty)
return result.observation if hasattr(result, "observation") else result
async def step(
self,
action_type: str,
question: Optional[str] = None,
test_name: Optional[str] = None,
diagnosis: Optional[str] = None,
) -> PatientObservation:
action = DiagnosticAction(
action_type=action_type,
question=question,
test_name=test_name,
diagnosis=diagnosis,
)
result = await self._env.step(action)
return result.observation if hasattr(result, "observation") else result
async def state(self):
return await self._env.state()
async def run_demo():
"""Example usage of the training wrapper."""
async with TrainingEnv() as env:
obs = await env.reset(difficulty="easy")
print("Reset observation:", obs.message)
result = await env.step(action_type="ask_question", question="Do you have a fever?")
print("Step result:", result.message)
if __name__ == "__main__":
asyncio.run(run_demo())
|