scripts/marlin_smoketest.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "vllm",
#     "requests",
#     "qwen-vl-utils",
# ]
# ///
"""Smoke test: NemoStation/Marlin-2B on vLLM, captioning Prelinger clips from a bucket."""

import os
import re

import requests

BUCKET = "davanstrien/prelinger-sample"
CLIPS = ["AboutBan1935.mp4", "JoanAvoi1947.mp4", "Sleepfor1950.mp4"]
MODEL = "NemoStation/Marlin-2B"

# Canonical training prompt from modeling_marlin.py — must match exactly.
CAPTION_PROMPT = (
    "Provide a spatial description of this clip followed by time-ranged events.\n"
    "For each event, give the time range as <start - end> and a short description."
)

THINK = re.compile(r"<think>.*?</think>\s*|^\s*<think>\s*\n*|</think>\s*", re.DOTALL)


def fetch_clips() -> list[str]:
    token = os.environ["HF_TOKEN"]
    os.makedirs("/tmp/clips", exist_ok=True)
    paths = []
    for name in CLIPS:
        url = f"https://huggingface.co/buckets/{BUCKET}/resolve/clips/{name}"
        dest = f"/tmp/clips/{name}"
        r = requests.get(url, headers={"Authorization": f"Bearer {token}"}, timeout=120)
        r.raise_for_status()
        with open(dest, "wb") as f:
            f.write(r.content)
        print(f"fetched {name}: {len(r.content)} bytes", flush=True)
        paths.append(dest)
    return paths


def main():
    paths = fetch_clips()

    from vllm import LLM, SamplingParams

    llm = LLM(
        model=MODEL,
        hf_overrides={"architectures": ["Qwen3_5ForConditionalGeneration"]},
        allowed_local_media_path="/tmp/clips",
        max_model_len=65536,
        limit_mm_per_prompt={"video": 1},
        enforce_eager=True,
    )

    params = SamplingParams(temperature=0, max_tokens=1024)
    for path in paths:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "video_url", "video_url": {"url": f"file://{path}"}},
                    {"type": "text", "text": CAPTION_PROMPT},
                ],
            }
        ]
        out = llm.chat(messages, params)
        text = THINK.sub("", out[0].outputs[0].text).strip()
        print(f"\n===== {os.path.basename(path)} =====\n{text}\n", flush=True)

    print("SMOKE TEST COMPLETE", flush=True)


if __name__ == "__main__":
    main()