| """ |
| Dummy algorithm for the ORena SAVE FOCUS challenge — SEGMENT track. |
| |
| Replace the marked blocks with your model. Every spot to change carries a ``TODO``. |
| |
| What one run receives |
| --------------------- |
| A single container run answers a **batch of questions**. |
| Inputs are mounted read-only at ``/input``: |
| |
| request.json a LIST of focus.Request objects — one per question |
| plain/<qID>.mp4 the video clip for question <qID> |
| overlayed/<qID>.mp4 the same clip with an HH:MM:SS timestamp burned in |
| FO_definitions.json the foreign-object class definitions, as plain text |
| batch.json a convenience index of the batch (qIDs and layout) |
| |
| and one file is written to ``/output``: |
| |
| answer.json a LIST of focus.Response objects — one per question |
| |
| ``request.json`` is authoritative: iterate it, and answer every ``qID`` it lists. |
| |
| Each question carries its own clip |
| ---------------------------------- |
| Every question has a **separate video file, already trimmed to that question's |
| window**, and no two questions in a batch come from the same source video. Your |
| answer to one question can therefore never depend on another question's footage. |
| |
| Two consequences for your pipeline: |
| |
| * ``request.start_time`` / ``end_time`` refer to the **original, untrimmed |
| procedure video** — the same timeline shown in the overlay variant. The clip you |
| are given has already been cut to that window, so decode the clip from its own |
| beginning; do **not** seek to ``start_time`` inside it. ``request.duration`` |
| still tells you how long the window is. |
| * There is nothing to share between questions — no common video to cache or encode |
| once. Each clip is decoded on its own. |
| |
| Both variants show the same footage. Use the plain clip, the overlay, or both. |
| |
| Clip encoding |
| ------------- |
| Clips are H.264 MP4 at exactly **5 fps** (one frame every 0.2 s), height-normalised |
| to at most 576 px. Widths follow the source's aspect ratio (e.g., 1024x576, 720x576 or |
| 640x360) — do not assume a fixed resolution. A |
| keyframe every 5 s keeps seeking cheap. The overlay clock shows the source |
| timeline's second, floored. |
| |
| Load your model once |
| -------------------- |
| Model loading, imports and warm-up happen **once per container run**, not once per |
| question. The latency budget is built around that, so keep model loading outside |
| the per-question loop — as this file does. This is the single most important thing |
| to get right for speed. |
| |
| The latency budget |
| ------------------ |
| The platform measures the execution time of the whole run itself; the ``latency`` |
| field you write into each Response is informational and is not used for scoring. |
| |
| The budget is pooled across the batch. For a batch of B questions: |
| |
| allowed = 120 s + B x 15 s |
| ^setup ^per-question processing (the SEGMENT budget) |
| |
| Going over does not fail the run outright. Instead, a share of your questions is |
| **forfeited** (scored incorrect), growing with the overrun: a small overrun costs a |
| few questions, and exceeding the budget by 20% forfeits the whole batch. |
| |
| The forfeited questions are selected for you — not necessarily the slow ones — so an |
| overrun costs you questions you would have answered correctly. If you are running out |
| of budget, answer the rest quickly rather than overrun. |
| |
| Keep going when a question fails |
| -------------------------------- |
| One bad question should not cost you the batch. Catch errors per question, still |
| emit a Response for that ``qID``, and carry on — as the loop below does. A question |
| with no answer is simply scored incorrect. |
| |
| Frame sampling |
| -------------- |
| The clip is decoded with `decord <https://pypi.org/project/decord/>`_. Full-frame |
| solid-blue signal-loss screens are detected from pixels, without relying on a |
| video identifier. Exactly 80 bin-centre frames are then sampled over the remaining |
| time and letterboxed to the Tr006 training resolution. Their timestamps are mapped |
| back to the original timeline from ``Request.start_time``. |
| |
| Trying it out |
| ------------- |
| A sample batch (``test/input/interface_1/``) ships with the template. This package |
| stores the merged Tr006 checkpoint under ``resources/model``; runtime inference is |
| fully local and does not download model files. |
| |
| ./do_test_run.sh # build the image and run it on the sample batch |
| ./do_save.sh # build and save a .tar.gz for upload |
| """ |
|
|
| import json |
| import logging |
| import re |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| |
| |
| |
| import decord |
| from focus import Request, Response, load_requests, save_items |
| from PIL import Image, ImageOps |
| from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration |
|
|
| logging.basicConfig( |
| stream=sys.stdout, |
| level=logging.DEBUG, |
| format="%(asctime)s [%(levelname)s] %(message)s", |
| datefmt="%Y-%m-%d %H:%M:%S", |
| ) |
| log = logging.getLogger(__name__) |
|
|
| RESOURCES_PATH = Path(__file__).parent / "resources" |
| MODEL_PATH = RESOURCES_PATH / "model" |
|
|
| INPUT_PATH = Path("/input") |
| OUTPUT_PATH = Path("/output") |
| N_FRAMES = 80 |
| FRAME_WIDTH = 420 |
| FRAME_HEIGHT = 224 |
| BLUE_DETECT_CHUNK = 64 |
| BLUE_DETECT_STRIDE = 16 |
| BLUE_REFINE_ATTEMPTS = 2 |
| MAX_NEW_TOKENS = 64 |
| USE_OVERLAY = False |
| VIDEO_DIR = INPUT_PATH / ("overlayed" if USE_OVERLAY else "plain") |
|
|
| |
| |
| |
| |
| SYSTEM_PROMPT = ( |
| "You are a surgical video analysis assistant. You are shown a short clip from a laparoscopic procedure and asked a single question about it.\n\n" |
| "The clip is given as a sequence of frames in chronological order. Each frame is preceded by its timestamp in the source video, written as " |
| "<SECONDS seconds> -- a frame marked <1234.5 seconds> was taken 1234.5 seconds (00:20:34) after the video began. Consecutive frames may be several seconds " |
| "apart, so an event can happen between two of them.\n\n" |
| "A foreign object (FO) is any object fully introduced into the patient's body\n" |
| "cavity during surgery that must be retrieved or accounted for. Importantly,\n" |
| "standard surgical instruments that remain connected to the external environment\n" |
| "(e.g., graspers, scissors, trocars, staplers, cameras) are not considered foreign\n" |
| "objects. Furthermore, we exclude detachable parts of surgical instruments,\n" |
| "particularly anvil components of staplers.\n\n" |
| "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent.\n\n" |
| "Reply with the answer and nothing else -- no reasoning, no preamble, no\n" |
| "explanation, no restating the question. A single short line.\n\n" |
| "Rules for the answer:\n" |
| "- Write the value only. No sentence, no explanation, no units, no trailing\n" |
| " period, and never repeat the question.\n" |
| "- Asks yes or no -> write exactly: yes or no\n" |
| "- Asks how many / for a count -> write digits only, e.g. 0 or 1 or 2.\n" |
| "- Asks which foreign object class(es) -> write class names exactly as spelled\n" |
| " in the list above, comma-separated (e.g. Clip, Sponge), or exactly: none\n" |
| " Never answer with a generic description such as \"surgical instrument\".\n" |
| "- Asks when something happens -> write hh:mm:ss, read off the frame timestamps.\n" |
| " Those timestamps count from the START OF THE VIDEO, not from the start of this\n" |
| " clip, so <1234.5 seconds> is 00:20:34. Convert carefully: 3600 seconds is one\n" |
| " hour, 60 seconds is one minute. If the event falls between two frames, give\n" |
| " your best estimate between their timestamps rather than snapping to one.\n" |
| " For several time points, separate them with commas: 00:09:19, 00:12:44\n" |
| "- Asks for a percentage or a share -> write the number only, e.g. 40\n" |
| "- Lists options to choose from -> copy exactly one of those options, verbatim.\n" |
| "- Anything else -> a short phrase, at most a few words.\n\n" |
| "If you are unsure, still commit to your single best answer in the required\n" |
| "form. An empty, hedged, or explanatory answer is scored as wrong.\n" |
| ) |
|
|
|
|
| def _flatten_token_ids(value) -> list[int]: |
| """Flatten tokenizer/config token-id values while dropping invalid entries.""" |
| if value is None or isinstance(value, bool): |
| return [] |
| if isinstance(value, (int, np.integer)): |
| return [int(value)] if int(value) >= 0 else [] |
| if hasattr(value, "detach"): |
| value = value.detach().cpu().tolist() |
| elif hasattr(value, "tolist") and not isinstance(value, (str, bytes)): |
| value = value.tolist() |
| if isinstance(value, (list, tuple, set)): |
| out = [] |
| for item in value: |
| out.extend(_flatten_token_ids(item)) |
| return out |
| return [] |
|
|
|
|
| def _ordered_unique(values: list[int]) -> list[int]: |
| return list(dict.fromkeys(int(value) for value in values if int(value) >= 0)) |
|
|
|
|
| def resolve_eos_token_ids(model, tokenizer) -> list[int]: |
| """Merge model and tokenizer EOS values, preserving order and removing duplicates.""" |
| values = [] |
| for owner in ( |
| getattr(model, "generation_config", None), |
| getattr(model, "config", None), |
| tokenizer, |
| ): |
| values.extend(_flatten_token_ids(getattr(owner, "eos_token_id", None))) |
| return _ordered_unique(values) |
|
|
|
|
| def resolve_im_start_token_ids(tokenizer) -> list[int]: |
| """Return the tokenizer's (possibly multi-token) ``<|im_start|>`` sequence.""" |
| try: |
| values = _flatten_token_ids(tokenizer.encode("<|im_start|>", add_special_tokens=False)) |
| except Exception: |
| values = [] |
| if values: |
| return values |
| try: |
| value = tokenizer.convert_tokens_to_ids("<|im_start|>") |
| except Exception: |
| value = None |
| return _flatten_token_ids(value) |
|
|
|
|
| def _token_rows(value) -> list[list[int]]: |
| if hasattr(value, "detach"): |
| value = value.detach().cpu().tolist() |
| elif hasattr(value, "tolist") and not isinstance(value, (str, bytes)): |
| value = value.tolist() |
| if not isinstance(value, list): |
| value = [value] |
| if value and not isinstance(value[0], list): |
| value = [value] |
| return [[int(token) for token in row] for row in value] |
|
|
|
|
| def extract_new_token_ids(generated, input_ids, attention_mask=None) -> list[list[int]]: |
| """Extract only generated IDs, accounting for padded prompt layouts.""" |
| outputs = _token_rows(generated) |
| prompts = _token_rows(input_ids) |
| masks = _token_rows(attention_mask) if attention_mask is not None else [] |
| if not prompts: |
| return outputs |
| extracted = [] |
| for index, output in enumerate(outputs): |
| prompt = prompts[min(index, len(prompts) - 1)] |
| mask = masks[min(index, len(masks) - 1)] if masks else [1] * len(prompt) |
| active_prompt = [token for token, flag in zip(prompt, mask) if flag] |
| if len(output) >= len(prompt) and output[: len(prompt)] == prompt: |
| offset = len(prompt) |
| elif active_prompt and len(output) >= len(active_prompt) and output[: len(active_prompt)] == active_prompt: |
| offset = len(active_prompt) |
| else: |
| offset = 0 |
| extracted.append(output[offset:]) |
| return extracted |
|
|
|
|
| def _first_subsequence(values: list[int], needle: list[int]) -> int | None: |
| if not needle or len(needle) > len(values): |
| return None |
| width = len(needle) |
| for index in range(len(values) - width + 1): |
| if values[index : index + width] == needle: |
| return index |
| return None |
|
|
|
|
| def trim_generated_token_ids( |
| raw_token_ids, |
| *, |
| eos_token_ids, |
| im_start_token_ids, |
| max_new_tokens: int, |
| ) -> dict: |
| """Stop at the earliest EOS or im_start sequence and retain a raw-token audit.""" |
| raw = _flatten_token_ids(raw_token_ids) |
| eos_ids = set(_flatten_token_ids(eos_token_ids)) |
| im_start = _flatten_token_ids(im_start_token_ids) |
| eos_positions = [index for index, token in enumerate(raw) if token in eos_ids] |
| eos_index = eos_positions[0] if eos_positions else None |
| im_start_index = _first_subsequence(raw, im_start) |
| candidates = [(index, "eos") for index in (eos_index,) if index is not None] |
| candidates.extend((index, "im_start") for index in (im_start_index,) if index is not None) |
| stop_index, stop_reason = min(candidates, default=(None, None), key=lambda item: (item[0], 0 if item[1] == "eos" else 1)) |
| hit_max = stop_index is None and len(raw) >= int(max_new_tokens) |
| if stop_index is None: |
| stop_reason = "max_new_tokens" if hit_max else "completed" |
| token_ids = raw |
| else: |
| token_ids = raw[:stop_index] |
| return { |
| "raw_token_ids": raw, |
| "token_ids": token_ids, |
| "stop_index": stop_index, |
| "eos_position": eos_index, |
| "eos_positions": eos_positions, |
| "stop_reason": stop_reason, |
| "hit_max_new_tokens": hit_max, |
| } |
|
|
|
|
| _LEADING_THINK = re.compile(r"\A\s*<think>.*?</think>\s*", re.IGNORECASE | re.DOTALL) |
| _TRAILING_SPECIAL = re.compile(r"(?:<\|[^>\n]+\|>|</think>|<think>)", re.IGNORECASE) |
| _ROLE_MARKER = re.compile(r"\A\s*(?:assistant|user|system)\s*(?::|\Z)", re.IGNORECASE) |
|
|
|
|
| def clean_decoded_answer(text: str) -> str: |
| """Remove only a complete leading think block and trailing template pollution.""" |
| value = str(text or "").strip() |
| leading = _LEADING_THINK.match(value) |
| if leading: |
| value = value[leading.end() :].strip() |
| elif value.lower().startswith("<think>"): |
| return value |
| value = value.splitlines()[0].strip() if value else "" |
| marker = _TRAILING_SPECIAL.search(value) |
| if marker: |
| value = value[: marker.start()].rstrip() |
| role = _ROLE_MARKER.match(value) |
| if role: |
| value = "" |
| return value |
|
|
|
|
| def log_environment(device: torch.device) -> None: |
| """Report the CUDA stack this container ended up with — once, at startup. |
| |
| A torch build carrying no kernels for the GPU it was handed is the classic |
| reason a run dies on the platform, and it stays invisible until the first |
| CUDA call. Printing this makes a failed job diagnosable from the logs alone. |
| Keep it: it costs milliseconds, once per run. |
| """ |
| log.info("--- Environment ---") |
| log.info(" torch : %s (built against CUDA %s)", torch.__version__, torch.version.cuda) |
|
|
| |
| |
| arch = torch.cuda.get_arch_list() or (torch._C._cuda_getArchFlags() or "").split() |
| log.info(" torch kernels : %s", " ".join(arch) or "unknown") |
|
|
| driver_file = Path("/proc/driver/nvidia/version") |
| log.info( |
| " host driver : %s", |
| driver_file.read_text().strip().splitlines()[0] |
| if driver_file.exists() |
| else "no NVIDIA driver visible", |
| ) |
|
|
| if device.type != "cuda": |
| log.warning(" No GPU visible — running on CPU. The platform always provides one.") |
| return |
|
|
| capability = torch.cuda.get_device_capability(0) |
| free, total = torch.cuda.mem_get_info(0) |
| log.info(" GPU : %s", torch.cuda.get_device_name(0)) |
| log.info(" capability : %s (sm_%d%d)", capability, *capability) |
| log.info(" VRAM : %.1f GiB free of %.1f GiB", free / 1024**3, total / 1024**3) |
|
|
|
|
| def clip_path_for(req: Request) -> Path: |
| """The video clip belonging to this question. |
| |
| Every question has exactly one clip, named after its ``qID`` and already cut to |
| that question's ``[start_time, end_time]`` window. |
| """ |
| return VIDEO_DIR / f"{req.qID}.mp4" |
|
|
|
|
| def blue_screen_flags(frames: np.ndarray) -> np.ndarray: |
| """Classify full-frame solid-blue signal loss on a sparse pixel grid.""" |
| sampled = frames[:, ::BLUE_DETECT_STRIDE, ::BLUE_DETECT_STRIDE] |
| pixels = sampled.reshape(len(sampled), -1, 3).astype(np.float32) |
| red, green, blue = pixels[:, :, 0], pixels[:, :, 1], pixels[:, :, 2] |
| blue_fraction = ((blue > 180) & (blue > 4 * (red + 1)) & (blue > 4 * (green + 1))).mean(1) |
| return (blue_fraction >= 0.98) & (pixels.std(1).mean(1) <= 15) |
|
|
|
|
| def sample_non_blue_indices(blue: np.ndarray) -> np.ndarray: |
| valid = np.flatnonzero(~blue) |
| if not len(valid): |
| log.warning("Every decoded frame looks like a blue screen; using the full clip") |
| valid = np.arange(len(blue)) |
| positions = (np.arange(N_FRAMES) + 0.5) * len(valid) / N_FRAMES |
| slots = np.floor(positions + 0.5 - 1e-9).astype(np.int64).clip(0, len(valid) - 1) |
| return valid[slots] |
|
|
|
|
| def detect_blue_intervals(vr: decord.VideoReader) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """Probe 80 time bins and expand blue probes to their represented intervals.""" |
| probes = sample_non_blue_indices(np.zeros(len(vr), dtype=bool)) |
| probe_frames = vr.get_batch(probes).asnumpy() |
| flags = blue_screen_flags(probe_frames) |
| if flags.all(): |
| blue = np.zeros(len(vr), dtype=bool) |
| for begin in range(0, len(vr), BLUE_DETECT_CHUNK): |
| indices = np.arange(begin, min(begin + BLUE_DETECT_CHUNK, len(vr))) |
| blue[indices] = blue_screen_flags(vr.get_batch(indices).asnumpy()) |
| return blue, probes, probe_frames |
| boundaries = np.concatenate(( |
| [0], ((probes[:-1] + probes[1:]) // 2 + 1), [len(vr)] |
| )) |
| blue = np.zeros(len(vr), dtype=bool) |
| for index in np.flatnonzero(flags): |
| blue[boundaries[index]:boundaries[index + 1]] = True |
| return blue, probes, probe_frames |
|
|
|
|
| def letterbox_frames(frames: np.ndarray) -> np.ndarray: |
| return np.stack([ |
| np.asarray(ImageOps.pad( |
| Image.fromarray(frame), |
| (FRAME_WIDTH, FRAME_HEIGHT), |
| method=Image.Resampling.BICUBIC, |
| color=(0, 0, 0), |
| centering=(0.5, 0.5), |
| )) |
| for frame in frames |
| ]) |
|
|
|
|
| def your_preprocessing(video_path: Path, req: Request) -> tuple[np.ndarray, dict]: |
| """Decode exactly ``N_FRAMES`` and attach original-timeline millisecond PTS. |
| |
| Official clips are already trimmed and encoded at 5fps. Blue signal-loss time |
| is detected online, then the same non-blue bin-centre and 420x224 letterbox |
| contracts used for Tr006 V2 are applied. ``frames_indices`` remain absolute. |
| """ |
| log.debug("Decoding %s (size on disk: %d bytes)", video_path, video_path.stat().st_size) |
|
|
| vr = decord.VideoReader(str(video_path), ctx=decord.cpu(0), num_threads=1) |
| base_fps = float(vr.get_avg_fps()) |
| n_frames = len(vr) |
| if n_frames <= 0 or base_fps <= 0: |
| raise RuntimeError(f"invalid video stream: n_frames={n_frames}, fps={base_fps}") |
| blue, indices, frames = detect_blue_intervals(vr) |
| if blue.any(): |
| for _ in range(BLUE_REFINE_ATTEMPTS + 1): |
| indices = sample_non_blue_indices(blue) |
| frames = vr.get_batch(indices).asnumpy() |
| selected_blue = blue_screen_flags(frames) |
| if not selected_blue.any() or blue.all(): |
| break |
| radius = max(1, int(np.ceil(n_frames / (2 * N_FRAMES)))) |
| for index in indices[selected_blue]: |
| blue[max(0, index - radius):min(n_frames, index + radius + 1)] = True |
| local_seconds = vr.get_frame_timestamp(indices)[:, 0] |
| frames = letterbox_frames(frames) |
| log.debug( |
| "Video: %d frames @ %.3f fps; excluded %d blue frames; selected %d non-blue bin centres", |
| n_frames, base_fps, int(blue.sum()), len(indices), |
| ) |
| del vr |
|
|
| |
| |
| |
| |
| absolute_ms = np.rint((float(req.start_time) + local_seconds) * 1000.0).astype(np.int64) |
| metadata = { |
| "fps": 1000, |
| "frames_indices": absolute_ms.tolist(), |
| "total_num_frames": int(absolute_ms.max()) + 1, |
| "width": FRAME_WIDTH, |
| "height": FRAME_HEIGHT, |
| "duration": (int(absolute_ms.max()) + 1) / 1000.0, |
| "video_backend": "precomputed_absolute_pts", |
| } |
| return frames, metadata |
|
|
|
|
| def answer_one(model, processor, req: Request) -> tuple[str, dict]: |
| """Run the merged Qwen3.5 checkpoint on one official request.""" |
| frames, metadata = your_preprocessing(clip_path_for(req), req) |
| messages = [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": [{"type": "video"}, {"type": "text", "text": req.question}]}, |
| ] |
| text = processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, enable_thinking=False |
| ) |
| inputs = processor( |
| text=[text], |
| videos=[frames], |
| video_metadata=[metadata], |
| do_sample_frames=False, |
| return_tensors="pt", |
| ) |
| model_device = next(model.parameters()).device |
| inputs = inputs.to(model_device) |
| tokenizer = getattr(processor, "tokenizer", processor) |
| eos_token_ids = resolve_eos_token_ids(model, tokenizer) |
| im_start_token_ids = resolve_im_start_token_ids(tokenizer) |
| generation_kwargs = {"max_new_tokens": MAX_NEW_TOKENS, "do_sample": False} |
| if eos_token_ids: |
| generation_kwargs["eos_token_id"] = eos_token_ids if len(eos_token_ids) > 1 else eos_token_ids[0] |
| with torch.inference_mode(): |
| generated = model.generate(**inputs, **generation_kwargs) |
| if hasattr(generated, "sequences"): |
| generated = generated.sequences |
| new_token_rows = extract_new_token_ids( |
| generated, |
| inputs["input_ids"], |
| inputs.get("attention_mask"), |
| ) |
| trimmed = trim_generated_token_ids( |
| new_token_rows[0] if new_token_rows else [], |
| eos_token_ids=eos_token_ids, |
| im_start_token_ids=im_start_token_ids, |
| max_new_tokens=MAX_NEW_TOKENS, |
| ) |
| decoded = processor.batch_decode( |
| [trimmed["token_ids"]], skip_special_tokens=False, clean_up_tokenization_spaces=False |
| )[0] |
| answer = clean_decoded_answer(decoded) |
| audit = { |
| "raw_token_ids": trimmed["raw_token_ids"], |
| "stop_index": trimmed["stop_index"], |
| "eos_position": trimmed["eos_position"], |
| "eos_positions": trimmed["eos_positions"], |
| "stop_reason": trimmed["stop_reason"], |
| "hit_max_new_tokens": trimmed["hit_max_new_tokens"], |
| "new_token_count": len(trimmed["raw_token_ids"]), |
| } |
| return answer, audit |
|
|
|
|
| def run() -> int: |
| t_start = time.monotonic() |
| log.info("=== ORena SAVE FOCUS — inference start ===") |
|
|
| |
| |
| |
| |
| |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| log.info("Device: %s", device) |
| log_environment(device) |
|
|
| |
| log.info("--- Loading inputs ---") |
| for name in ("batch.json", "request.json", "FO_definitions.json"): |
| p = INPUT_PATH / name |
| status = f"{p.stat().st_size} bytes" if p.exists() else "MISSING" |
| log.info(" %-24s %s", name, status) |
| log.info(" %-24s %s", f"{VIDEO_DIR.name}/", f"{len(list(VIDEO_DIR.glob('*.mp4')))} clip(s)" if VIDEO_DIR.is_dir() else "MISSING") |
|
|
| |
| requests = load_requests(INPUT_PATH / "request.json") |
| if not requests: |
| log.error("request.json contains no requests") |
| return 1 |
| log.info("Batch of %d question(s)", len(requests)) |
|
|
| |
| |
| fo_definitions = json.loads((INPUT_PATH / "FO_definitions.json").read_text()) |
| log.info("FO definitions: %d chars", len(fo_definitions)) |
|
|
| |
| |
| |
| log.info("--- Loading model (once for the batch) ---") |
| if not MODEL_PATH.exists(): |
| log.error("Merged model directory not found: %s", MODEL_PATH) |
| return 1 |
| log.info("Merged model path: %s", MODEL_PATH) |
| processor = AutoProcessor.from_pretrained(MODEL_PATH, local_files_only=True) |
| model = Qwen3_5ForConditionalGeneration.from_pretrained( |
| MODEL_PATH, |
| torch_dtype=torch.bfloat16, |
| device_map="auto", |
| low_cpu_mem_usage=True, |
| local_files_only=True, |
| ) |
| model.eval() |
| n_params = sum(p.numel() for p in model.parameters()) |
| log.info("Model loaded and set to eval mode: %d parameters (setup %.2f s)", n_params, time.monotonic() - t_start) |
|
|
| if device.type == "cuda": |
| allocated = torch.cuda.memory_allocated(0) / 1024**2 |
| reserved = torch.cuda.memory_reserved(0) / 1024**2 |
| log.info(" VRAM after model load: %.1f MiB allocated / %.1f MiB reserved", allocated, reserved) |
|
|
| |
| log.info("--- Running inference over %d question(s) ---", len(requests)) |
| responses = [] |
| generation_audits = [] |
| n_failed = 0 |
| t_batch = time.monotonic() |
|
|
| for i, req in enumerate(requests, start=1): |
| log.info( |
| "[%d/%d] qID=%s videoID=%s window=[%.2f, %.2f] s (%.2f s)", |
| i, len(requests), req.qID, req.videoID, req.start_time, req.end_time, req.duration, |
| ) |
| t0 = time.monotonic() |
| try: |
| answer, audit = answer_one(model, processor, req) |
| except Exception: |
| |
| |
| n_failed += 1 |
| log.exception("[%d/%d] qID=%s failed; emitting empty answer", i, len(requests), req.qID) |
| answer = "" |
| audit = { |
| "raw_token_ids": [], |
| "stop_index": None, |
| "eos_position": None, |
| "eos_positions": [], |
| "stop_reason": "error", |
| "hit_max_new_tokens": False, |
| "new_token_count": 0, |
| } |
| latency = time.monotonic() - t0 |
|
|
| |
| |
| responses.append(Response(qID=req.qID, content=answer, latency=latency)) |
| generation_audits.append({"qID": req.qID, **audit}) |
| log.info( |
| "[%d/%d] qID=%s answered in %.4f s answer_len=%d stop_reason=%s", |
| i, |
| len(requests), |
| req.qID, |
| latency, |
| len(answer), |
| audit["stop_reason"], |
| ) |
|
|
| batch_seconds = time.monotonic() - t_batch |
| log.info( |
| "Inference complete: %d answered (%d failed) in %.2f s (%.3f s/question)", |
| len(responses), n_failed, batch_seconds, batch_seconds / max(len(responses), 1), |
| ) |
|
|
| if device.type == "cuda": |
| allocated = torch.cuda.memory_allocated(0) / 1024**2 |
| log.debug(" VRAM after inference: %.1f MiB allocated", allocated) |
|
|
| |
| log.info("--- Writing output ---") |
| OUTPUT_PATH.mkdir(parents=True, exist_ok=True) |
| output_path = OUTPUT_PATH / "answer.json" |
| save_items(responses, output_path) |
| log.info("Wrote %d response(s) to %s", len(responses), output_path) |
| audit_path = OUTPUT_PATH / "generation_audit.json" |
| audit_path.write_text(json.dumps(generation_audits, ensure_ascii=False, indent=2) + "\n") |
| log.info("Wrote generation audit for %d response(s) to %s", len(generation_audits), audit_path) |
| log.info("=== inference done in %.2f s total ===", time.monotonic() - t_start) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(run()) |
|
|
