AdamK29/Meta-OpenENV-Hackathon

1import torch2import asyncio3 4from client import EmailEnvClient5from core.models import EmailAction6from train_agent import PolicyNet, encode_state, ACTIONS7 8import ssl9import certifi10 11ssl._create_default_https_context = ssl._create_unverified_context12 13# ---------- LOAD MODEL ----------14 15model = PolicyNet()16model.load_state_dict(torch.load("email_agent.pth"))17model.eval()18 19 20def predict(obs):21    state = encode_state(obs)22    probs = model(state)23    action = ACTIONS[probs.argmax().item()]24    return action25 26 27# ---------- TEST 1: ENVIRONMENT ----------28 29async def test_environment():30 31    print("\n===== TESTING ON ENVIRONMENT =====\n")32 33    env = EmailEnvClient(base_url="https://adamk29-meta-openenv-hackathon.hf.space")34    await env.__aenter__()35 36    try:37        result = await env.reset(task="medium")38 39        step = 040        total_reward = 041 42        while True:43            step += 144 45            email = result.observation.email_text46            action = predict(result.observation)47 48            result = await env.step(49                EmailAction(action_type="classify", content=action)50            )51 52            reward = result.reward or 0.053            total_reward += reward54 55            print(f"""56STEP {step}57Email: {email}58Predicted: {action}59Reward: {reward:.2f}60--------------------------61""")62 63            if result.done:64                break65 66        print(f"\n✅ TOTAL REWARD: {total_reward:.2f}")67 68    finally:69        await env.__aexit__(None, None, None)70 71 72# ---------- TEST 2: CUSTOM EXAMPLES ----------73 74def test_custom_examples():75 76    print("\n===== TESTING CUSTOM EMAILS =====\n")77 78    class DummyObs:79        def __init__(self, text):80            self.email_text = text81 82    samples = [83        "Win a free iPhone now!!!",84        "Meeting at 5pm today",85        "Fix the server ASAP, urgent issue",86        "Client feedback received, please review",87        "Limited time offer, click now!!!"88    ]89 90    for text in samples:91        obs = DummyObs(text)92        action = predict(obs)93 94        print(f"""95Email: {text}96Predicted Label: {action}97--------------------------98""")99 100 101# ---------- MAIN ----------102 103async def main():104    await test_environment()105    test_custom_examples()106 107 108if __name__ == "__main__":109    asyncio.run(main())