""" inference.py ──────────── CogTraceEnv — Baseline Inference Script Meta × PyTorch OpenEnv Hackathon Required environment variables: API_BASE_URL — e.g. https://api-inference.huggingface.co/v1 MODEL_NAME — e.g. meta-llama/Meta-Llama-3.1-8B-Instruct HF_TOKEN — your Hugging Face token STDOUT FORMAT (strictly followed): [START] task= env= model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= score= rewards= """ from __future__ import annotations import os import sys import json import time import re import traceback from typing import Any, Dict, List, Optional from openai import OpenAI # ── Env config ──────────────────────────────────────────────────────────────── API_BASE_URL = os.getenv("API_BASE_URL", "https://api-inference.huggingface.co/v1") MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") HF_TOKEN = os.getenv("HF_TOKEN", "") BENCHMARK = "CogTraceEnv" NUM_SEEDS_EASY = 10 NUM_SEEDS_MEDIUM = 10 NUM_SEEDS_HARD = 5 SUCCESS_SCORE_THRESHOLD = 0.5 # score >= 0.5 counts as success client = OpenAI(api_key=HF_TOKEN or "dummy", base_url=API_BASE_URL) # ── Structured log helpers (exact format required by grader) ────────────────── def log_start(task: str, env: str, model: str) -> None: print(f"[START] task={task} env={env} model={model}", flush=True) def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: error_val = error if error else "null" done_val = str(done).lower() print( f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}", flush=True, ) def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True, ) # ── LLM helpers ─────────────────────────────────────────────────────────────── def _chat(system: str, user: str, max_tokens: int = 64, retries: int = 3) -> str: for attempt in range(retries): try: resp = client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": system}, {"role": "user", "content": user}, ], max_tokens=max_tokens, temperature=0.0, ) return resp.choices[0].message.content.strip() except Exception as e: if attempt < retries - 1: time.sleep(2 ** attempt) else: print(f"[DEBUG] LLM error: {e}", file=sys.stderr, flush=True) return "0" return "0" def _parse_int(text: str, lo: int = 0, hi: int = 3) -> int: nums = re.findall(r"\d+", text) if not nums: return 0 return max(lo, min(hi, int(nums[0]))) # ── Task agents ─────────────────────────────────────────────────────────────── TASK1_SYSTEM = """\ You are a clinical AI assistant. You will receive behavioral metrics for a patient being monitored for Alzheimer's disease. Your task: classify the patient's CDR (Clinical Dementia Rating) stage. 0 = Healthy / no impairment 1 = Very mild impairment 2 = Mild impairment 3 = Moderate impairment 4 = Severe impairment Respond with a SINGLE integer 0-4. No explanation.""" def task1_agent(obs: Dict[str, Any]) -> int: user = f"Patient behavioral snapshot:\n{json.dumps(obs, indent=2)}\n\nStage (0-4):" raw = _chat(TASK1_SYSTEM, user, max_tokens=8) return _parse_int(raw, lo=0, hi=4) TASK2_SYSTEM = """\ You are a clinical AI monitoring a patient for Alzheimer's-related deterioration. Each day you receive behavioral sensor readings. You must choose an action: 0 = do_nothing 1 = soft_alert 2 = medium_alert 3 = escalate Raise your FIRST alert (action >= 1) as soon as you detect an anomaly. Respond with a SINGLE integer 0-3. No explanation.""" def task2_agent(obs: Dict[str, Any], step: int) -> int: user = ( f"Day {step + 1} observation:\n{json.dumps(obs, indent=2)}\n\n" "Action (0=do_nothing, 1=soft_alert, 2=medium_alert, 3=escalate):" ) raw = _chat(TASK2_SYSTEM, user, max_tokens=8) return _parse_int(raw, lo=0, hi=3) TASK3_SYSTEM = """\ You are a clinical AI monitoring a patient over 30 days. Each day, choose the most appropriate action: 0 = do_nothing (patient appears stable) 1 = soft_alert (minor concern, flag for review) 2 = medium_alert (notable deterioration, schedule checkup) 3 = escalate (acute event, immediate intervention) - Most days the patient is stable, prefer do_nothing to avoid alert fatigue. - When multiple signals deteriorate together, escalate proportionally. - Avoid raising more than 4 alerts in any 7-day window. Respond with a SINGLE integer 0-3. No explanation.""" def task3_agent(obs: Dict[str, Any], step: int, history: List[int]) -> int: recent = history[-7:] if len(history) >= 7 else history user = ( f"Day {step + 1}/30 observation:\n{json.dumps(obs, indent=2)}\n" f"Recent actions (last {len(recent)} days): {recent}\n\n" "Action (0-3):" ) raw = _chat(TASK3_SYSTEM, user, max_tokens=8) return _parse_int(raw, lo=0, hi=3) # ── Task runners ────────────────────────────────────────────────────────────── def run_task1() -> float: from task1_easy import run_episode, grade, build_env log_start(task="task1_easy", env=BENCHMARK, model=MODEL_NAME) results = [] rewards: List[float] = [] steps_taken = 0 score = 0.0 success = False try: for seed in range(NUM_SEEDS_EASY): env = build_env(seed) result = run_episode(env, agent_fn=task1_agent) results.append(result) reward = result["score"] rewards.append(reward) steps_taken += 1 log_step( step=seed + 1, action=f"predict_stage({result['predicted_stage']})", reward=reward, done=(seed == NUM_SEEDS_EASY - 1), error=None, ) score = grade(results) success = score >= SUCCESS_SCORE_THRESHOLD except Exception as e: traceback.print_exc() log_end(success=success, steps=steps_taken, score=score, rewards=rewards) return score def run_task2() -> float: from task2_medium import run_episode, grade, build_env log_start(task="task2_medium", env=BENCHMARK, model=MODEL_NAME) results = [] rewards: List[float] = [] steps_taken = 0 score = 0.0 success = False try: for seed in range(NUM_SEEDS_MEDIUM): env = build_env(seed) result = run_episode(env, agent_fn=task2_agent) results.append(result) reward = result["score"] rewards.append(reward) steps_taken += 1 log_step( step=seed + 1, action=f"first_alert_day({result['first_alert_day']})", reward=reward, done=(seed == NUM_SEEDS_MEDIUM - 1), error=None, ) score = grade(results) success = score >= SUCCESS_SCORE_THRESHOLD except Exception as e: traceback.print_exc() log_end(success=success, steps=steps_taken, score=score, rewards=rewards) return score def run_task3() -> float: from task3_hard import run_episode, grade, build_env log_start(task="task3_hard", env=BENCHMARK, model=MODEL_NAME) results = [] rewards: List[float] = [] steps_taken = 0 score = 0.0 success = False try: for seed in range(NUM_SEEDS_HARD): env = build_env(seed) # task3 needs history passed into agent; wrap to match run_episode signature history: List[int] = [] def agent_with_history(obs_dict, step, _hist=None): action = task3_agent(obs_dict, step, list(history)) history.append(action) return action result = run_episode(env, agent_fn=agent_with_history) results.append(result) reward = result["score"] rewards.append(reward) steps_taken += 1 log_step( step=seed + 1, action=f"episode(tp={result['tp']},fp={result['fp']},fn={result['fn']})", reward=reward, done=(seed == NUM_SEEDS_HARD - 1), error=None, ) score = grade(results) success = score >= SUCCESS_SCORE_THRESHOLD except Exception as e: traceback.print_exc() log_end(success=success, steps=steps_taken, score=score, rewards=rewards) return score # ── Main ────────────────────────────────────────────────────────────────────── def main(): task1_score = run_task1() task2_score = run_task2() task3_score = run_task3() mean = (task1_score + task2_score + task3_score) / 3 print(f"[DEBUG] Mean score across all tasks: {mean:.4f}", flush=True) if __name__ == "__main__": main()