""" Inference Script — Data Cleaning OpenEnv ========================================= Runs a baseline LLM agent against all 3 tasks and produces reproducible scores. Required environment variables: API_BASE_URL — LLM API endpoint MODEL_NAME — model identifier HF_TOKEN — Hugging Face API key """ import os import sys import json import time from typing import List, Dict, Any from openai import OpenAI # ───────────────────────────────────────── # CONFIG # ───────────────────────────────────────── API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1") API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") or os.getenv("API_KEY", "") MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini") MAX_STEPS = 15 TEMPERATURE = 0.1 MAX_TOKENS = 400 VALID_TASKS = [ "easy_dedup_rename", "medium_missing_dtype", "hard_full_pipeline", "expert_sales_pipeline" ] SYSTEM_PROMPT = """ You are an expert data cleaning agent. You will receive information about a dirty dataset and must clean it step by step using the available operations. Available operations: - remove_duplicates: {"operation": "remove_duplicates", "parameters": {}} - fill_missing: {"operation": "fill_missing", "parameters": {"strategy": "mean|median|mode"}} - fix_dtype: {"operation": "fix_dtype", "parameters": {"dtype": "auto"}} - remove_outliers: {"operation": "remove_outliers", "parameters": {"method": "iqr"}} - rename_columns: {"operation": "rename_columns", "parameters": {}} - validate_schema: {"operation": "validate_schema", "parameters": {}} - finish: {"operation": "finish", "parameters": {}} Rules: 1. Always respond with ONLY a valid JSON object 2. No explanations, no markdown, no extra text 3. Just the JSON action object 4. Call finish when you think the dataset is clean Example response: {"operation": "remove_duplicates", "parameters": {}} """ def build_user_prompt(observation: Dict[str, Any]) -> str: return f""" Task: {observation.get("task_description", "")} Step: {observation.get("step", 0)} Last message: {observation.get("message", "")} Current dataset info: - Shape: {observation.get("shape", [])} - Columns: {observation.get("columns", [])} - Duplicate rows: {observation.get("duplicate_count", 0)} - Missing values: {observation.get("missing_values", {})} - Data types: {observation.get("dtypes", {})} Sample rows (first 3): {json.dumps(observation.get("sample_rows", []), indent=2)} Available operations: {observation.get("available_operations", [])} What is your next action? Respond with JSON only. """ def parse_action(response_text: str) -> Dict[str, Any]: """Parse LLM response into action dict.""" text = response_text.strip() # Remove markdown if present if "```json" in text: text = text.split("```json")[1].split("```")[0].strip() elif "```" in text: text = text.split("```")[1].split("```")[0].strip() try: action = json.loads(text) if "operation" not in action: return {"operation": "finish", "parameters": {}} if "parameters" not in action: action["parameters"] = {} return action except json.JSONDecodeError: # Try to find JSON in text import re match = re.search(r"\{.*\}", text, re.DOTALL) if match: try: return json.loads(match.group()) except Exception: pass return {"operation": "finish", "parameters": {}} def run_task( client: OpenAI, env_module, task_id: str ) -> Dict[str, Any]: """Run one full episode for a task.""" print(f"\n{'='*50}") print(f" Task: {task_id}") print(f"{'='*50}") # Import here to use local environment from environment import DataCleaningEnv from models import Action env = DataCleaningEnv(task_id=task_id) # Reset result = env.reset() obs = result.observation.model_dump() done = result.done step = 0 rewards = [] actions_taken = [] print(f" Description: {obs.get('task_description', '')[:80]}...") print(f" Initial shape: {obs.get('shape', [])}") print(f" Duplicates: {obs.get('duplicate_count', 0)}") print(f" Missing: {sum(obs.get('missing_values', {}).values())}") while not done and step < MAX_STEPS: step += 1 # Build prompt user_prompt = build_user_prompt(obs) messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt} ] # Call LLM try: completion = client.chat.completions.create( model=MODEL_NAME, messages=messages, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False ) response_text = completion.choices[0].message.content or "" except Exception as e: print(f" [Step {step}] LLM error: {e}") response_text = '{"operation": "finish", "parameters": {}}' # Parse action action_dict = parse_action(response_text) action = Action( operation=action_dict.get("operation", "finish"), parameters=action_dict.get("parameters", {}) ) print(f" [Step {step}] Action: {action.operation} {action.parameters}") # Step environment result = env.step(action) obs = result.observation.model_dump() done = result.done reward = result.reward.total rewards.append(reward) actions_taken.append(action.operation) print(f" Reward: {reward:.4f} | Message: {obs.get('message', '')[:60]}") if done: print(f" Episode done at step {step}") break # Small delay to avoid rate limiting time.sleep(0.5) final_reward = rewards[-1] if rewards else 0.0 print(f"\n Final Score: {final_reward:.4f}") print(f" Steps taken: {step}") print(f" Actions: {actions_taken}") return { "task_id": task_id, "final_score": round(final_reward, 4), "steps": step, "rewards": rewards, "actions": actions_taken, "done": done } def main(): print("\n" + "="*50) print(" Data Cleaning OpenEnv — Baseline Inference") print("="*50) print(f" Model: {MODEL_NAME}") print(f" API URL: {API_BASE_URL}") print(f" Tasks: {VALID_TASKS}") print("="*50) # Validate config if not API_KEY: print("ERROR: HF_TOKEN or API_KEY not set") sys.exit(1) if not MODEL_NAME: print("ERROR: MODEL_NAME not set") sys.exit(1) # Init client client = OpenAI( base_url=API_BASE_URL, api_key=API_KEY ) # Run all tasks all_results = [] for task_id in VALID_TASKS: try: result = run_task(client, None, task_id) all_results.append(result) except Exception as e: print(f" ERROR on task {task_id}: {e}") all_results.append({ "task_id": task_id, "final_score": 0.0, "error": str(e) }) # Summary print("\n" + "="*50) print(" FINAL RESULTS SUMMARY") print("="*50) total_score = 0.0 for r in all_results: score = r.get("final_score", 0.0) total_score += score status = "ERROR" if "error" in r else "OK" print(f" {r['task_id']:<30} Score: {score:.4f} [{status}]") avg_score = total_score / len(all_results) print(f"\n Average Score: {avg_score:.4f}") print("="*50) # Save results locally output = { "model": MODEL_NAME, "tasks": all_results, "average_score": round(avg_score, 4) } with open("baseline_results.json", "w") as f: json.dump(output, f, indent=2) print("\n Results saved to baseline_results.json") # Auto submit to leaderboard HF_SPACE_URL = os.getenv( "HF_SPACE_URL", "https://thorodin103-data-cleaning-openenv.hf.space" ) print("\n Submitting scores to leaderboard...") try: import urllib.request for r in all_results: if "error" not in r: entry = { "model_name": MODEL_NAME, "task_id": r["task_id"], "score": r["final_score"], "steps": r["steps"] } data = json.dumps(entry).encode("utf-8") req = urllib.request.Request( f"{HF_SPACE_URL}/leaderboard/submit", data=data, headers={"Content-Type": "application/json"}, method="POST" ) with urllib.request.urlopen(req, timeout=10) as resp: result = json.loads(resp.read()) print(f" ✅ Submitted {r['task_id']}: {r['final_score']}") print(" Leaderboard updated!") print(f" View at: {HF_SPACE_URL}/ui") except Exception as e: print(f" ⚠️ Leaderboard submit failed: {e}") print(" Scores saved locally in baseline_results.json") print(" Done!") if __name__ == "__main__": main()