File size: 3,608 Bytes
344db23 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
REFERENCE ONLY — This is the official sample inference script from the hackathon dashboard.
It's for BrowserGym, NOT our Security Audit env. But the PATTERN is what we must follow.
KEY TAKEAWAYS FOR OUR inference.py:
1. API_BASE_URL defaults to "https://router.huggingface.co/v1" (HuggingFace Inference)
2. API_KEY = HF_TOKEN or API_KEY (support both)
3. MODEL_NAME from env var
4. OpenAI client with base_url=API_BASE_URL
5. Loop: reset → step → step → ... until done or max_steps
6. Build user prompt from observation + history
7. Parse LLM response into action
8. Fallback action if LLM fails
9. Track history of actions/results
10. Print reward/done/error at each step
"""
import os
import re
import textwrap
from typing import List, Optional, Dict
from openai import OpenAI
# --- ENV VARS (MANDATORY) ---
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
MODEL_NAME = os.getenv("MODEL_NAME")
# --- CONFIG ---
MAX_STEPS = 8
TEMPERATURE = 0.2
MAX_TOKENS = 200
FALLBACK_ACTION = "noop()"
# --- SYSTEM PROMPT (customize for your env) ---
SYSTEM_PROMPT = textwrap.dedent("""
You control a web browser through BrowserGym.
Reply with exactly one action string.
...
""").strip()
def build_user_prompt(step, observation, history):
"""Format the observation into a prompt for the LLM."""
# ... customize for your env
pass
def parse_model_action(response_text):
"""Extract action from LLM's text response."""
# ... customize for your env
pass
def main():
client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
# Connect to your environment
env = ... # YourEnv(base_url="https://your-space.hf.space").sync()
history = []
try:
result = env.reset()
observation = result.observation
for step in range(1, MAX_STEPS + 1):
if result.done:
break
# Build prompt from observation
user_prompt = build_user_prompt(step, observation, history)
messages = [
{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
{"role": "user", "content": [{"type": "text", "text": user_prompt}]},
]
# Call LLM
try:
completion = client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
temperature=TEMPERATURE,
max_tokens=MAX_TOKENS,
stream=False,
)
response_text = completion.choices[0].message.content or ""
except Exception as exc:
print(f"Model request failed ({exc}). Using fallback.")
response_text = FALLBACK_ACTION
# Parse action from LLM output
action_str = parse_model_action(response_text)
print(f"Step {step}: model suggested -> {action_str}")
# Step the environment
result = env.step(action_str) # adapt to your Action type
observation = result.observation
# Track
reward = result.reward or 0.0
history.append(f"Step {step}: {action_str} -> reward {reward:+.2f}")
print(f" Reward: {reward:+.2f} | Done: {result.done}")
if result.done:
print("Episode complete.")
break
else:
print(f"Reached max steps ({MAX_STEPS}).")
finally:
env.close()
if __name__ == "__main__":
main()
|