Spaces:
Sleeping
Sleeping
| import os | |
| import requests | |
| from openai import OpenAI | |
| from typing import List, Optional | |
| # Load .env file if present (for local development) | |
| try: | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| except ImportError: | |
| pass # dotenv not required in production | |
| # ── Config ────────────────────────────────────────────────────────────── | |
| # Mandatory environment variables per hackathon requirements | |
| API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:7860") # Default to local env | |
| MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini") # Default model | |
| HF_TOKEN = os.getenv("HF_TOKEN") # No default - required | |
| # OpenAI client configured with environment variables | |
| client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL) | |
| ENV_URL = os.getenv("ENV_URL", "http://localhost:7860") | |
| BENCHMARK = "data-cleaning-env" | |
| TASKS = ["easy", "medium", "hard"] | |
| # ── Structured Logging Functions (MANDATORY FORMAT) ───────────────────── | |
| def log_start(task: str, env: str, model: str) -> None: | |
| """Emit [START] line at episode begin.""" | |
| print(f"[START] task={task} env={env} model={model}", flush=True) | |
| def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: | |
| """Emit [STEP] line after each env.step() with exact format.""" | |
| error_val = error if error else "null" | |
| done_val = str(done).lower() # Must be lowercase: true/false | |
| 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: | |
| """Emit [END] line after episode completes. Always called even on exception.""" | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) | |
| success_val = str(success).lower() # Must be lowercase: true/false | |
| print( | |
| f"[END] success={success_val} steps={steps} score={score:.3f} rewards={rewards_str}", | |
| flush=True, | |
| ) | |
| def ask_llm(observation: dict) -> dict: | |
| """Ask the LLM what cleaning action to take.""" | |
| # Get first column as fallback | |
| columns = list(observation['current_data'].keys()) | |
| first_column = columns[0] if columns else "id" | |
| prompt = f""" | |
| You are a data cleaning agent. You receive a row of data and must clean it. | |
| Current row data: {observation['current_data']} | |
| Issues detected: {observation['issues_detected']} | |
| Legal actions: {observation['legal_actions']} | |
| Available columns: {columns} | |
| Choose ONE action from the legal actions list. | |
| Respond in this exact JSON format with ALL three fields always present: | |
| {{ | |
| "action_type": "fill_missing", | |
| "column": "age", | |
| "value": 30 | |
| }} | |
| Rules: | |
| - action_type must be from legal_actions list | |
| - column must ALWAYS be one of: {columns} | |
| - value is required for fill_missing, optional for others (use null if not needed) | |
| - For fill_missing: provide a sensible default value | |
| - For skip: still include column field using first available column | |
| - ALWAYS include all three fields: action_type, column, value | |
| """ | |
| response = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role": "user", "content": prompt}], | |
| temperature=0.0, | |
| ) | |
| import json | |
| text = response.choices[0].message.content.strip() | |
| text = text.replace("```json", "").replace("```", "").strip() | |
| parsed = json.loads(text) | |
| # Safety fallbacks | |
| if "column" not in parsed or not parsed["column"]: | |
| parsed["column"] = first_column | |
| if "value" not in parsed: | |
| parsed["value"] = None | |
| if parsed.get("column") not in columns: | |
| parsed["column"] = first_column | |
| return parsed | |
| def run_task(task_name: str) -> float: | |
| """Run one full episode for a task and return score in [0, 1] range.""" | |
| rewards: List[float] = [] | |
| steps_taken = 0 | |
| score = 0.0 | |
| success = False | |
| log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME) | |
| try: | |
| # Reset environment | |
| obs = requests.post( | |
| f"{ENV_URL}/reset", | |
| json={"task_name": task_name} | |
| ).json() | |
| # Run episode until done | |
| while not obs["done"]: | |
| steps_taken += 1 | |
| # Get action from LLM | |
| error_msg = None | |
| try: | |
| action = ask_llm(obs) | |
| action_str = f"{action['action_type']}('{action['column']}',{action['value']})" | |
| except Exception as e: | |
| error_msg = str(e) | |
| # Fallback action on LLM error | |
| action = { | |
| "action_type": "skip", | |
| "column": list(obs["current_data"].keys())[0] if obs["current_data"] else "id", | |
| "value": None | |
| } | |
| action_str = f"skip('{action['column']}',null)" | |
| # Take step in environment | |
| try: | |
| obs = requests.post( | |
| f"{ENV_URL}/step", | |
| json=action | |
| ).json() | |
| reward = obs.get("reward", 0.0) | |
| done = obs.get("done", False) | |
| except Exception as e: | |
| error_msg = str(e) | |
| reward = 0.0 | |
| done = True | |
| obs = {"done": True, "reward": 0.0} | |
| rewards.append(reward) | |
| log_step(step=steps_taken, action=action_str, reward=reward, done=done, error=error_msg) | |
| if done: | |
| break | |
| # Get final score from environment grader (must be in [0, 1]) | |
| try: | |
| state = requests.get(f"{ENV_URL}/state").json() | |
| score = state.get("score", 0.0) | |
| # Clamp score to [0, 1] range as required | |
| score = min(max(score, 0.0), 1.0) | |
| # Success if score meets reasonable threshold | |
| success = score >= 0.5 | |
| except Exception as e: | |
| score = 0.0 | |
| success = False | |
| except Exception as e: | |
| # Even on exception, we must emit [END] | |
| pass | |
| finally: | |
| # Always emit [END] line, even on exception | |
| log_end(success=success, steps=steps_taken, score=score, rewards=rewards) | |
| return score | |
| def main(): | |
| """Main inference script entry point.""" | |
| # Health check - verify server is running | |
| try: | |
| health = requests.get(f"{ENV_URL}/health", timeout=5).json() | |
| if health.get("status") != "ok": | |
| print(f"[ERROR] Server not healthy at {ENV_URL}", flush=True) | |
| return | |
| except Exception as e: | |
| print(f"[ERROR] Server not running at {ENV_URL}: {e}", flush=True) | |
| return | |
| # Run inference on all tasks | |
| scores = {} | |
| for task in TASKS: | |
| scores[task] = run_task(task) | |
| # Print summary (not part of structured output, just for convenience) | |
| print(f"\n{'='*60}", flush=True) | |
| print("BASELINE RESULTS SUMMARY", flush=True) | |
| print(f"{'='*60}", flush=True) | |
| for task, score in scores.items(): | |
| print(f" {task:10s}: {score:.3f}", flush=True) | |
| avg = sum(scores.values()) / len(scores) if scores else 0.0 | |
| print(f" {'average':10s}: {avg:.3f}", flush=True) | |
| print(f"{'='*60}", flush=True) | |
| if __name__ == "__main__": | |
| main() |