#!/usr/bin/env python3 """ Baseline inference script for the Data Cleaning Environment. Uses the OpenAI API client to run an LLM agent against the environment for all 3 tasks (easy, medium, hard) and prints reproducible scores. Usage: # Set your API key export OPENAI_API_KEY=sk-... # Run against local server (default) python baseline.py # Run against a deployed HF Space python baseline.py --base-url https://your-username-data-cleaning-env.hf.space Requirements: pip install openai requests """ import argparse import json import os import sys import requests try: from dotenv import load_dotenv load_dotenv() except ImportError: pass from huggingface_hub import InferenceClient try: from openai import OpenAI except ImportError: print("openai package not found. Install with: pip install openai") sys.exit(1) try: from openai import OpenAI except ImportError: print("openai package not found. Install with: pip install openai") sys.exit(1) # --------------------------------------------------------------------------- # Deterministic rule-based agent (no LLM needed for baseline) # --------------------------------------------------------------------------- RULE_POLICIES = { "easy": ["impute_mean", "impute_mode", "drop_missing_rows"], "medium": ["remove_duplicates", "fix_type_errors", "drop_missing_rows"], "hard": [ "fill_quantity_mean", "drop_missing_rows", "remove_duplicates", "fix_type_errors", "remove_outliers", "normalize_text", ], } def run_rule_baseline(base_url: str) -> dict[str, float]: """Run deterministic rule-based baseline — no LLM required.""" scores = {} for task in ["easy", "medium", "hard"]: # Reset resp = requests.post(f"{base_url}/reset", json={"task": task}, timeout=10) resp.raise_for_status() # Apply each operation in the policy for op in RULE_POLICIES[task]: resp = requests.post( f"{base_url}/step", json={"action": {"operation": op}}, timeout=10, ) resp.raise_for_status() data = resp.json() if data.get("done"): break # Grade resp = requests.post(f"{base_url}/grader", timeout=10) resp.raise_for_status() result = resp.json() scores[task] = result["score"] return scores # --------------------------------------------------------------------------- # LLM agent (uses OpenAI API) # --------------------------------------------------------------------------- SYSTEM_PROMPT = """You are a data cleaning agent. You will be shown a dirty dataset as a text table and must choose ONE cleaning operation to apply per turn. Available operations: impute_mean – Fill numeric missing values with the column mean impute_mode – Fill categorical missing values with the most common value drop_missing_rows – Drop all rows that have any missing value remove_duplicates – Remove exact duplicate rows fix_type_errors – Coerce non-numeric values in numeric columns to float remove_outliers – Drop rows where price <= 0 or price >= 500 normalize_text – Strip whitespace and title-case all string columns fill_quantity_mean – Fill missing quantity values with the column mean Respond ONLY with a JSON object like: {"operation": "remove_duplicates"} or with an optional column: {"operation": "impute_mean", "column": "age"} No explanation. JSON only.""" def run_llm_baseline(base_url: str, api_key: str, max_steps: int = 10) -> dict[str, float]: """Run an LLM agent (GPT-4o-mini) against the environment.""" client = OpenAI(api_key=api_key) # client = InferenceClient(api_key=api_key) scores = {} for task in ["easy", "medium", "hard"]: print(f"\n [LLM] Task: {task}") resp = requests.post(f"{base_url}/reset", json={"task": task}, timeout=10) resp.raise_for_status() obs = resp.json() for step in range(max_steps): current_text = obs["observation"].get("current_text", "") metadata = obs["observation"].get("metadata", {}) quality = metadata.get("quality_score", "?") valid_ops = metadata.get("valid_operations", []) user_msg = ( f"Current dataset (quality score: {quality}):\n" f"{current_text}\n\n" f"Valid operations: {valid_ops}\n" f"Choose ONE operation to improve data quality." ) completion = client.chat.completions.create( model="gpt-4o-mini", messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_msg}, ], temperature=0, max_tokens=64, ) raw = completion.choices[0].message.content.strip() try: action = json.loads(raw) except json.JSONDecodeError: # Extract JSON from response if wrapped in markdown import re match = re.search(r"\{.*\}", raw, re.DOTALL) action = json.loads(match.group()) if match else {"operation": "drop_missing_rows"} print(f" step {step+1}: {action}") resp = requests.post( f"{base_url}/step", json={"action": action}, timeout=10, ) resp.raise_for_status() obs = resp.json() if obs.get("done"): print(f" Episode done at step {step+1}") break # Grade resp = requests.post(f"{base_url}/grader", timeout=10) resp.raise_for_status() result = resp.json() scores[task] = result["score"] print(f" [LLM] {task} score: {scores[task]}") return scores # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser(description="Data Cleaning Env baseline script") parser.add_argument( "--base-url", default="http://localhost:8000", help="Base URL of the running environment server", ) parser.add_argument( "--mode", choices=["rule", "llm", "both"], default="rule", help="Baseline mode: 'rule' (no API key needed), 'llm' (needs OPENAI_API_KEY), 'both'", ) args = parser.parse_args() base_url = args.base_url.rstrip("/") # Health check try: r = requests.get(f"{base_url}/health", timeout=5) r.raise_for_status() print(f"✓ Server healthy at {base_url}") except Exception as e: print(f"✗ Cannot reach server at {base_url}: {e}") sys.exit(1) # ── Rule-based baseline (always runs) ────────────────────────────────── if args.mode in ("rule", "both"): print("\n=== Rule-based Baseline ===") try: scores = run_rule_baseline(base_url) print("\nScores:") for task, score in scores.items(): bar = "█" * int(score * 20) print(f" {task:<8} {score:.4f} {bar}") print(f"\n Mean: {sum(scores.values()) / len(scores):.4f}") except Exception as e: print(f"Rule baseline failed: {e}") # ── LLM baseline ─────────────────────────────────────────────────────── if args.mode in ("llm", "both"): api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") if not api_key: print("\nSkipping LLM baseline: OPENAI_API_KEY not set.") else: print("\n=== LLM Baseline (gpt-4o-mini) ===") try: scores = run_llm_baseline(base_url, api_key) print("\nScores:") for task, score in scores.items(): bar = "█" * int(score * 20) print(f" {task:<8} {score:.4f} {bar}") print(f"\n Mean: {sum(scores.values()) / len(scores):.4f}") except Exception as e: print(f"LLM baseline failed: {e}") if __name__ == "__main__": main()