Spaces:
Sleeping
Sleeping
| """ | |
| inference.py — Baseline inference script for the Medical Diagnostic Environment. | |
| This script demonstrates how to run the environment with an LLM (using OpenAI client) | |
| and logs results in the exact format required by the hackathon. | |
| Format requirements: | |
| [START] task=<task_name> env=<env_name> model=<model_name> | |
| [STEP] step=<n> action=<action_str> reward=<float> done=<bool> error=<str|null> | |
| [END] success=<bool> steps=<n> score=<float> rewards=<comma_separated_list> | |
| All fields on a single line with NO NEWLINES within a line. | |
| """ | |
| import asyncio | |
| import os | |
| import json | |
| import textwrap | |
| from typing import List, Optional, Dict | |
| from openai import OpenAI | |
| import sys | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from client import DiagnosticEnv | |
| from models import DiagnosticAction | |
| # ============================================================================== | |
| # CONFIGURATION | |
| # ============================================================================== | |
| # Configuration | |
| API_KEY = os.environ.get("API_KEY") | |
| API_BASE_URL = os.environ.get("API_BASE_URL") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct") | |
| LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME", "medical-diagnostic-env:latest") | |
| ENV_URL = os.getenv("ENV_URL", "ws://localhost:8000/ws") | |
| BENCHMARK = os.getenv("BENCHMARK", "medical_diagnostic_env") | |
| # Inference configuration | |
| MAX_STEPS = 15 # Maximum steps per episode | |
| TEMPERATURE = 0.7 # LLM temperature for reasoning | |
| MAX_TOKENS = 256 # Max tokens per completion | |
| TASK_NAMES = ["easy_diagnosis", "medium_diagnosis", "hard_diagnosis"] | |
| DIFFICULTY_LEVELS = ["easy", "medium", "hard"] | |
| # ============================================================================== | |
| # LOGGING FUNCTIONS | |
| # ============================================================================== | |
| def log_start(task: str, env: str, model: str) -> None: | |
| """Log episode start in required format.""" | |
| # Clean model name for logging | |
| model_clean = model.split("/")[-1] if "/" in model else model | |
| print(f"[START] task={task} env={env} model={model_clean}", flush=True) | |
| def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: | |
| """Log single step in required format.""" | |
| error_val = f'"{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: | |
| """Log episode end in required format.""" | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) | |
| success_val = str(success).lower() | |
| print( | |
| f"[END] success={success_val} steps={steps} score={score:.3f} rewards={rewards_str}", | |
| flush=True, | |
| ) | |
| # ============================================================================== | |
| # LLM INTERACTION | |
| # ============================================================================== | |
| def create_system_prompt() -> str: | |
| """Create system prompt for medical diagnostic reasoning.""" | |
| return textwrap.dedent(""" | |
| You are an expert medical diagnostic AI assistant. Your role is to: | |
| 1. GATHER INFORMATION: Ask relevant clinical questions about symptoms, | |
| history, and presentation. | |
| 2. ORDER TESTS: Request appropriate diagnostic tests based on the | |
| differential diagnosis. | |
| 3. REASON DIAGNOSTICALLY: Consider the patient's presentation, | |
| synthesize findings, and make a diagnosis. | |
| Your reasoning should follow clinical guidelines and prioritize: | |
| - Life-threatening conditions first (red flags) | |
| - Most common diagnoses for the presentation | |
| - Efficiency (minimize unnecessary tests) | |
| When responding, use EXACTLY ONE of these actions: | |
| ACTION: ask_question | |
| QUESTION: <your question here> | |
| OR | |
| ACTION: order_test | |
| TEST: <test name> | |
| OR | |
| ACTION: submit_diagnosis | |
| DIAGNOSIS: <final diagnosis> | |
| Be concise. Diagnose within 10-15 steps if possible. | |
| """).strip() | |
| def extract_action_from_response(response: str) -> Optional[Dict]: | |
| """ | |
| Extract structured action from LLM response. | |
| Returns dict with: | |
| { | |
| "action_type": "ask_question" | "order_test" | "submit_diagnosis", | |
| "question": str or None, | |
| "test_name": str or None, | |
| "diagnosis": str or None, | |
| } | |
| """ | |
| response_lower = response.lower() | |
| # Try to find ACTION: directive | |
| if "action:" in response_lower: | |
| lines = response.split("\n") | |
| action_type = None | |
| question = None | |
| test_name = None | |
| diagnosis = None | |
| for i, line in enumerate(lines): | |
| if "action:" in line.lower(): | |
| action_part = line.split(":", 1)[1].strip().lower() | |
| if "question" in action_part: | |
| action_type = "ask_question" | |
| elif "test" in action_part: | |
| action_type = "order_test" | |
| elif "diagnosis" in action_part: | |
| action_type = "submit_diagnosis" | |
| if "question:" in line.lower(): | |
| question = line.split(":", 1)[1].strip() | |
| elif "test:" in line.lower(): | |
| test_name = line.split(":", 1)[1].strip() | |
| elif "diagnosis:" in line.lower(): | |
| diagnosis = line.split(":", 1)[1].strip() | |
| if action_type: | |
| return { | |
| "action_type": action_type, | |
| "question": question, | |
| "test_name": test_name, | |
| "diagnosis": diagnosis, | |
| } | |
| # Fallback: try to infer action type | |
| if "question" in response_lower or "ask" in response_lower: | |
| # Extract the question | |
| for line in response.split("\n"): | |
| if "?" in line: | |
| return { | |
| "action_type": "ask_question", | |
| "question": line.strip(), | |
| "test_name": None, | |
| "diagnosis": None, | |
| } | |
| if "test" in response_lower or "order" in response_lower: | |
| # Try to extract test name | |
| words = response.split() | |
| for i, word in enumerate(words): | |
| if "test" in word.lower() and i + 1 < len(words): | |
| test = words[i + 1] | |
| return { | |
| "action_type": "order_test", | |
| "question": None, | |
| "test_name": test, | |
| "diagnosis": None, | |
| } | |
| if "diagnos" in response_lower: | |
| # Try to extract diagnosis | |
| for word in response.split(): | |
| if len(word) > 3: # Filter out small words | |
| return { | |
| "action_type": "submit_diagnosis", | |
| "question": None, | |
| "test_name": None, | |
| "diagnosis": response.strip(), | |
| } | |
| return None | |
| def build_conversation_history(episode_history: List[Dict]) -> List[Dict]: | |
| """Build conversation history for multi-turn interaction.""" | |
| conversation = [ | |
| { | |
| "role": "system", | |
| "content": create_system_prompt(), | |
| } | |
| ] | |
| for turn in episode_history: | |
| # Add assistant message | |
| if turn.get("agent_action"): | |
| conversation.append({ | |
| "role": "assistant", | |
| "content": turn["agent_action"], | |
| }) | |
| # Add environment feedback | |
| if turn.get("environment_feedback"): | |
| conversation.append({ | |
| "role": "user", | |
| "content": turn["environment_feedback"], | |
| }) | |
| return conversation | |
| # ============================================================================== | |
| # EPISODE EXECUTION | |
| # ============================================================================== | |
| async def run_episode_async( | |
| client: OpenAI, | |
| image_name: str, | |
| difficulty: str, | |
| task_name: str, | |
| ) -> Dict: | |
| """ | |
| Run a single episode with asyncio. | |
| Returns: { | |
| "task": task_name, | |
| "success": bool, | |
| "steps_taken": int, | |
| "total_reward": float, | |
| "episode_rewards": [float], | |
| "final_diagnosis_accuracy": float, | |
| } | |
| """ | |
| log_start(task_name, "medical_diagnostic_env", MODEL_NAME) | |
| # Reset environment | |
| async with DiagnosticEnv.from_docker_image(image_name=image_name, base_url=ENV_URL) as env: | |
| obs_result = await env.reset(difficulty=difficulty) | |
| obs = obs_result.observation if hasattr(obs_result, 'observation') else obs_result | |
| episode_history = [] | |
| episode_rewards = [] | |
| step_count = 0 | |
| error_occurred = False | |
| # Initial environment message | |
| initial_message = f"Patient presentation: {obs.message}" | |
| while step_count < MAX_STEPS and not obs.done: | |
| step_count += 1 | |
| # Build conversation with history | |
| conversation = [ | |
| { | |
| "role": "system", | |
| "content": create_system_prompt(), | |
| } | |
| ] | |
| # Add conversation history | |
| for turn in episode_history: | |
| if turn.get("agent_thought"): | |
| conversation.append({ | |
| "role": "assistant", | |
| "content": f"Thinking: {turn['agent_thought']}\nAction: {turn['agent_action']}", | |
| }) | |
| if turn.get("environment_response"): | |
| conversation.append({ | |
| "role": "user", | |
| "content": turn["environment_response"], | |
| }) | |
| # Add current observation if first step | |
| if step_count == 1: | |
| conversation.append({ | |
| "role": "user", | |
| "content": initial_message, | |
| }) | |
| try: | |
| # Get LLM response | |
| response = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=conversation, | |
| temperature=TEMPERATURE, | |
| max_tokens=MAX_TOKENS, | |
| ) | |
| llm_response = response.choices[0].message.content | |
| # Extract action from response | |
| action_dict = extract_action_from_response(llm_response) | |
| if not action_dict: | |
| error_msg = "Could not parse action from response" | |
| log_step(step_count, "parse_error", 0.0, False, error_msg) | |
| error_occurred = True | |
| break | |
| # Create action | |
| action = DiagnosticAction( | |
| action_type=action_dict["action_type"], | |
| question=action_dict.get("question"), | |
| test_name=action_dict.get("test_name"), | |
| diagnosis=action_dict.get("diagnosis"), | |
| ) | |
| # Execute action with short action string for logging | |
| action_str = f"{action.action_type}" | |
| if action.question: | |
| action_str += f"('{action.question[:30]}...')" | |
| elif action.test_name: | |
| action_str += f"('{action.test_name}')" | |
| elif action.diagnosis: | |
| action_str += f"('{action.diagnosis[:40]}...')" | |
| # Take step in environment | |
| step_result = await env.step(action) | |
| obs = step_result.observation if hasattr(step_result, 'observation') else step_result | |
| reward = obs.reward or 0.0 | |
| episode_rewards.append(reward) | |
| log_step(step_count, action_str, reward, obs.done, None) | |
| # Store in history | |
| episode_history.append({ | |
| "agent_thought": llm_response[:100], | |
| "agent_action": action_str, | |
| "environment_response": obs.message[:200], | |
| }) | |
| except Exception as e: | |
| error_msg = str(e)[:100] | |
| log_step(step_count, "error", 0.0, True, error_msg) | |
| error_occurred = True | |
| break | |
| # Get final state | |
| try: | |
| state = await env.state() | |
| final_accuracy = state.final_accuracy if hasattr(state, 'final_accuracy') else 0.0 | |
| except: | |
| final_accuracy = 0.0 | |
| # Calculate results | |
| total_reward = sum(episode_rewards) | |
| success = obs.done and final_accuracy > 0.3 | |
| log_end(success, step_count, final_accuracy, episode_rewards) | |
| return { | |
| "task": task_name, | |
| "success": success, | |
| "steps_taken": step_count, | |
| "total_reward": total_reward, | |
| "episode_rewards": episode_rewards, | |
| "final_diagnosis_accuracy": final_accuracy, | |
| } | |
| # ============================================================================== | |
| # MAIN ORCHESTRATION | |
| # ============================================================================== | |
| async def run_all_tasks() -> Dict: | |
| """Run all 3 difficulty levels and report overall results.""" | |
| if not API_KEY: | |
| print("ERROR: API key not found. Set API_KEY to the provided API key.", flush=True) | |
| return {} | |
| if not API_BASE_URL: | |
| print("ERROR: API_BASE_URL is not set. Set API_BASE_URL to the provided LLM proxy endpoint.", flush=True) | |
| return {} | |
| if not ENV_URL: | |
| print("ERROR: ENV_URL is not set. Set ENV_URL to the environment WebSocket URL.", flush=True) | |
| return {} | |
| # Initialize OpenAI client | |
| client = OpenAI( | |
| api_key=os.environ["API_KEY"], | |
| base_url=os.environ["API_BASE_URL"], | |
| ) | |
| results = { | |
| "timestamp": None, | |
| "model": MODEL_NAME, | |
| "environment": "medical_diagnostic_env", | |
| "tasks_completed": 0, | |
| "task_results": [], | |
| "overall_score": 0.0, | |
| } | |
| # Run each task | |
| for i, (task_name, difficulty) in enumerate(zip(TASK_NAMES, DIFFICULTY_LEVELS)): | |
| print(f"\n--- Task {i+1}/3: {difficulty} difficulty ---", flush=True) | |
| try: | |
| result = await run_episode_async( | |
| client, | |
| LOCAL_IMAGE_NAME, | |
| difficulty=difficulty, | |
| task_name=task_name, | |
| ) | |
| results["task_results"].append(result) | |
| results["tasks_completed"] += 1 | |
| except Exception as e: | |
| print(f"ERROR in task {task_name}: {str(e)}", flush=True) | |
| # Calculate overall score | |
| if results["tasks_completed"] > 0: | |
| accuracies = [r["final_diagnosis_accuracy"] for r in results["task_results"]] | |
| results["overall_score"] = sum(accuracies) / len(accuracies) | |
| # Print summary | |
| print("\n" + "="*60, flush=True) | |
| print(f"Baseline Inference Complete", flush=True) | |
| print(f"Tasks completed: {results['tasks_completed']}/3", flush=True) | |
| print(f"Overall diagnostic accuracy: {results['overall_score']:.3f}", flush=True) | |
| print("="*60, flush=True) | |
| return results | |
| def main(): | |
| """Entry point.""" | |
| # Run async loop | |
| results = asyncio.run(run_all_tasks()) | |
| if __name__ == "__main__": | |
| main() | |