Skip to content

Latest commit

 

History

History
98 lines (82 loc) · 5.06 KB

File metadata and controls

98 lines (82 loc) · 5.06 KB

EgoSafetyBench — the model abstraction (egosafetybench.models)

One provider-agnostic way to call any model — local open-weight or hosted API — shared by annotate, eval, and scenegen. Adding a model is dropping one YAML; adding a provider is one transport class.

The one factory

from egosafetybench.models.spec import load_spec
from egosafetybench.models.registry import build_model
from egosafetybench.models.message import build_window_message  # eval
# (annotate builds its own Message via egosafetybench.annotate.prompts)

spec  = load_spec("configs/models/qwen3_5_0_8b.yaml")
model = build_model(spec, ledger=ledger, limiter=limiter)   # ledger/limiter optional
model.load()
result = model.generate(message, spec.gen)                  # -> GenerationResult(.text, .usage, ...)

build_model is the only constructor. It returns:

  • a local transport (HFLocalTransport + a per-family hook) for kind: local_hf
  • an API transport wrapped in MeteredModel (cost/budget/retry) for kind: api

ModelSpec YAML

name: qwen3_5_0_8b            # short name (used in output dirs / roster)
kind: local_hf               # local_hf | api
provider: hf                 # hf | openrouter | google_genai | anthropic
model_id: Qwen/Qwen3.5-0.8B  # HF repo id, OpenRouter slug, or Anthropic model id
slug: qwen3_5_0_8b           # optional API slug (defaults to model_id)
group: Open-source           # optional: aggregate-table roster group
api_key_env: ANTHROPIC_API_KEY   # api only: env var holding the key
base_url: https://openrouter.ai/api/v1   # openai-compatible only
gen:                         # generation params (ONE source of truth)
  max_output_tokens: 192
  temperature: 0.0           # ignored by transports that reject sampling params
  do_sample: false
limits:
  max_concurrency: 8
  rpm: null                  # token-bucket rate limits (per provider+model)
  tpm: null
  max_retries: 5
  backoff_base_s: 1.0
  timeout_s: 120
pricing:                     # used by the cost estimate + budget ledger
  usd_per_mtok_in: 0.25
  usd_per_mtok_out: 1.00
supports_window: true        # eval window protocol capability
supports_json_mode: true
extra:                       # free-form, e.g. local dtype/device or capability flags
  family: qwen3_5            # local_hf: which family hook to use
  dtype: bfloat16
  device_map: auto
  attn_implementation: sdpa
  vcm_capable: false         # eval: model has no VCM head (e.g. Cosmos) -> skip axis

Providers (transports)

provider transport notes
hf HFLocalTransport + families/<name> local open-weight; extra.family selects the per-family render+generate hook (qwen3_5 / qwen3_vl / gemma3 / internvl3_5 / cosmos)
openrouter (openai_compat) OpenAICompatTransport OpenAI Chat Completions shape; opts into usage accounting so usage.cost is populated
google_genai GoogleGenAITransport Gemini; surfaces safety/finish-reason instead of swallowing as empty
anthropic AnthropicTransport public pinned Claude id, base64 image blocks; never sends temperature/top_p/budget_tokens (Opus 4.7/4.8 reject them)

Middleware (shared, wraps every API transport)

  • CostModel.estimate_call — text tokens (tiktoken/heuristic) + image tokens (per-provider rule, calibrated to ~1.1k tok per full-res frame) → projected USD.
  • preflight — sums the estimate over a planned call set: {n_calls, est_tokens_in/out, est_usd}, gate on --max-cost/--yes before any spend.
  • BudgetLedger — one per run, all providers, thread-safe. reserve(est) before each call (raises BudgetExceeded → clean resume-safe stop), settle with the provider-reported actual (falls back to the estimate, e.g. Google).
  • RateLimiter — token-bucket per (provider, model) from limits.rpm/tpm.
  • RetryPolicy — classifies on typed exceptions / HTTP status (not string matching): transient → backoff; rate-limited → backoff then clean stop; content-blocked/auth/budget → stop; fatal → no retry.

Local models skip the middleware (no network, no billing).

Adding things

  • A model → drop a configs/models/<name>.yaml. Nothing else.
  • A local family → add models/families/<name>.py implementing load() + generate() (lift the per-family generate tail), register it in registry._local_family.
  • A provider → add models/transports/<provider>.py (a _one(parts, gen) + _init_client), register it in registry.build_model.

Gotchas baked in (so you don't rediscover them)

  • Opus 4.7/4.8 (and Fable 5) 400 on temperature/top_p/top_k/budget_tokens — the Anthropic transport omits them.
  • OpenRouter only reports usage.cost if you opt in (extra_body={"usage":{"include":True}}) — the transport does this; without it a budget cap silently never fires.
  • Image payloads in a neutral Message may be PIL images (eval, encoded at send) or raw JPEG bytes (annotation's ffmpeg frames, passed through) — to_jpeg_parts handles both so the on-the-wire bytes match what produced the GT.