""" inference.py ──────────── Baseline inference script for CogTraceEnv. Runs an LLM agent (via OpenAI-compatible client) against all 3 tasks and prints reproducible scores. 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 (used as API key) Usage: python inference.py Runtime: < 20 min on 2 vCPU / 8 GB RAM """ from __future__ import annotations import os import sys import json import time import traceback from typing import Any, Dict, List, Optional from openai import OpenAI # ── Env config ──────────────────────────────────────────────────────────────── API_BASE_URL = os.environ.get("API_BASE_URL", "https://api-inference.huggingface.co/v1") MODEL_NAME = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct") HF_TOKEN = os.environ.get("HF_TOKEN", "") NUM_SEEDS_EASY = 10 NUM_SEEDS_MEDIUM = 10 NUM_SEEDS_HARD = 5 # fewer because 30-step episodes take longer client = OpenAI(api_key=HF_TOKEN or "dummy", base_url=API_BASE_URL) # ── LLM helpers ─────────────────────────────────────────────────────────────── def _chat(system: str, user: str, max_tokens: int = 64, retries: int = 3) -> str: """Call the LLM and return the stripped response text.""" 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" [LLM ERROR] {e}", file=sys.stderr) return "0" return "0" def _parse_int(text: str, lo: int = 0, hi: int = 3) -> int: """Extract first integer from model output, clamped to [lo, hi].""" import re nums = re.findall(r"\d+", text) if not nums: return 0 return max(lo, min(hi, int(nums[0]))) # ── Task 1 — Stage Classification ──────────────────────────────────────────── 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) # ── Task 2 — Anomaly Timing ─────────────────────────────────────────────────── 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 — a sudden worsening across multiple signals compared to the patient's baseline. Be neither too early (normal variation) nor too late (delayed response). 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) # ── Task 3 — Full Triage ────────────────────────────────────────────────────── 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) Key principles: - 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 unless truly warranted 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) # ── Runner ──────────────────────────────────────────────────────────────────── def run_task1() -> float: from tasks.task1_easy import run_episode, grade, build_env results = [] for seed in range(NUM_SEEDS_EASY): env = build_env(seed) result = run_episode(env, agent_fn=task1_agent) results.append(result) print(f" seed={seed:02d} true={result['true_stage']} pred={result['predicted_stage']} score={result['score']:.2f}") return grade(results) def run_task2() -> float: from tasks.task2_medium import run_episode, grade, build_env results = [] for seed in range(NUM_SEEDS_MEDIUM): env = build_env(seed) result = run_episode(env, agent_fn=task2_agent) results.append(result) print(f" seed={seed:02d} anomaly_day={result['anomaly_day']} first_alert={result['first_alert_day']} score={result['score']:.2f}") return grade(results) def run_task3() -> float: from tasks.task3_hard import run_episode, grade, build_env results = [] for seed in range(NUM_SEEDS_HARD): env = build_env(seed) result = run_episode(env, agent_fn=task3_agent) results.append(result) print(f" seed={seed:02d} tp={result['tp']} fp={result['fp']} fn={result['fn']} f1={result['f1_score']:.2f} score={result['score']:.2f}") return grade(results) def main(): print("=" * 60) print("CogTraceEnv — Baseline Inference") print(f"Model : {MODEL_NAME}") print(f"API : {API_BASE_URL}") print("=" * 60) scores: Dict[str, Optional[float]] = {} for task_id, runner, label in [ ("task1_easy", run_task1, "Task 1 (Easy) — Stage Classification"), ("task2_medium", run_task2, "Task 2 (Medium) — Anomaly Timing"), ("task3_hard", run_task3, "Task 3 (Hard) — Full Triage"), ]: print(f"\n▶ {label}") t0 = time.time() try: score = runner() elapsed = time.time() - t0 scores[task_id] = score print(f" ✓ SCORE: {score:.4f} ({elapsed:.1f}s)") except Exception: print(f" ✗ FAILED") traceback.print_exc() scores[task_id] = None print("\n" + "=" * 60) print("FINAL SCORES") print("=" * 60) valid = [s for s in scores.values() if s is not None] for task_id, score in scores.items(): val = f"{score:.4f}" if score is not None else "ERROR" print(f" {task_id:<20} {val}") if valid: print(f"\n Mean score: {sum(valid) / len(valid):.4f}") print("=" * 60) if __name__ == "__main__": main()