| """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 |
|
|
| |
| |
| MODEL_ENV = "FLORA_REWARDER_REPO" |
| REPOSITORY_COPY = Path(__file__).absolute().parent.parent |
| |
| |
| |
| |
| 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, |
| |
| |
| 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) |
|
|
