Spaces:
Running
Running
| """ | |
| Inference Script — Stocker OpenEnv | |
| ================================== | |
| MANDATORY | |
| - Environment variables: | |
| API_BASE_URL The API endpoint for the LLM (default: HuggingFace Router) | |
| MODEL_NAME The model identifier (default: Qwen/Qwen2.5-7B-Instruct) | |
| HF_TOKEN Your HuggingFace / API key | |
| API_KEY Alternative to HF_TOKEN | |
| - Place inference.py in the project root. Uses OpenAI Client for LLM calls. | |
| STDOUT FORMAT | |
| [START] task=<task_name> env=<benchmark> model=<model_name> | |
| [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null> | |
| [END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn> | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import textwrap | |
| from typing import Optional | |
| from openai import OpenAI | |
| from app.core.environment import StockerEnv | |
| from app.core.tasks import list_task_ids | |
| API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") | |
| API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1" | |
| MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-7B-Instruct" | |
| BENCHMARK = "stocker" | |
| MAX_STEPS = 50 | |
| TEMPERATURE = 0.2 | |
| MAX_TOKENS = 200 | |
| SUCCESS_SCORE_THRESHOLD = 0.0 # any non-negative net P&L counts as success | |
| SYSTEM_PROMPT = textwrap.dedent("""\ | |
| You are a disciplined stock trader. At each step you observe the current | |
| price, recent price history, your cash, and your share position. You must | |
| decide: buy, sell, or hold — and if buying or selling, how many shares. | |
| Rules: | |
| - You cannot buy more than your cash allows. | |
| - You cannot sell more shares than you own. | |
| - Aim to maximize the final portfolio value. | |
| Respond with ONLY a valid JSON object (no markdown, no extra text): | |
| {"side": "buy|sell|hold", "quantity": <int>}""") | |
| 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: | |
| print( | |
| f"[STEP] step={step} action={action} reward={reward:.4f} done={str(done).lower()} " | |
| f"error={error if error else 'null'}", | |
| flush=True, | |
| ) | |
| def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None: | |
| rewards_str = ",".join(f"{r:.4f}" for r in rewards) | |
| print( | |
| f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards={rewards_str}", | |
| flush=True, | |
| ) | |
| def build_user_prompt(observation: dict) -> str: | |
| return f"""Trade decision for {observation['ticker']} on {observation['date']}. | |
| Current price: {observation['price']:.2f} | |
| Recent prices: {observation['price_history']} | |
| Fundamentals: {observation['fundamentals']} | |
| Cash: {observation['cash']:.2f} | |
| Position: {observation['position']} shares | |
| Portfolio value: {observation['portfolio_value']:.2f} | |
| Step {observation['step_number']} of {observation['total_steps']}""" | |
| def parse_response(text: str) -> dict: | |
| try: | |
| cleaned = re.sub(r"```(?:json)?\s*", "", text).strip().rstrip("`") | |
| start = cleaned.find("{") | |
| if start >= 0: | |
| depth = 0 | |
| for i in range(start, len(cleaned)): | |
| if cleaned[i] == "{": | |
| depth += 1 | |
| elif cleaned[i] == "}": | |
| depth -= 1 | |
| if depth == 0: | |
| parsed = json.loads(cleaned[start : i + 1]) | |
| side = str(parsed.get("side", "hold")).lower() | |
| if side not in ("buy", "sell", "hold"): | |
| side = "hold" | |
| qty = int(parsed.get("quantity", 0)) | |
| return {"side": side, "quantity": max(0, qty)} | |
| except (json.JSONDecodeError, ValueError, TypeError): | |
| pass | |
| side_match = re.search(r"side[\"']?\s*[:=]\s*[\"']?(buy|sell|hold)", text, re.I) | |
| qty_match = re.search(r"quantity[\"']?\s*[:=]\s*[\"']?(-?\d+)", text, re.I) | |
| if side_match: | |
| return { | |
| "side": side_match.group(1).lower(), | |
| "quantity": max(0, int(qty_match.group(1))) if qty_match else 0, | |
| } | |
| return {"side": "hold", "quantity": 0} | |
| def run_episode( | |
| client: OpenAI, model: str, task_id: str, use_json_mode: bool = True | |
| ) -> dict: | |
| env = StockerEnv(task_id=task_id) | |
| reset_result = env.reset() | |
| obs = reset_result.observation.model_dump() | |
| rewards: list[float] = [] | |
| step_details: list[dict] = [] | |
| steps_taken = 0 | |
| log_start(task=task_id, env=BENCHMARK, model=model) | |
| try: | |
| for step in range(1, MAX_STEPS + 1): | |
| messages = [ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": build_user_prompt(obs)}, | |
| ] | |
| error: Optional[str] = None | |
| try: | |
| kwargs = { | |
| "model": model, | |
| "messages": messages, | |
| "temperature": TEMPERATURE, | |
| "max_tokens": MAX_TOKENS, | |
| } | |
| if use_json_mode: | |
| kwargs["response_format"] = {"type": "json_object"} | |
| response = client.chat.completions.create(**kwargs) | |
| llm_text = (response.choices[0].message.content or "").strip() | |
| except Exception as e: | |
| err = str(e) | |
| if use_json_mode and ("response_format" in err or "json" in err.lower()): | |
| use_json_mode = False | |
| try: | |
| response = client.chat.completions.create( | |
| model=model, messages=messages, | |
| temperature=TEMPERATURE, max_tokens=MAX_TOKENS, | |
| ) | |
| llm_text = (response.choices[0].message.content or "").strip() | |
| except Exception as e2: | |
| error = str(e2) | |
| llm_text = '{"side": "hold", "quantity": 0}' | |
| else: | |
| error = err | |
| llm_text = '{"side": "hold", "quantity": 0}' | |
| action_dict = parse_response(llm_text) | |
| step_result = env.step(action_dict) | |
| reward = step_result.reward | |
| done = step_result.done | |
| rewards.append(reward) | |
| steps_taken = step | |
| action_str = f"{action_dict['side']}({action_dict['quantity']})" | |
| log_step(step, action_str, reward, done, error) | |
| step_details.append({ | |
| "step": step, | |
| "side": action_dict["side"], | |
| "quantity": action_dict["quantity"], | |
| "reward": reward, | |
| "info": step_result.info, | |
| "done": done, | |
| }) | |
| if done: | |
| break | |
| obs = step_result.observation.model_dump() | |
| score = sum(rewards) | |
| success = score >= SUCCESS_SCORE_THRESHOLD | |
| finally: | |
| log_end(success=success, steps=steps_taken, score=score, rewards=rewards) | |
| return { | |
| "task_id": task_id, | |
| "score": round(score, 4), | |
| "total_reward": round(sum(rewards), 4), | |
| "steps": steps_taken, | |
| "success": success, | |
| "step_details": step_details, | |
| } | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Run LLM inference on Stocker") | |
| parser.add_argument("--task", default="all", help="Task name or 'all'") | |
| parser.add_argument("--model", default=None, help=f"Model (default: {MODEL_NAME})") | |
| parser.add_argument("--output", default=None, help="Path to write JSON results") | |
| parser.add_argument( | |
| "--no-json-mode", action="store_true", | |
| help="Disable JSON response mode for models that don't support it", | |
| ) | |
| args = parser.parse_args() | |
| if not API_KEY: | |
| print("ERROR: No API key. Set HF_TOKEN, API_KEY, or OPENAI_API_KEY.") | |
| raise SystemExit(1) | |
| client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) | |
| model = args.model or MODEL_NAME | |
| tasks = list_task_ids() if args.task == "all" else [args.task] | |
| results = [] | |
| for task_id in tasks: | |
| results.append(run_episode(client, model, task_id, use_json_mode=not args.no_json_mode)) | |
| print(f"\n{'='*56}", file=sys.stderr) | |
| print(f"{'Task':<18} {'Score':>10} {'Reward':>10} {'Steps':>6} {'Pass':>6}", file=sys.stderr) | |
| for r in results: | |
| print( | |
| f"{r['task_id']:<18} {r['score']:>10.4f} {r['total_reward']:>10.4f} " | |
| f"{r['steps']:>6} {'yes' if r['success'] else 'no':>6}", | |
| file=sys.stderr, | |
| ) | |
| if args.output: | |
| with open(args.output, "w") as f: | |
| json.dump({ | |
| "model": model, | |
| "api_base_url": API_BASE_URL, | |
| "benchmark": BENCHMARK, | |
| "tasks": results, | |
| "total_score": round(sum(r["score"] for r in results) / max(len(results), 1), 4), | |
| }, f, indent=2) | |
| print(f"\nResults written to {args.output}", file=sys.stderr) | |
| if __name__ == "__main__": | |
| main() | |