Spaces:
Sleeping
Sleeping
| """ | |
| MedCodeRL - Baseline Inference Script | |
| Uses the OpenAI API client to run a model against the MedCodeRL environment. | |
| Reads API credentials from environment variables. | |
| Usage: | |
| export HF_TOKEN="your-key" | |
| export API_BASE_URL="https://api.openai.com/v1" | |
| export MODEL_NAME="gpt-4o-mini" | |
| python inference.py | |
| """ | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import time | |
| from typing import Optional | |
| from openai import OpenAI | |
| # Add project root to path so we can import the environment directly | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from server.my_env_environment import MyEnvironment, _load_task_cases | |
| from models import MedAction | |
| # ----- Configuration ----- | |
| API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1") | |
| MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini") | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| if not HF_TOKEN: | |
| print("ERROR: Set HF_TOKEN environment variable (Your Hugging Face / API key).") | |
| sys.exit(1) | |
| # Number of cases to evaluate per difficulty level | |
| CASES_PER_DIFFICULTY = int(os.environ.get("CASES_PER_DIFFICULTY", "5")) | |
| MAX_RETRIES = 2 | |
| def create_client() -> OpenAI: | |
| """Create an OpenAI-compatible client.""" | |
| return OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL) | |
| def extract_json_from_response(text: str) -> Optional[dict]: | |
| """Extract JSON from an LLM response, handling markdown code blocks.""" | |
| text = text.strip() | |
| if text.startswith("{"): | |
| try: | |
| return json.loads(text) | |
| except json.JSONDecodeError: | |
| pass | |
| patterns = [ | |
| r"```json\s*\n?(.*?)\n?\s*```", | |
| r"```\s*\n?(.*?)\n?\s*```", | |
| r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", | |
| ] | |
| for pattern in patterns: | |
| matches = re.findall(pattern, text, re.DOTALL) | |
| for match in matches: | |
| try: | |
| return json.loads(match) | |
| except json.JSONDecodeError: | |
| continue | |
| return None | |
| SYSTEM_PROMPT = """You are an expert medical coding and billing compliance specialist. | |
| Your role is to: | |
| 1. Analyze clinical documentation | |
| 2. Assign appropriate ICD-10-CM diagnosis codes | |
| 3. Assign appropriate CPT procedure codes | |
| 4. Make billing compliance decisions (approve, reject, or flag for review) | |
| 5. Identify compliance risk flags | |
| 6. Provide clinical reasoning for your decisions | |
| You have deep knowledge of: | |
| - ICD-10-CM coding guidelines and conventions | |
| - CPT coding and modifier usage | |
| - Medicare/Medicaid billing rules | |
| - Medical necessity requirements | |
| - Common compliance violations (upcoding, unbundling, fraudulent billing) | |
| - Clinical documentation integrity | |
| - Prior authorization requirements | |
| Always respond with ONLY a valid JSON object in the exact format requested. | |
| Be precise with your ICD-10 and CPT codes. | |
| Consider the clinical documentation, symptoms, treatments, and insurance type when making decisions. | |
| Flag any compliance concerns in risk_flags.""" | |
| def format_observation(obs) -> str: | |
| """Format an observation into a readable prompt for the LLM.""" | |
| parts = [ | |
| f"## Medical Coding Case: {obs.case_id}", | |
| f"**Difficulty:** {obs.difficulty}", | |
| f"**Visit Type:** {obs.visit_type}", | |
| f"**Provider Specialty:** {obs.provider_specialty}", | |
| "", | |
| "### Patient Information", | |
| f"- **Age:** {obs.patient_age} | **Sex:** {obs.patient_sex}", | |
| f"- **Insurance:** {obs.insurance_type}", | |
| f"- **Prior Authorization Required:** {'Yes' if obs.prior_auth_required else 'No'}", | |
| f"- **Treatment Cost Tier:** {obs.treatment_cost}", | |
| "", | |
| "### Clinical Note", | |
| obs.clinical_note, | |
| "", | |
| "### Symptoms", | |
| ", ".join(obs.symptoms) if obs.symptoms else "None reported", | |
| "", | |
| "### Treatments", | |
| ", ".join(obs.treatments) if obs.treatments else "None", | |
| ] | |
| if obs.comorbidities: | |
| parts += ["", "### Comorbidities", ", ".join(obs.comorbidities)] | |
| if obs.lab_results: | |
| parts += ["", "### Lab Results", obs.lab_results] | |
| if obs.medications: | |
| parts += ["", "### Current Medications", ", ".join(obs.medications)] | |
| return "\n".join(parts) | |
| ACTION_PROMPT = """ | |
| Based on the clinical case above, provide your medical coding and billing compliance assessment. | |
| You MUST respond with a valid JSON object containing exactly these fields: | |
| { | |
| "diagnosis_codes": ["<ICD-10 code(s)>"], | |
| "procedure_codes": ["<CPT code(s) if applicable, or empty list>"], | |
| "decision": "<approve|reject|review>", | |
| "confidence": <0.0 to 1.0>, | |
| "reasoning": "<15-500 character clinical justification>", | |
| "modifier_codes": ["<optional CPT modifiers, or empty list>"], | |
| "risk_flags": ["<compliance risk flags identified, or empty list>"] | |
| } | |
| Guidelines: | |
| - Use standard ICD-10-CM codes (e.g., J06.9 for upper respiratory infection) | |
| - Use standard CPT codes (5 digits, e.g., 99213 for office visit) | |
| - decision: "approve" if coding is appropriate, "reject" if non-compliant, "review" if ambiguous | |
| - confidence: your certainty (0.0 = unsure, 1.0 = certain) | |
| - reasoning: explain WHY you chose these codes and this decision | |
| - risk_flags: compliance risks (e.g., "upcoding_risk", "missing_documentation", "bundling_violation") | |
| IMPORTANT: Respond ONLY with the JSON object, no additional text. | |
| """ | |
| def call_llm(client: OpenAI, obs) -> Optional[dict]: | |
| """Call the LLM to get a coding decision for a clinical case.""" | |
| formatted = format_observation(obs) | |
| for attempt in range(MAX_RETRIES + 1): | |
| try: | |
| response = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": formatted + "\n\n" + ACTION_PROMPT}, | |
| ], | |
| temperature=0.1, | |
| max_tokens=800, | |
| ) | |
| content = response.choices[0].message.content | |
| if not content: | |
| print(f" [Attempt {attempt+1}] Empty response from LLM") | |
| continue | |
| action = extract_json_from_response(content) | |
| if action is None: | |
| print(f" [Attempt {attempt+1}] Failed to parse JSON from response") | |
| if attempt < MAX_RETRIES: | |
| time.sleep(1) | |
| continue | |
| # Sanitize fields | |
| if "diagnosis_codes" not in action or not isinstance(action["diagnosis_codes"], list): | |
| action["diagnosis_codes"] = [action["diagnosis_codes"]] if isinstance(action.get("diagnosis_codes"), str) else ["R69"] | |
| if "procedure_codes" not in action: | |
| action["procedure_codes"] = [] | |
| if isinstance(action.get("procedure_codes"), str): | |
| action["procedure_codes"] = [action["procedure_codes"]] | |
| if "decision" not in action: | |
| action["decision"] = "review" | |
| if "confidence" not in action: | |
| action["confidence"] = 0.5 | |
| action["confidence"] = max(0.0, min(1.0, float(action["confidence"]))) | |
| if "reasoning" not in action or len(str(action.get("reasoning", ""))) < 15: | |
| action["reasoning"] = "Medical coding assessment based on clinical documentation review and compliance guidelines." | |
| if "modifier_codes" not in action: | |
| action["modifier_codes"] = [] | |
| if "risk_flags" not in action: | |
| action["risk_flags"] = [] | |
| return action | |
| except Exception as e: | |
| print(f" [Attempt {attempt+1}] API error: {e}") | |
| if attempt < MAX_RETRIES: | |
| time.sleep(2 ** attempt) | |
| return None | |
| def get_fallback_action() -> dict: | |
| """Return a safe fallback action if LLM fails.""" | |
| return { | |
| "diagnosis_codes": ["R69"], | |
| "procedure_codes": ["99213"], | |
| "decision": "review", | |
| "confidence": 0.1, | |
| "reasoning": "Unable to obtain LLM response. Flagging for manual review as a safety measure.", | |
| "modifier_codes": [], | |
| "risk_flags": ["llm_failure"], | |
| } | |
| def run_evaluation(): | |
| """Run the baseline evaluation across all difficulty levels.""" | |
| print("=" * 70) | |
| print("MedCodeRL - Baseline Inference Script") | |
| print("=" * 70) | |
| print(f"API Base URL: {API_BASE_URL}") | |
| print(f"Model: {MODEL_NAME}") | |
| print(f"Cases per difficulty: {CASES_PER_DIFFICULTY}") | |
| print("=" * 70) | |
| client = create_client() | |
| env = MyEnvironment() | |
| all_scores = [] | |
| results_by_difficulty = {} | |
| for difficulty in ["easy", "medium", "hard"]: | |
| print(f"\n{'─' * 50}") | |
| print(f" Running {difficulty.upper()} tasks") | |
| print(f"{'─' * 50}") | |
| difficulty_scores = [] | |
| available = len(env._task_cases.get(difficulty, [])) | |
| num_cases = min(CASES_PER_DIFFICULTY, available) | |
| if num_cases == 0: | |
| print(f" No cases available for {difficulty}") | |
| continue | |
| for i in range(num_cases): | |
| # Reset environment for this difficulty | |
| obs = env.reset(task_id=difficulty) | |
| print(f"\n Case {i+1}/{num_cases}: {obs.case_id}") | |
| # Get LLM action | |
| action_dict = call_llm(client, obs) | |
| if action_dict is None: | |
| print(" ⚠ LLM failed, using fallback action") | |
| action_dict = get_fallback_action() | |
| # Build MedAction | |
| med_action = MedAction( | |
| diagnosis_codes=action_dict["diagnosis_codes"], | |
| procedure_codes=action_dict.get("procedure_codes", []), | |
| decision=action_dict["decision"], | |
| confidence=action_dict["confidence"], | |
| reasoning=action_dict["reasoning"], | |
| modifier_codes=action_dict.get("modifier_codes", []), | |
| risk_flags=action_dict.get("risk_flags", []), | |
| ) | |
| # Step the environment | |
| try: | |
| result_obs = env.step(med_action) | |
| score = result_obs.reward if result_obs.reward is not None else 0.0 | |
| difficulty_scores.append(score) | |
| all_scores.append(score) | |
| print(f" Score: {score:.4f}") | |
| print(f" Decision: {action_dict.get('decision', 'N/A')}") | |
| print(f" Diagnosis: {action_dict.get('diagnosis_codes', [])}") | |
| print(f" Procedure: {action_dict.get('procedure_codes', [])}") | |
| if result_obs.reward_breakdown: | |
| gc = result_obs.reward_breakdown.get("grade_components", {}) | |
| if gc: | |
| print(f" Components: diag={gc.get('diagnosis_accuracy', 0):.2f} " | |
| f"proc={gc.get('procedure_accuracy', 0):.2f} " | |
| f"dec={gc.get('decision_accuracy', 0):.2f}") | |
| pens = result_obs.reward_breakdown.get("penalties", {}) | |
| if pens: | |
| print(f" Penalties: {list(pens.keys())}") | |
| if result_obs.feedback: | |
| print(f" Feedback: {result_obs.feedback}") | |
| except Exception as e: | |
| print(f" ✗ Step failed: {e}") | |
| difficulty_scores.append(0.0) | |
| all_scores.append(0.0) | |
| # Rate limiting | |
| time.sleep(0.5) | |
| if difficulty_scores: | |
| avg = sum(difficulty_scores) / len(difficulty_scores) | |
| results_by_difficulty[difficulty] = { | |
| "scores": difficulty_scores, | |
| "average": round(avg, 4), | |
| "count": len(difficulty_scores), | |
| } | |
| print(f"\n {difficulty.upper()} Average: {avg:.4f} ({len(difficulty_scores)} cases)") | |
| # Final summary | |
| print(f"\n{'=' * 70}") | |
| print(" FINAL RESULTS") | |
| print(f"{'=' * 70}") | |
| for diff, result in results_by_difficulty.items(): | |
| print(f" {diff.upper():>8}: {result['average']:.4f} ({result['count']} cases)") | |
| if all_scores: | |
| overall = sum(all_scores) / len(all_scores) | |
| print(f"\n {'OVERALL':>8}: {overall:.4f} ({len(all_scores)} total cases)") | |
| else: | |
| overall = 0.0 | |
| print("\n No scores recorded.") | |
| print(f"{'=' * 70}") | |
| # Write results to file | |
| results_output = { | |
| "model": MODEL_NAME, | |
| "api_base_url": API_BASE_URL, | |
| "cases_per_difficulty": CASES_PER_DIFFICULTY, | |
| "results_by_difficulty": results_by_difficulty, | |
| "overall_score": round(overall, 4), | |
| "total_cases": len(all_scores), | |
| } | |
| with open("baseline_results.json", "w") as f: | |
| json.dump(results_output, f, indent=2) | |
| print(f"\nResults saved to baseline_results.json") | |
| return overall | |
| if __name__ == "__main__": | |
| score = run_evaluation() | |
| sys.exit(0 if score > 0 else 1) | |