pharma-vigilance / run_demo.py
modelbuilderhq's picture
Upload folder using huggingface_hub
f2beac3 verified
raw
history blame contribute delete
721 Bytes
from agent import RuleBasedPharmaAgent
from env import PharmaVigilanceEnv
def main() -> None:
env = PharmaVigilanceEnv()
agent = RuleBasedPharmaAgent()
for task_id in ("known_signal_easy", "cluster_signal_medium", "confounded_hard"):
observation = env.reset(task_id)
action = agent.act(observation)
observation, reward, done, info = env.step(action)
print(f"\nTask: {task_id}")
print(f"Action: {action.classification} / {action.suspect_drug}")
print(f"Reward: {reward.total:.2f}")
print(f"Done: {done}")
print(f"Feedback: {observation.feedback}")
print(f"Info: {info}")
if __name__ == "__main__":
main()