""" Inference Script — DataCleanEnv =================================== MANDATORY - Before submitting, ensure the following variables are defined in your environment configuration: API_BASE_URL The API endpoint for the LLM. MODEL_NAME The model identifier to use for inference. HF_TOKEN Your Hugging Face / API key. - The inference script must be named `inference.py` and placed in the root directory of the project - Participants must use OpenAI Client for all LLM calls using above variables This script emits exactly these stdout line types: - [START] ... - [STEP] ... (one per step) - [END] ... (always) """ from __future__ import annotations import json import os import re import textwrap from typing import Any, Dict, List, Optional from openai import OpenAI from client import DataCleanEnv from models import DataCleanAction, DataCleanObservation # --------------------------------------------------------------------------- # Config — read at import time but API vars re-read in main() to catch # late-injected env vars from the validator. # --------------------------------------------------------------------------- BENCHMARK_URL = os.getenv( "BENCHMARK_URL", os.getenv("ENV_URL", "https://tns-openenv-data-clean.hf.space"), ) BENCHMARK = os.getenv("BENCHMARK", "data_clean_env") TASKS = ["customer_contacts", "sales_records", "employee_records", "financial_transactions"] # --------------------------------------------------------------------------- # Structured logging # --------------------------------------------------------------------------- def log_start(task: str, env: str, model: str) -> None: print(f"[START] task={task} env={env} model={model}", flush=True) def log_step(step: int, action: str, reward: float, done: bool, error: str | None) -> None: err = _single_line(error) if error else "null" print( f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={err}", flush=True, ) def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None: reward_csv = ",".join(f"{r:.2f}" for r in rewards) if rewards else "0.00" print( f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={reward_csv}", flush=True, ) def _single_line(text: str | None) -> str: return (text or "").replace("\n", " ").replace("\r", " ").strip() # --------------------------------------------------------------------------- # System prompt — Conservative plan-then-execute # --------------------------------------------------------------------------- PLANNING_PROMPT = textwrap.dedent("""\ You are a data quality analyst. You will receive a dataset, inspection results, and validation rules. Produce a PRECISE fix plan as a JSON array. CRITICAL RULES: - ONLY fix cells that inspection flagged as having issues (suspicious values, wrong format, etc.) - If inspection shows "Issues remaining in this column: 0", do NOT touch that column - Do NOT fix cells that already have correct values - Each wrong fix costs -0.05 penalty. Be CONSERVATIVE. - For duplicate rows (two identical rows), use "delete" on the LATER row - List all "fix" actions first, then all "delete" actions - Delete from highest row index to lowest VALIDATION RULES: - Emails: user@domain.tld (no [at], no @@, no spaces, no missing domain) - Phones: digits and dashes only, 10+ digits (no letters) - Dates: YYYY-MM-DD only (not MM/DD/YYYY, not slashes, valid calendar date) - Empty values: provide a reasonable non-empty value - Negative numbers: use the absolute value (make positive) - Outliers: fix to a reasonable mid-range value within the stated bounds - Inconsistent format: use the EXACT canonical form listed in the task description - Whitespace: trim leading/trailing, collapse double spaces to single - Salaries: must be $20,000-$500,000 - Performance scores: must be 0.0-10.0 - Currency: must be ISO code (USD, EUR, GBP, JPY, CAD) - Reviewer IDs: approved/flagged status requires a reviewer_id OUTPUT: Respond with ONLY a JSON array. No explanation, no markdown, no text before or after. EXAMPLE for a 3-issue dataset: [{"action":"fix","row":3,"column":"email","value":"alice@mail.com"},{"action":"fix","row":7,"column":"phone","value":"555-012-3408"},{"action":"delete","row":14}] """) # --------------------------------------------------------------------------- # JSON plan extraction # --------------------------------------------------------------------------- def extract_json_plan(text: str) -> Optional[List[Dict]]: if not text: return None text = re.sub(r"^```(?:json)?\s*\n?", "", text.strip()) text = re.sub(r"\n?```\s*$", "", text.strip()) try: plan = json.loads(text) if isinstance(plan, list): return plan except json.JSONDecodeError: pass match = re.search(r"\[[\s\S]*\]", text) if match: try: plan = json.loads(match.group()) if isinstance(plan, list): return plan except json.JSONDecodeError: pass return None def plan_to_command(action: Dict) -> Optional[str]: act_type = action.get("action", "") if act_type == "fix": row = action.get("row", 0) col = action.get("column", "") val = str(action.get("value", "")) return f'fix({row}, "{col}", "{val}")' elif act_type == "delete": row = action.get("row", 0) return f"delete({row})" return None # --------------------------------------------------------------------------- # Fallback: single-action extraction # --------------------------------------------------------------------------- ACTION_RE = re.compile(r"(inspect|fix|delete|submit)\s*\(", re.IGNORECASE) def extract_action(response_text: str) -> str: if not response_text: return "submit()" for line in response_text.strip().splitlines(): line = line.strip() if not line: continue line = re.sub(r"^```\w*\s*", "", line) line = re.sub(r"\s*```$", "", line) line = re.sub(r"^(?:action|next action)\s*[:\-]\s*", "", line, flags=re.IGNORECASE) if ACTION_RE.search(line): m = ACTION_RE.search(line) start = m.start() depth = 0 for i in range(start, len(line)): if line[i] == "(": depth += 1 elif line[i] == ")": depth -= 1 if depth == 0: return line[start : i + 1] return line[start:] + ")" return "submit()" # --------------------------------------------------------------------------- # Run a single task # --------------------------------------------------------------------------- def run_task(client: OpenAI, env, task_id: str, model_name: str = "") -> None: rewards: list[float] = [] step_count = 0 score = 0.0 success = False log_start(task=task_id, env=BENCHMARK, model=model_name) try: # --- Reset --- result = env.reset(task_id=task_id) obs = result.observation done = result.done if done: score = obs.current_score return total_issues = obs.total_issues # --- Phase 1: Auto-inspect all columns --- columns = [] for line in obs.column_info.strip().splitlines(): line = line.strip() if ":" in line: col_name = line.split(":")[0].strip() if col_name: columns.append(col_name) inspection_results = {} for col in columns: if done: break step_count += 1 cmd = f'inspect("{col}")' result = env.step(DataCleanAction(command=cmd)) obs = result.observation done = result.done reward = float(result.reward or 0.0) rewards.append(reward) log_step(step=step_count, action=cmd, reward=reward, done=done, error=None) inspection_results[col] = obs.feedback if done: score = obs.current_score success = score >= 0.5 return # --- Phase 1.5: Filter to only columns WITH issues --- flagged_inspections = {} for col, feedback in inspection_results.items(): m = re.search(r"Issues remaining in this column:\s*(\d+)", feedback) issue_count = int(m.group(1)) if m else 0 if issue_count > 0: flagged_inspections[col] = feedback for col, feedback in inspection_results.items(): if col not in flagged_inspections and "Suspicious:" in feedback: flagged_inspections[col] = feedback # --- Phase 2: Ask LLM to plan fixes --- if flagged_inspections: inspection_text = "\n\n".join( f"[{col}]\n{fb}" for col, fb in flagged_inspections.items() ) else: inspection_text = "(No specific column issues flagged. Check for duplicate rows.)" planning_message = ( f"Task: {obs.task_id} ({obs.difficulty})\n" f"Total issues to find and fix: {total_issues}\n\n" f"Task description:\n{obs.task_description}\n\n" f"Column definitions:\n{obs.column_info}\n\n" f"FLAGGED COLUMNS (only fix cells in these columns or duplicate rows):\n{inspection_text}\n\n" f"Current data:\n{obs.data_preview}\n\n" f"Produce a JSON array with EXACTLY the fixes needed. " f"Expected: around {total_issues} actions (fixes + deletes). " f"Do NOT produce more than {total_issues + 3} actions." ) try: completion = client.chat.completions.create( model=model_name, messages=[ {"role": "system", "content": PLANNING_PROMPT}, {"role": "user", "content": planning_message}, ], temperature=0.0, max_tokens=2000, stream=False, ) plan_text = completion.choices[0].message.content or "" except Exception as exc: # LLM error — submit immediately step_count += 1 cmd = "submit()" result = env.step(DataCleanAction(command=cmd)) obs = result.observation done = result.done reward = float(result.reward or 0.0) rewards.append(reward) log_step(step=step_count, action=cmd, reward=reward, done=True, error=_single_line(str(exc))) score = obs.current_score return plan = extract_json_plan(plan_text) # --- Sanity check: reject bloated plans --- if plan and len(plan) > total_issues + 5: plan = plan[:total_issues + 3] if not plan: # --- Fallback: single-action mode --- fallback_messages = [ {"role": "system", "content": ( "You are a data quality analyst. Respond with EXACTLY ONE command per turn.\n" "Commands: inspect(\"col\"), fix(row, \"col\", \"val\"), delete(row), submit()\n" "ONLY fix cells with actual issues. Do NOT fix correct data.\n" "Respond with ONLY the command." )}, {"role": "user", "content": planning_message}, ] remaining = obs.actions_remaining while not done and remaining > 0: try: comp = client.chat.completions.create( model=model_name, messages=fallback_messages, temperature=0.0, max_tokens=300, stream=False, ) resp_text = comp.choices[0].message.content or "" except Exception: resp_text = "submit()" action_cmd = extract_action(resp_text) fallback_messages.append({"role": "assistant", "content": action_cmd}) step_count += 1 result = env.step(DataCleanAction(command=action_cmd)) obs = result.observation done = result.done reward = float(result.reward or 0.0) rewards.append(reward) log_step(step=step_count, action=action_cmd, reward=reward, done=done, error=None) remaining = obs.actions_remaining if not done: fb = obs.feedback fallback_messages.append({"role": "user", "content": f"Result: {fb}\nFixed: {obs.issues_fixed}/{obs.total_issues}. Remaining steps: {remaining}."}) if len(fallback_messages) > 30: fallback_messages = [fallback_messages[0]] + fallback_messages[-28:] score = obs.current_score success = score >= 0.5 return # --- Phase 3: Execute plan --- remaining = obs.actions_remaining for action_item in plan: if done or remaining <= 1: break cmd = plan_to_command(action_item) if not cmd: continue step_count += 1 result = env.step(DataCleanAction(command=cmd)) obs = result.observation done = result.done reward = float(result.reward or 0.0) rewards.append(reward) log_step(step=step_count, action=cmd, reward=reward, done=done, error=None) remaining = obs.actions_remaining # --- Phase 4: Submit --- if not done: step_count += 1 cmd = "submit()" result = env.step(DataCleanAction(command=cmd)) obs = result.observation reward = float(result.reward or 0.0) rewards.append(reward) log_step(step=step_count, action=cmd, reward=reward, done=True, error=None) score = obs.current_score success = score >= 0.5 except Exception as exc: log_step(step=step_count + 1, action="error", reward=0.0, done=True, error=_single_line(str(exc))) success = False finally: log_end(success=success, steps=step_count, score=score, rewards=rewards) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> None: import sys # Read API vars fresh — validator injects these at runtime api_base_url = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1") api_key = os.environ.get("API_KEY") or os.environ.get("HF_TOKEN", "") model_name = os.environ.get("MODEL_NAME", "") # Debug: log config to stderr (never stdout — validator parses that) print(f"[CONFIG] API_BASE_URL={api_base_url}", file=sys.stderr, flush=True) print(f"[CONFIG] API_KEY={'set('+api_key[:8]+'...)' if api_key else 'EMPTY'}", file=sys.stderr, flush=True) print(f"[CONFIG] MODEL_NAME={model_name}", file=sys.stderr, flush=True) print(f"[CONFIG] BENCHMARK_URL={BENCHMARK_URL}", file=sys.stderr, flush=True) client = OpenAI(base_url=api_base_url, api_key=api_key) env_client = DataCleanEnv(base_url=BENCHMARK_URL) with env_client.sync() as env: for task_id in TASKS: run_task(client, env, task_id, model_name) if __name__ == "__main__": main()