Update inference.py

inference.py

@@ -1,8 +1,6 @@
 import os
 import sys
-import json
 import textwrap
-import asyncio
 from typing import List
 
 from openai import OpenAI
@@ -16,9 +14,6 @@ MAX_STEPS = 5
 TEMPERATURE = 0.2
 MAX_TOKENS = 200
 FALLBACK_ACTION = "skip"
-TASK_NAME = "code_review"          # or "easy", "medium", etc. – but we'll vary per task
-ENV_NAME = "code_review_env"
-SUCCESS_THRESHOLD = 0.5            # if final score >= 0.5, consider success
 
 SYSTEM_PROMPT = textwrap.dedent(
     """
@@ -88,27 +83,14 @@ def parse_model_action(response_text: str) -> Action:
         return Action(action_type="propose_fix", fix_code=fix)
     return Action(action_type="write_comment", comment_text=raw)
 
-def log_start(task: str, env: str, model: str):
-    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):
-    error_str = f" error={error}" if error else ""
-    print(f"[STEP] step={step} action={action} reward={reward:.3f} done={done}{error_str}", flush=True)
-
-def log_end(success: bool, steps: int, score: float, rewards: List[float]):
-    # rewards as JSON array to match sample (though sample shows `rewards=rewards` which might be a list)
-    rewards_str = json.dumps(rewards, separators=(',', ':'))
-    print(f"[END] success={success} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)
-
-async def main() -> None:
+def main() -> None:
     if not API_BASE_URL or not API_KEY or not MODEL_NAME:
+        # Print error to stderr (won't affect stdout)
         print("Error: API_BASE_URL, HF_TOKEN/API_KEY, and MODEL_NAME must be set.", file=sys.stderr)
         return
 
     client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-    env = CodeReviewEnv()   # synchronous environment – we'll wrap calls in async? The sample uses async, but we can run sync and just not await.
-    # Since our env is sync, we'll call methods directly. The logging is the key part.
-
+    env = CodeReviewEnv()
     tasks = ["easy", "medium", "hard", "harder", "hardest"]
 
     for task in tasks:
@@ -117,10 +99,10 @@ async def main() -> None:
         history: List[str] = []
         done = False
         step = 0
-        rewards = []
-        final_score = 0.0
+        final_reward = 0.0
 
-        log_start(task=task, env=ENV_NAME, model=MODEL_NAME)
+        # START marker
+        print(f"[START] task={task}", flush=True)
 
         while not done and step < MAX_STEPS:
             step += 1
@@ -138,28 +120,21 @@ async def main() -> None:
                 )
                 response_text = completion.choices[0].message.content or ""
             except Exception as exc:
+                # Print error to stderr only
                 print(f"Request failed: {exc}. Using fallback.", file=sys.stderr)
                 response_text = FALLBACK_ACTION
 
             action = parse_model_action(response_text)
             obs, reward, done, info = env.step(action)
-            reward_val = reward.value
-            rewards.append(reward_val)
+            final_reward = reward.value
 
-            log_step(step=step, action=action.action_type, reward=reward_val, done=done)
+            # STEP marker
+            print(f"[STEP] step={step} reward={final_reward:.3f}", flush=True)
 
             history.append(f"Step {step}: {action.action_type}")
 
-        # Compute overall score (e.g., average reward or final reward? Sample uses sum(rewards)/MAX_TOTAL_REWARD.
-        # We'll use average reward per step as score, clamped to [0,1].
-        if rewards:
-            avg_reward = sum(rewards) / len(rewards)
-            score = max(0.0, min(1.0, avg_reward))
-        else:
-            score = 0.0
-        success = score >= SUCCESS_THRESHOLD
-
-        log_end(success=success, steps=step, score=score, rewards=rewards)
+        # END marker
+        print(f"[END] task={task} score={final_reward:.3f} steps={step}", flush=True)
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()