File size: 6,459 Bytes
79dd27b | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Shared command-line plumbing and the offline vLLM engine."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import re
import sys
from typing import Any, Iterable
# Without --model: $FLORA_REWARDER_REPO, else the downloaded copy of the repository that holds
# batch_utils/ (absolute, not resolved: Hub cache files are symlinks).
MODEL_ENV = "FLORA_REWARDER_REPO"
REPOSITORY_COPY = Path(__file__).absolute().parent.parent
# Pixel budget for collage clips (terminal and dense): at least ``shortest_edge`` and at most
# ``longest_edge`` pixels over all frames of one clip. Collages are planned inside this budget,
# so frames reach the model at their planned size.
# Configuration and tokenizer files: enough to render prompts for --dry-run.
PROCESSOR_FILES = ["*.json", "*.jinja", "*.txt"]
COLLAGE_VIDEO_SIZE = {"shortest_edge": 4096, "longest_edge": 29_360_128}
def default_model() -> str | None:
if os.environ.get(MODEL_ENV):
return os.environ[MODEL_ENV]
if (REPOSITORY_COPY / "model.safetensors.index.json").is_file():
return str(REPOSITORY_COPY)
return None
def base_parser(description: str, *, tensor_parallel_size: int = 1) -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description=description)
source = p.add_mutually_exclusive_group(required=True)
source.add_argument("--jobs", type=Path, help="JSONL file, one episode per line (paths relative to the file).")
source.add_argument("--video", type=Path, nargs="+", help="Camera video(s) of one episode, in order.")
p.add_argument("--task", help="Task instruction for --video.")
p.add_argument("--output", type=Path, required=True, help="Output JSONL file.")
p.add_argument("--overwrite", action="store_true", help="Replace an existing output file.")
p.add_argument("--model", default=default_model(),
help=f"Hub repository id or local directory (default: ${MODEL_ENV}, else this repository's copy).")
p.add_argument("--revision", default=None, help="Hub revision (branch, tag or commit).")
p.add_argument("--tensor-parallel-size", type=int, default=tensor_parallel_size)
p.add_argument("--gpu-memory-utilization", type=float, default=0.88)
p.add_argument("--max-num-seqs", type=int, default=8, help="Concurrent requests in the engine.")
p.add_argument("--seed", type=int, default=0)
p.add_argument("--dry-run", action="store_true", help="Read the videos and print the first prompt; no model.")
return p
def read_jobs(args: argparse.Namespace, extra_fields: Iterable[str] = ()) -> list[dict[str, Any]]:
"""Jobs from ``--jobs`` or from ``--video``/``--task``; video paths become absolute."""
if args.video:
if not args.task:
raise SystemExit("--task is required with --video")
return [{"id": "episode", "task": args.task, "videos": [str(p.expanduser().resolve()) for p in args.video]}]
root = args.jobs.expanduser().resolve().parent
allowed = {"id", "task", "videos", *extra_fields}
jobs, ids = [], set()
for number, line in enumerate(args.jobs.read_text().splitlines(), 1):
if not line.strip():
continue
job = json.loads(line)
if not isinstance(job, dict) or set(job) - allowed or not {"task", "videos"} <= set(job):
raise SystemExit(f"{args.jobs}:{number}: expected keys {sorted(allowed)}")
job.setdefault("id", f"episode_{number:06d}")
if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}", str(job["id"])) or job["id"] in ids:
raise SystemExit(f"{args.jobs}:{number}: invalid or duplicate id {job['id']!r}")
ids.add(job["id"])
if not isinstance(job["videos"], list) or not job["videos"]:
raise SystemExit(f"{args.jobs}:{number}: 'videos' must be a non-empty list")
job["videos"] = [str((root / Path(v).expanduser()).resolve()) for v in job["videos"]]
jobs.append(job)
if not jobs:
raise SystemExit(f"{args.jobs} has no jobs")
return jobs
def open_output(args: argparse.Namespace):
if args.output.exists() and not args.overwrite:
raise SystemExit(f"{args.output} exists; pass --overwrite to replace it")
args.output.parent.mkdir(parents=True, exist_ok=True)
return args.output.open("w")
def write_record(stream, record: dict[str, Any]) -> None:
stream.write(json.dumps(record, ensure_ascii=False) + "\n")
stream.flush()
def resolve_model(model: str | None, revision: str | None = None, *, weights: bool = True) -> str:
"""A local directory with the model; ``weights=False`` fetches only the processor files."""
if not model:
raise SystemExit(f"Pass --model (a Hub repository id or a local directory) or set {MODEL_ENV}.")
path = Path(model).expanduser()
if path.is_dir():
return str(path.resolve())
from huggingface_hub import snapshot_download
return snapshot_download(model, revision=revision, allow_patterns=None if weights else PROCESSOR_FILES)
def load_processor(path: str):
from transformers import AutoProcessor
return AutoProcessor.from_pretrained(path)
def load_llm(
path: str,
args: argparse.Namespace,
*,
max_model_len: int,
videos_per_prompt: int = 1,
mm_processor_kwargs: dict[str, Any] | None = None,
**engine_kwargs: Any,
):
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
from vllm import LLM
return LLM(
model=path,
tensor_parallel_size=args.tensor_parallel_size,
dtype="bfloat16",
gpu_memory_utilization=args.gpu_memory_utilization,
max_model_len=max_model_len,
max_num_seqs=args.max_num_seqs,
limit_mm_per_prompt={"image": 0, "video": videos_per_prompt},
mm_processor_kwargs=mm_processor_kwargs,
# Triton prefill for the linear-attention layers: the configuration the reported
# results were measured with.
gdn_prefill_backend="triton",
enable_prefix_caching=True,
seed=args.seed,
**engine_kwargs,
)
def render(processor, messages, *, thinking: bool, vision_ids: bool = False) -> str:
return processor.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=thinking,
add_vision_id=vision_ids,
)
def log(message: str) -> None:
print(message, file=sys.stderr, flush=True)
|