fix:update infernnce to support validator

inference.py

@@ -1,13 +1,26 @@
 import os
-from openai import OpenAI, AuthenticationError
-from typing import Dict
-from env_server import TASKS, KernelOptimization_env, grade_episode
-from models import Action
-import json
 import sys
+from typing import List, Optional
+
 from dotenv import load_dotenv
+from openai import OpenAI
+
+from env_server import KernelOptimization_env, TASKS, grade_episode
+from models import Action
 
 load_dotenv()
+
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
+API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+TASK_NAME = os.getenv("TASK_ID", "vector_add_easy")
+BENCHMARK = "kernel_optimization"
+
+
+def one_line(text: str) -> str:
+    return " ".join((text or "").split())
+
+
 def extract_code(text: str) -> str:
     if "```" not in text:
         return text
@@ -18,56 +31,101 @@ def extract_code(text: str) -> str:
         return chunk.split("\n", 1)[1]
     return chunk
 
-def choose_action(client: OpenAI, model: str, observation: Dict) -> Action:
-    prompt = f"""Optimize this CUDA kernel.
-    Task: {observation['task_name']}
-    Pending checks: {observation['pending_checks']}
-    Baseline:
-    {observation['baseline_code']}
-    Current best speedup: {observation['current_best_speedup']}x
-    Return only optimized CUDA code.
-"""
-    response = client.chat.completions.create(
-        model=model,
+
+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: Optional[str]) -> None:
+    error_val = one_line(error) if error else "null"
+    done_val = str(done).lower()
+    action_val = one_line(action)
+    print(
+        f"[STEP] step={step} action={action_val} reward={reward:.2f} done={done_val} error={error_val}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
+    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+    print(
+        f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}",
+        flush=True,
+    )
+
+
+def choose_action(client: OpenAI, observation: dict) -> Action:
+    prompt = (
+        "Optimize this CUDA kernel.\n"
+        f"Task: {observation['task_name']}\n"
+        f"Pending checks: {observation['pending_checks']}\n"
+        f"Current code:\n{observation['current_best_code']}\n"
+        "Return only optimized CUDA code."
+    )
+    completion = client.chat.completions.create(
+        model=MODEL_NAME,
         temperature=0.0,
         messages=[
             {"role": "system", "content": "You are a CUDA optimization expert. Return code only."},
             {"role": "user", "content": prompt},
         ],
     )
-    text = (response.choices[0].message.content or "").strip()
-    code = extract_code(text).strip() or observation["current_best_code"]
+    content = (completion.choices[0].message.content or "").strip()
+    code = extract_code(content).strip() or observation["current_best_code"]
     return Action(optimized_code=code, strategy="llm_proposed")
 
-def run_task(client: OpenAI, model: str, task_id: str) -> float:
+
+def main() -> int:
+    task_id = TASK_NAME if TASK_NAME in TASKS else "vector_add_easy"
     env = KernelOptimization_env()
-    obs = env.reset(task_id=task_id)["observation"]
-    done = False
-    while not done:
-        action = choose_action(client, model, obs)
-        step_result = env.step(action)
-        obs = step_result.observation.model_dump()
-        done = step_result.done
-    return grade_episode(task_id, env.state.completed_checks, env.state.best_speedup, env.state.step_count, env.state.max_steps)
-def main()->int:
-    if not os.getenv("OPENAI_API_KEY"):
-        print("openai key not set")
-
-    model =os.getenv("MODEL_NAME", "gemma-3-4b")
-    client =OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url =os.getenv("API_BASE_URL", "https://api.openai.com/v1"))
-
-    scores: Dict[str, float] = {}
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
+
     try:
-        for task_id in TASKS:
-            scores[task_id] = run_task(client, model, task_id)
-            print(f"[TASK] {task_id} score={scores[task_id]:.4f}")
-    except AuthenticationError:
-        print("ERROR: OpenAI authentication failed. Check OPENAI_API_KEY.", file=sys.stderr)
+        if not API_KEY:
+            raise RuntimeError("Missing OPENAI_API_KEY")
+
+        client = OpenAI(api_key=API_KEY, base_url=API_BASE_URL)
+        obs = env.reset(task_id=task_id)["observation"]
+        done = False
+
+        while not done:
+            action = choose_action(client, obs)
+            action_str = action.optimized_code
+            step_result = env.step(action)
+            done = step_result.done
+            obs = step_result.observation.model_dump()
+            reward = step_result.reward.value
+            rewards.append(reward)
+            steps_taken = obs["step_count"]
+            log_step(step=steps_taken, action=action_str, reward=reward, done=done, error=None)
+
+        score = grade_episode(
+            task_id,
+            env.state.completed_checks,
+            env.state.best_speedup,
+            env.state.step_count,
+            env.state.max_steps,
+        )
+        score = min(max(score, 0.0), 1.0)
+        success = score >= 0.1
+        return 0
+    except Exception as exc:
+        log_step(
+            step=max(1, steps_taken + 1),
+            action="error",
+            reward=0.0,
+            done=True,
+            error=str(exc),
+        )
         return 1
+    finally:
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
 
-    avg = sum(scores.values()) / len(scores)
-    print(f"[BASELINE] model={model} average_score={avg:.4f}")
-    print(json.dumps({"scores": scores, "average": round(avg, 4)}))
-    return 0
-if __name__=="__main__":
+if __name__ == "__main__":
     sys.exit(main())