Initial fp32 release

Browse files

Files changed (13) hide show

.gitattributes +1 -0
chat_template.jinja +154 -0
config.json +118 -0
configuration_argus.py +71 -0
eval_vidore_v1_v2.py +224 -0
model-00001-of-00002.safetensors +3 -0
model-00002-of-00002.safetensors +3 -0
model.safetensors.index.json +666 -0
modeling_argus.py +559 -0
processing_argus.py +260 -0
processor_config.json +64 -0
tokenizer.json +3 -0
tokenizer_config.json +33 -0

.gitattributes CHANGED Viewed

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja ADDED Viewed

	@@ -0,0 +1,154 @@

+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is true %}
+        {{- '<think>\n' }}
+    {%- else %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json ADDED Viewed

	@@ -0,0 +1,118 @@

+{
+  "architectures": [
+    "ArgusForRetrieval"
+  ],
+  "dtype": "bfloat16",
+  "image_token_id": 248056,
+  "model_type": "argus_colqwen35",
+  "rope_parameters": {},
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 6144,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 16,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 248320
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.6.1",
+  "use_cache": false,
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 24,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1024,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "model_type": "qwen3_5_vision",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 2048,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053,
+  "auto_map": {
+    "AutoConfig": "configuration_argus.ArgusConfig",
+    "AutoModel": "modeling_argus.ArgusForRetrieval",
+    "AutoProcessor": "processing_argus.ArgusProcessor"
+  },
+  "retrieval_dim": 1024,
+  "num_specialists": 4,
+  "top_k_experts": 2,
+  "region_size": 4,
+  "router_layer_index": -5,
+  "router_temperature": 0.8,
+  "router_noise_std": 0.0,
+  "mask_non_image_embeddings": true,
+  "shared_gate_init": 0.0,
+  "specialist_gate_init": 0.0
+}

configuration_argus.py ADDED Viewed

	@@ -0,0 +1,71 @@

+"""Argus: Region-Aware Query-Conditioned Mixture of Experts for Visual Document Retrieval.
+Config class. Subclasses the Qwen3.5-VL config and adds the Argus-specific
+retrieval + MoE hyperparameters. Used by ``AutoConfig.from_pretrained`` via the
+``auto_map`` field in ``config.json`` (requires ``trust_remote_code=True``).
+"""
+from __future__ import annotations
+try:
+    from transformers.models.qwen3_5 import Qwen3_5Config as _BackboneConfig
+except ImportError:
+    try:
+        from transformers.models.qwen3_5 import Qwen35Config as _BackboneConfig
+    except ImportError as exc:
+        raise ImportError(
+            "Argus requires a transformers build that exposes the Qwen3.5 VL "
+            "classes (transformers.models.qwen3_5). Upgrade to transformers "
+            ">= 4.57.0.dev0."
+        ) from exc
+class ArgusConfig(_BackboneConfig):
+    """Top-level config for Argus-Colqwen3.5-9B.
+    Holds the standard Qwen3.5-VL fields (text_config, vision_config, image
+    token ids, etc.) plus Argus-specific retrieval + MoE knobs:
+    - ``retrieval_dim``: output dimensionality of the multi-vector retrieval
+      head (``custom_text_proj``). Default: 768.
+    - ``num_specialists``: number of latent spatial experts in the MoE stack.
+    - ``top_k_experts``: sparsity of the router (top-k routing).
+    - ``region_size``: spatial pooling window (patches) for region tokens.
+    - ``router_layer_index``: hidden-state layer used as input to the router.
+    - ``router_temperature``: softmax temperature of the router.
+    - ``mask_non_image_embeddings``: zero out embedding positions that are
+      not image tokens at encode time (document side).
+    - ``shared_gate_init`` / ``specialist_gate_init``: logit-space init for
+      the gate scalars (sigmoid of these multiplies shared/specialist expert
+      contributions).
+    """
+    model_type = "argus_colqwen35"
+    def __init__(
+        self,
+        retrieval_dim: int = 768,
+        num_specialists: int = 4,
+        top_k_experts: int = 2,
+        region_size: int = 4,
+        router_layer_index: int = -5,
+        router_temperature: float = 0.8,
+        router_noise_std: float = 0.0,
+        mask_non_image_embeddings: bool = True,
+        shared_gate_init: float = 0.0,
+        specialist_gate_init: float = 0.0,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.retrieval_dim = int(retrieval_dim)
+        self.num_specialists = int(num_specialists)
+        self.top_k_experts = int(top_k_experts)
+        self.region_size = int(region_size)
+        self.router_layer_index = int(router_layer_index)
+        self.router_temperature = float(router_temperature)
+        self.router_noise_std = float(router_noise_std)
+        self.mask_non_image_embeddings = bool(mask_non_image_embeddings)
+        self.shared_gate_init = float(shared_gate_init)
+        self.specialist_gate_init = float(specialist_gate_init)
+__all__ = ["ArgusConfig"]

eval_vidore_v1_v2.py ADDED Viewed

	@@ -0,0 +1,224 @@

+#!/usr/bin/env python3
+"""Evaluate Argus-Colqwen3.5-9B on ViDoRe V1 + V2 using the official
+``vidore-benchmark`` library straight from the HuggingFace hub.
+Why this wrapper exists
+-----------------------
+The reference evaluators live in https://github.com/illuin-tech/vidore-benchmark
+— every ColPali / Nemotron / vidore leaderboard submission is scored against
+``ViDoReEvaluatorQA`` / ``ViDoReEvaluatorBEIR``. By delegating to those
+evaluators here (instead of re-implementing nDCG/Recall/MRR locally) we
+guarantee:
+- ``None`` queries are filtered correctly (Shift, all SyntheticDocQA subsets).
+- The full image corpus is preserved (distractors stay in the retrieval pool).
+- MTEB-style metrics (ndcg/map/recall/precision/mrr at every k) match the
+  canonical leaderboard numbers bit-for-bit.
+Usage
+-----
+    pip install vidore-benchmark  # or: pip install git+https://github.com/illuin-tech/vidore-benchmark
+    python eval_vidore_v1_v2.py \\
+        --model ./argus-colqwen3.5-9b-v0 \\
+        --benchmarks v1 v2 \\
+        --batch-query 4 \\
+        --batch-passage 2
+Use ``--model DataScience-UIBK/Argus-Colqwen3.5-9B-v0`` once uploaded.
+"""
+from __future__ import annotations
+import argparse
+import json
+from pathlib import Path
+from typing import Dict
+import torch
+# ---------------------- ViDoRe dataset catalog ---------------------- #
+# ViDoRe V1 (QA format). Each HF dataset has a single ``test`` split with
+# columns: query, image, image_filename. Some rows contain ``query=None``
+# (distractors); the library handles this.
+V1_DATASETS: Dict[str, str] = {
+    "ArxivQ": "vidore/arxivqa_test_subsampled",
+    "DocQ":   "vidore/docvqa_test_subsampled",
+    "InfoQ":  "vidore/infovqa_test_subsampled",
+    "TabF":   "vidore/tabfquad_test_subsampled",
+    "TATQ":   "vidore/tatdqa_test",
+    "Shift":  "vidore/shiftproject_test",
+    "AI":     "vidore/syntheticDocQA_artificial_intelligence_test",
+    "Energy": "vidore/syntheticDocQA_energy_test",
+    "Gov":    "vidore/syntheticDocQA_government_reports_test",
+    "Health": "vidore/syntheticDocQA_healthcare_industry_test",
+}
+# ViDoRe V2 (BEIR format). Each HF repo exposes 3 dataset configs:
+# ``corpus`` (images + corpus-id), ``queries`` (query text + query-id), and
+# ``qrels`` (query-id, corpus-id, score). The library's ``ViDoReEvaluatorBEIR``
+# expects that exact shape.
+V2_DATASETS: Dict[str, str] = {
+    "MIT_Biomedical_Multi":        "vidore/biomedical_lectures_v2",
+    "Economics_Macro_Multi":       "vidore/economics_reports_v2",
+    "ESG_Restaurant_Human_EN":     "vidore/esg_reports_human_labeled_v2",
+    "ESG_Restaurant_Synth_Multi":  "vidore/esg_reports_v2",
+}
+# ---------------------- helpers ---------------------- #
+def _load_model_and_processor(args: argparse.Namespace):
+    from transformers import AutoModel, AutoProcessor
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    print(f"[eval] loading model: {args.model} ({args.dtype}, attn={args.attn_implementation})")
+    # ``dtype`` on transformers >= 4.57; older builds still use ``torch_dtype``.
+    load_kwargs = {"trust_remote_code": True, "attn_implementation": args.attn_implementation}
+    try:
+        model = AutoModel.from_pretrained(args.model, dtype=dtype, **load_kwargs).eval().cuda()
+    except TypeError:
+        model = AutoModel.from_pretrained(args.model, torch_dtype=dtype, **load_kwargs).eval().cuda()
+    processor = AutoProcessor.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+        max_num_visual_tokens=args.max_num_visual_tokens,
+    )
+    return model, processor
+class _EmbeddingOnlyWrapper(torch.nn.Module):
+    """Adapter that exposes the plain embeddings tensor to vidore-benchmark.
+    ``VisionRetriever.forward_queries`` / ``forward_passages`` call
+    ``self.model(**batch).to("cpu")``, i.e. they assume the model returns a
+    Tensor. ``ArgusForRetrieval.forward`` returns an ``ArgusOutput`` dataclass
+    (embeddings + region_embeddings + routing info) to keep the MoE analysis
+    surface. This wrapper unwraps ``.embeddings`` so the library sees the
+    expected shape without us having to touch the model class.
+    """
+    def __init__(self, inner: torch.nn.Module):
+        super().__init__()
+        self.inner = inner
+    def __getattr__(self, name):
+        # Delegate .device / .dtype / .eval() / etc. to the wrapped model.
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.inner, name)
+    def forward(self, **kwargs) -> torch.Tensor:
+        return self.inner(**kwargs).embeddings
+def _build_retriever(model, processor):
+    from vidore_benchmark.retrievers import VisionRetriever
+    wrapped = _EmbeddingOnlyWrapper(model).eval()
+    # Older vidore-benchmark releases don't accept ``num_workers`` at all;
+    # newer ones do. Try-with-kwarg for portability.
+    try:
+        return VisionRetriever(model=wrapped, processor=processor, num_workers=0)
+    except TypeError:
+        return VisionRetriever(model=wrapped, processor=processor)
+def _eval_v1(retriever, args: argparse.Namespace) -> Dict[str, Dict[str, float]]:
+    from datasets import load_dataset
+    from vidore_benchmark.evaluation.vidore_evaluators import ViDoReEvaluatorQA
+    evaluator = ViDoReEvaluatorQA(retriever)
+    results: Dict[str, Dict[str, float]] = {}
+    print("\n========== V1 ==========")
+    for short, repo_id in V1_DATASETS.items():
+        if args.datasets and short not in args.datasets:
+            continue
+        print(f"\n[V1:{short}] {repo_id}")
+        ds = load_dataset(repo_id, split="test")
+        metrics = evaluator.evaluate_dataset(
+            ds,
+            batch_query=args.batch_query,
+            batch_passage=args.batch_passage,
+            batch_score=args.batch_score,
+        )
+        results[short] = metrics
+        print(f"  nDCG@5 = {metrics.get('ndcg_at_5', 0.0):.4f}")
+    return results
+def _eval_v2(retriever, args: argparse.Namespace) -> Dict[str, Dict[str, float]]:
+    from datasets import load_dataset
+    from vidore_benchmark.evaluation.vidore_evaluators import ViDoReEvaluatorBEIR
+    evaluator = ViDoReEvaluatorBEIR(retriever)
+    results: Dict[str, Dict[str, float]] = {}
+    print("\n========== V2 ==========")
+    for short, repo_id in V2_DATASETS.items():
+        if args.datasets and short not in args.datasets:
+            continue
+        print(f"\n[V2:{short}] {repo_id}")
+        ds = {
+            "corpus":  load_dataset(repo_id, "corpus",  split="test"),
+            "queries": load_dataset(repo_id, "queries", split="test"),
+            "qrels":   load_dataset(repo_id, "qrels",   split="test"),
+        }
+        metrics = evaluator.evaluate_dataset(
+            ds,
+            batch_query=args.batch_query,
+            batch_passage=args.batch_passage,
+            batch_score=args.batch_score,
+        )
+        results[short] = metrics
+        print(f"  nDCG@5 = {metrics.get('ndcg_at_5', 0.0):.4f}")
+    return results
+# ---------------------- main ---------------------- #
+def run(args: argparse.Namespace) -> None:
+    model, processor = _load_model_and_processor(args)
+    retriever = _build_retriever(model, processor)
+    all_results: Dict[str, Dict[str, Dict[str, float]]] = {"v1": {}, "v2": {}}
+    if "v1" in args.benchmarks:
+        all_results["v1"] = _eval_v1(retriever, args)
+    if "v2" in args.benchmarks:
+        all_results["v2"] = _eval_v2(retriever, args)
+    # Summary
+    print("\n========== summary ==========")
+    for bench, per_ds in all_results.items():
+        if not per_ds:
+            continue
+        avg = sum(m.get("ndcg_at_5", 0.0) for m in per_ds.values()) / max(len(per_ds), 1)
+        print(f"{bench.upper()} avg nDCG@5 = {avg:.4f}   ({len(per_ds)} datasets)")
+    if args.output_json:
+        Path(args.output_json).write_text(json.dumps(all_results, indent=2, default=float))
+        print(f"[eval] saved: {args.output_json}")
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", required=True,
+                   help="HF repo id or local release folder.")
+    p.add_argument("--benchmarks", nargs="+", default=["v1", "v2"], choices=["v1", "v2"])
+    p.add_argument("--datasets", nargs="*", default=None,
+                   help="Optional subset by short key (e.g. ArxivQ DocQ Shift).")
+    p.add_argument("--batch-query", type=int, default=4)
+    p.add_argument("--batch-passage", type=int, default=2)
+    p.add_argument("--batch-score", type=int, default=4)
+    p.add_argument("--max-num-visual-tokens", type=int, default=2048)
+    p.add_argument("--attn-implementation", default="flash_attention_2",
+                   choices=["flash_attention_2", "sdpa", "eager"])
+    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    p.add_argument("--output-json", default=None)
+    return p.parse_args()
+if __name__ == "__main__":
+    run(parse_args())

model-00001-of-00002.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:cb248f0ce6c71b0341ecabe41b00e4cda90f17a14c3c4b028b62ffd235eb4e58
+size 4981231800

model-00002-of-00002.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:128445b209dc65b659e62ca73f07aea67eed7f4cf5c7282164387a01611aa3e0
+size 4308333160

model.safetensors.index.json ADDED Viewed

	@@ -0,0 +1,666 @@

+{
+  "metadata": {
+    "total_size": 9289489688
+  },
+  "weight_map": {
+    "visual.patch_embed.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.patch_embed.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.pos_embed.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.0.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.1.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.2.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.3.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.4.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.5.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.6.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.7.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.8.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.9.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.10.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.11.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.12.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.13.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.14.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.15.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.16.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.17.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.18.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.19.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.20.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.21.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.22.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.norm1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.norm1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.norm2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.norm2.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.mlp.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.mlp.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.mlp.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.blocks.23.mlp.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "visual.merger.norm.weight": "model-00001-of-00002.safetensors",
+    "visual.merger.norm.bias": "model-00001-of-00002.safetensors",
+    "visual.merger.linear_fc1.weight": "model-00001-of-00002.safetensors",
+    "visual.merger.linear_fc1.bias": "model-00001-of-00002.safetensors",
+    "visual.merger.linear_fc2.weight": "model-00001-of-00002.safetensors",
+    "visual.merger.linear_fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.dt_bias": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.A_log": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.conv1d.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_qkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_z.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_b.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_a.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.dt_bias": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.A_log": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.conv1d.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_qkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_z.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_b.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_a.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.dt_bias": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.A_log": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.conv1d.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_qkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_z.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_b.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_a.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.dt_bias": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.A_log": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.conv1d.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_qkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_z.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_b.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_a.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.dt_bias": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.A_log": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.conv1d.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_qkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_z.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_b.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_a.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.dt_bias": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.A_log": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.conv1d.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_qkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_z.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_b.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_a.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "language_model.layers.7.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.7.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.8.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.9.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.10.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.11.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.12.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.13.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.14.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.15.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.16.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.17.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.18.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.19.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.20.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.21.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.dt_bias": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.A_log": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.conv1d.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_qkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_z.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_b.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_a.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "language_model.norm.weight": "model-00002-of-00002.safetensors",
+    "custom_text_proj.weight": "model-00002-of-00002.safetensors",
+    "custom_text_proj.bias": "model-00002-of-00002.safetensors",
+    "shared_expert.net.0.weight": "model-00002-of-00002.safetensors",
+    "shared_expert.net.0.bias": "model-00002-of-00002.safetensors",
+    "shared_expert.net.1.weight": "model-00002-of-00002.safetensors",
+    "shared_expert.net.1.bias": "model-00002-of-00002.safetensors",
+    "shared_expert.net.3.weight": "model-00002-of-00002.safetensors",
+    "shared_expert.net.3.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.0.net.0.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.0.net.0.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.0.net.1.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.0.net.1.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.0.net.3.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.0.net.3.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.1.net.0.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.1.net.0.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.1.net.1.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.1.net.1.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.1.net.3.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.1.net.3.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.2.net.0.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.2.net.0.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.2.net.1.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.2.net.1.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.2.net.3.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.2.net.3.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.3.net.0.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.3.net.0.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.3.net.1.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.3.net.1.bias": "model-00002-of-00002.safetensors",
+    "latent_experts.3.net.3.weight": "model-00002-of-00002.safetensors",
+    "latent_experts.3.net.3.bias": "model-00002-of-00002.safetensors",
+    "region_router.0.weight": "model-00002-of-00002.safetensors",
+    "region_router.0.bias": "model-00002-of-00002.safetensors",
+    "region_router.1.weight": "model-00002-of-00002.safetensors",
+    "region_router.1.bias": "model-00002-of-00002.safetensors",
+    "region_router.3.weight": "model-00002-of-00002.safetensors",
+    "region_router.3.bias": "model-00002-of-00002.safetensors",
+    "region_coord_proj.weight": "model-00002-of-00002.safetensors",
+    "query_context_proj.weight": "model-00002-of-00002.safetensors",
+    "gate_scalars.shared": "model-00002-of-00002.safetensors",
+    "gate_scalars.specialist": "model-00002-of-00002.safetensors"
+  }
+}

modeling_argus.py ADDED Viewed

	@@ -0,0 +1,559 @@

+"""Argus: Region-Aware Query-Conditioned Mixture of Experts for Visual Document Retrieval.
+Self-contained model implementation for the Argus-Colqwen3.5-9B release.
+Usage
+-----
+>>> from transformers import AutoModel, AutoProcessor
+>>> model = AutoModel.from_pretrained(
+...     "DataScience-UIBK/Argus-Colqwen3.5-9B-v0",
+...     trust_remote_code=True,
+...     torch_dtype="bfloat16",
+... ).eval().cuda()
+>>> proc  = AutoProcessor.from_pretrained(
+...     "DataScience-UIBK/Argus-Colqwen3.5-9B-v0",
+...     trust_remote_code=True,
+... )
+>>> q_emb = model.encode_queries(proc, ["what is the revenue in 2019?"])
+>>> d_emb = model.encode_images(proc, [pil_image_1, pil_image_2])
+>>> scores = model.score(q_emb, d_emb)   # shape [num_queries, num_docs]
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from math import ceil
+from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.utils import ModelOutput
+try:
+    from transformers.models.qwen3_5 import Qwen3_5Config, Qwen3_5Model
+except ImportError:
+    try:
+        from transformers.models.qwen3_5 import Qwen35Config as Qwen3_5Config
+        from transformers.models.qwen3_5 import Qwen35Model as Qwen3_5Model
+    except ImportError as exc:
+        raise ImportError(
+            "Argus requires a transformers build that exposes the Qwen3.5 VL "
+            "classes (transformers.models.qwen3_5). Upgrade to transformers "
+            ">= 4.57.0.dev0."
+        ) from exc
+from .configuration_argus import ArgusConfig
+# --------------------------------------------------------------------------- #
+# Output container
+# --------------------------------------------------------------------------- #
+@dataclass
+class ArgusOutput(ModelOutput):
+    """Output of :meth:`ArgusForRetrieval.forward`.
+    Attributes:
+        embeddings: multi-vector token embeddings [B, T, D]. Use ``score`` /
+            ``score_multi_vector`` against queries encoded the same way.
+        region_embeddings: pooled region-level document embeddings [B, R, D]
+            (only populated when images are in the batch).
+        region_mask: valid mask for region_embeddings, shape [B, R].
+        routing_logits: raw MoE router logits [B, R, E] (per-region, per-expert).
+    """
+    embeddings: torch.Tensor
+    region_embeddings: Optional[torch.Tensor] = None
+    region_mask: Optional[torch.Tensor] = None
+    routing_logits: Optional[torch.Tensor] = None
+# --------------------------------------------------------------------------- #
+# MoE building blocks
+# --------------------------------------------------------------------------- #
+def _ceil_to_multiple(value: int, multiple: int) -> int:
+    return int(ceil(value / multiple) * multiple)
+class SharedDenseExpert(nn.Module):
+    """Shared expert applied to every spatial location."""
+    def __init__(self, hidden_dim: int, expansion: int = 4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim * expansion),
+            nn.GELU(),
+            nn.Linear(hidden_dim * expansion, hidden_dim),
+        )
+    def forward(self, grid: torch.Tensor) -> torch.Tensor:
+        return self.net(grid)
+class LatentSpatialExpert(nn.Module):
+    """One of ``num_specialists`` region-level experts routed by the query."""
+    def __init__(self, hidden_dim: int, expansion: int = 2):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim * expansion),
+            nn.GELU(),
+            nn.Linear(hidden_dim * expansion, hidden_dim),
+        )
+    def forward(self, grid: torch.Tensor) -> torch.Tensor:
+        return self.net(grid)
+class GateScalars(nn.Module):
+    """Two learnable scalars whose sigmoids weight shared / specialist expert
+    contributions onto the final hidden states.
+    """
+    def __init__(self, shared_init: float = 0.0, specialist_init: float = 0.0):
+        super().__init__()
+        self.shared = nn.Parameter(torch.tensor(float(shared_init), dtype=torch.float32))
+        self.specialist = nn.Parameter(torch.tensor(float(specialist_init), dtype=torch.float32))
+    def _apply(self, fn):  # noqa: D401 - keep fp32 even after .to(dtype)
+        super()._apply(fn)
+        for name in ("shared", "specialist"):
+            param = getattr(self, name)
+            if param.dtype != torch.float32:
+                param.data = param.data.to(torch.float32)
+        return self
+    def sigmoid(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        return torch.sigmoid(self.shared), torch.sigmoid(self.specialist)
+# --------------------------------------------------------------------------- #
+# Argus model
+# --------------------------------------------------------------------------- #
+class ArgusForRetrieval(Qwen3_5Model):
+    """Argus multi-vector visual document retriever.
+    Structure:
+    - Backbone: Qwen3.5-VL (9B) — produces per-token hidden states.
+    - Region pool: non-overlapping ``region_size × region_size`` blocks over
+      the vision-token grid; gives a compact region-level view.
+    - Router: per-region MLP → ``num_specialists`` logits; the query (if given)
+      biases the logits via ``query_context_proj``. Top-k sparse softmax.
+    - Experts: one shared expert (applied everywhere) + ``num_specialists``
+      latent spatial experts (per-region weighted sum).
+    - Fusion: ``final_hidden = final_hidden + σ(gate_shared) · shared_expert
+      + σ(gate_specialist) · specialist_sum``.
+    - Retrieval head: ``custom_text_proj`` projects fused hidden states to
+      ``retrieval_dim`` multi-vectors, L2-normalized.
+    - Query side: no MoE; just backbone + ``custom_text_proj``.
+    The user-facing helpers are ``encode_images``, ``encode_queries``, and
+    ``score`` (MaxSim). All live on this class so a downstream user can do
+    everything via ``model.<method>``.
+    """
+    config_class = ArgusConfig
+    main_input_name: ClassVar[str] = "input_ids"
+    def __init__(self, config: Union[ArgusConfig, Qwen3_5Config], **kwargs):
+        # Accept either an ArgusConfig or a plain Qwen3_5Config with extra attrs
+        # (transformers sometimes hands us a base-class instance during Auto*
+        # dispatch before config_class kicks in).
+        if not isinstance(config, ArgusConfig):
+            promoted = ArgusConfig(**config.to_dict())
+            config = promoted
+        dtype = kwargs.pop("dtype", kwargs.pop("torch_dtype", None))
+        attn_impl = kwargs.pop("attn_implementation", None)
+        use_cache = kwargs.pop("use_cache", None)
+        if hasattr(config, "text_config") and getattr(config.text_config, "rope_scaling", None) is None:
+            config.text_config.rope_scaling = {}
+        if getattr(config, "rope_scaling", None) is None:
+            config.rope_scaling = {}
+        super().__init__(config=config)
+        hidden_size = getattr(config, "hidden_size", None) or getattr(config.text_config, "hidden_size", None)
+        if hidden_size is None:
+            raise ValueError("Argus: could not determine backbone hidden_size from config.")
+        self.retrieval_dim = int(config.retrieval_dim)
+        self.num_specialists = int(config.num_specialists)
+        self.top_k_experts = max(1, min(int(config.top_k_experts), self.num_specialists))
+        self.region_size = int(config.region_size)
+        self.router_layer_index = int(config.router_layer_index)
+        self.router_temperature = float(config.router_temperature)
+        self.router_noise_std = float(config.router_noise_std)
+        self.mask_non_image_embeddings = bool(config.mask_non_image_embeddings)
+        self.spatial_merge_size = getattr(config.vision_config, "spatial_merge_size", 1)
+        self.padding_side = "left"
+        self.custom_text_proj = nn.Linear(hidden_size, self.retrieval_dim)
+        self.shared_expert = SharedDenseExpert(hidden_size)
+        self.latent_experts = nn.ModuleList(
+            LatentSpatialExpert(hidden_size) for _ in range(self.num_specialists)
+        )
+        self.region_router = nn.Sequential(
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, self.num_specialists),
+        )
+        self.region_coord_proj = nn.Linear(4, hidden_size, bias=False)
+        self.query_context_proj = nn.Linear(self.retrieval_dim, hidden_size, bias=False)
+        self.gate_scalars = GateScalars(
+            shared_init=config.shared_gate_init,
+            specialist_init=config.specialist_gate_init,
+        )
+        self.post_init()
+        if dtype is not None:
+            self.to(dtype=dtype)
+        if use_cache is not None:
+            self.config.use_cache = use_cache
+        if attn_impl is not None and hasattr(self, "set_attn_implementation"):
+            self.set_attn_implementation(attn_impl)
+    # ----------------------------------------------------------------- #
+    # Forward
+    # ----------------------------------------------------------------- #
+    def build_query_router_context(
+        self,
+        query_embeddings: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Pool query multi-vectors into one normalized vector per query.
+        Used to bias the MoE router when the query is known at doc-encode
+        time (cross-encoder-style, optional). Safe to call with query-only
+        outputs of :meth:`forward`.
+        """
+        if attention_mask is None:
+            pooled = query_embeddings.mean(dim=1)
+        else:
+            weights = attention_mask.unsqueeze(-1).to(query_embeddings.dtype)
+            pooled = (query_embeddings * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1.0)
+        return pooled / pooled.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+    def forward(self, *args, **kwargs) -> ArgusOutput:
+        """Run backbone + MoE + retrieval head.
+        Inputs follow the standard Qwen3-VL processor outputs:
+        ``input_ids``, ``attention_mask``, and (for images) ``pixel_values``
+        + ``image_grid_thw``. ``query_context`` is optional and, when given,
+        biases the router for this batch.
+        """
+        kwargs.pop("region_labels", None)
+        kwargs.pop("region_mask", None)
+        query_context = kwargs.pop("query_context", None)
+        image_grid_thw = kwargs.get("image_grid_thw")
+        # Processor may return per-image padded pixel tensors; the backbone
+        # wants them flat-concatenated.
+        if "pixel_values" in kwargs and image_grid_thw is not None:
+            offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2]
+            kwargs["pixel_values"] = torch.cat(
+                [pv[:off] for pv, off in zip(kwargs["pixel_values"], offsets)],
+                dim=0,
+            )
+        kwargs.pop("return_dict", True)
+        kwargs.pop("output_hidden_states", None)
+        kwargs.pop("use_cache", None)
+        outputs = super().forward(
+            *args,
+            **kwargs,
+            use_cache=False,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        final_hidden = outputs.last_hidden_state
+        router_hidden = outputs.hidden_states[self.router_layer_index]
+        del outputs.hidden_states
+        attention_mask = kwargs["attention_mask"]
+        region_embeddings_list: List[torch.Tensor] = []
+        routing_logits_list: List[torch.Tensor] = []
+        routing_mask_list: List[torch.Tensor] = []
+        if "pixel_values" in kwargs and "input_ids" in kwargs:
+            image_mask = kwargs["input_ids"] == self.config.image_token_id
+            for batch_idx in range(final_hidden.size(0)):
+                image_positions = image_mask[batch_idx].nonzero(as_tuple=False).squeeze(-1)
+                if image_positions.numel() == 0:
+                    region_embeddings_list.append(final_hidden.new_zeros(0, self.retrieval_dim))
+                    routing_logits_list.append(final_hidden.new_zeros(0, self.num_specialists))
+                    routing_mask_list.append(final_hidden.new_zeros(0, dtype=torch.bool))
+                    continue
+                grid_t = int(image_grid_thw[batch_idx, 0].item())
+                raw_grid_h = int(image_grid_thw[batch_idx, 1].item())
+                raw_grid_w = int(image_grid_thw[batch_idx, 2].item())
+                grid_h = max(1, raw_grid_h // self.spatial_merge_size)
+                grid_w = max(1, raw_grid_w // self.spatial_merge_size)
+                num_image_tokens = min(grid_t * grid_h * grid_w, image_positions.numel())
+                image_positions = image_positions[:num_image_tokens]
+                early_grid = router_hidden[batch_idx, image_positions].view(grid_t, grid_h, grid_w, -1).mean(dim=0)
+                final_grid = final_hidden[batch_idx, image_positions].view(grid_t, grid_h, grid_w, -1).mean(dim=0)
+                query_context_i = None if query_context is None else query_context[batch_idx]
+                fused_grid, pooled_regions, pooled_mask, logits = self._apply_query_conditioned_moe(
+                    early_grid=early_grid,
+                    final_grid=final_grid,
+                    query_context=query_context_i,
+                )
+                fused_tokens = (
+                    fused_grid.unsqueeze(0)
+                    .expand(grid_t, -1, -1, -1)
+                    .reshape(num_image_tokens, -1)
+                    .to(final_hidden.dtype)
+                )
+                final_hidden[batch_idx, image_positions] = fused_tokens
+                projected_regions = self.custom_text_proj(pooled_regions)
+                projected_regions = projected_regions / projected_regions.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+                region_embeddings_list.append(projected_regions)
+                routing_logits_list.append(logits)
+                routing_mask_list.append(pooled_mask)
+        embeddings = self.custom_text_proj(final_hidden)
+        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+        embeddings = embeddings * attention_mask.unsqueeze(-1)
+        if "pixel_values" in kwargs and self.mask_non_image_embeddings and "input_ids" in kwargs:
+            embeddings = embeddings * (kwargs["input_ids"] == self.config.image_token_id).unsqueeze(-1)
+        region_embeddings, padded_routing_logits, padded_routing_mask = self._pad_regions(
+            region_embeddings_list,
+            routing_logits_list,
+            routing_mask_list,
+            device=embeddings.device,
+            dtype=embeddings.dtype,
+        )
+        return ArgusOutput(
+            embeddings=embeddings,
+            region_embeddings=region_embeddings,
+            region_mask=padded_routing_mask,
+            routing_logits=padded_routing_logits,
+        )
+    # ----------------------------------------------------------------- #
+    # MoE internals
+    # ----------------------------------------------------------------- #
+    def _topk_sparse_probs(self, routing_logits: torch.Tensor) -> torch.Tensor:
+        logits = routing_logits.float()
+        if self.training and self.router_noise_std > 0:
+            logits = logits + self.router_noise_std * torch.randn_like(logits)
+        if self.top_k_experts >= self.num_specialists:
+            return F.softmax(logits / max(self.router_temperature, 1e-6), dim=-1).to(routing_logits.dtype)
+        topk_values, topk_indices = torch.topk(logits, k=self.top_k_experts, dim=-1)
+        sparse_logits = torch.full_like(logits, float("-inf"))
+        sparse_logits.scatter_(-1, topk_indices, topk_values)
+        probs = F.softmax(sparse_logits / max(self.router_temperature, 1e-6), dim=-1)
+        return probs.to(routing_logits.dtype)
+    def _apply_query_conditioned_moe(
+        self,
+        early_grid: torch.Tensor,
+        final_grid: torch.Tensor,
+        query_context: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        region_tokens, pooled_mask, coords, region_shape = self._pool_regions(early_grid)
+        router_input = region_tokens + self.region_coord_proj(coords.to(region_tokens.dtype))
+        if query_context is not None:
+            query_bias = self.query_context_proj(query_context.to(region_tokens.dtype)).unsqueeze(0)
+            router_input = router_input + query_bias
+        routing_logits = self.region_router(router_input)
+        routing_probs = self._topk_sparse_probs(routing_logits)
+        shared_out = self.shared_expert(final_grid)
+        specialist_outputs = torch.stack([expert(final_grid) for expert in self.latent_experts], dim=-2)
+        patch_probs = self._broadcast_region_probs(routing_probs, region_shape, final_grid.shape[:2])
+        specialist_out = (specialist_outputs * patch_probs.unsqueeze(-1)).sum(dim=-2)
+        shared_sig, specialist_sig = self.gate_scalars.sigmoid()
+        fused_grid = (
+            final_grid
+            + shared_sig.to(final_grid.dtype) * shared_out
+            + specialist_sig.to(final_grid.dtype) * specialist_out
+        )
+        pooled_regions, pooled_region_mask, _, _ = self._pool_regions(fused_grid)
+        return fused_grid, pooled_regions, pooled_region_mask, routing_logits
+    def _pool_regions(
+        self,
+        grid: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Tuple[int, int]]:
+        h, w, dim = grid.shape
+        rs = self.region_size
+        hp = _ceil_to_multiple(h, rs)
+        wp = _ceil_to_multiple(w, rs)
+        padded = grid.new_zeros(hp, wp, dim)
+        padded[:h, :w] = grid
+        valid = grid.new_zeros(hp, wp, 1)
+        valid[:h, :w] = 1
+        num_h = hp // rs
+        num_w = wp // rs
+        blocks = padded.view(num_h, rs, num_w, rs, dim).permute(0, 2, 1, 3, 4).reshape(num_h * num_w, rs * rs, dim)
+        valid_blocks = valid.view(num_h, rs, num_w, rs, 1).permute(0, 2, 1, 3, 4).reshape(num_h * num_w, rs * rs, 1)
+        counts = valid_blocks.sum(dim=1).clamp_min(1.0)
+        pooled = (blocks * valid_blocks).sum(dim=1) / counts
+        mask = counts.squeeze(-1) > 0.5
+        coords = []
+        for ry in range(num_h):
+            for rx in range(num_w):
+                y0 = (ry * rs) / max(h, 1)
+                x0 = (rx * rs) / max(w, 1)
+                y1 = min((ry + 1) * rs, h) / max(h, 1)
+                x1 = min((rx + 1) * rs, w) / max(w, 1)
+                coords.append([x0, y0, x1, y1])
+        coord_tensor = torch.tensor(coords, device=grid.device, dtype=grid.dtype)
+        return pooled, mask, coord_tensor, (num_h, num_w)
+    def _broadcast_region_probs(
+        self,
+        region_probs: torch.Tensor,
+        region_shape: Tuple[int, int],
+        grid_shape: Tuple[int, int],
+    ) -> torch.Tensor:
+        num_h, num_w = region_shape
+        h, w = grid_shape
+        rs = self.region_size
+        hp = num_h * rs
+        wp = num_w * rs
+        probs = region_probs.view(num_h, num_w, self.num_specialists)
+        probs = probs[:, :, None, None, :].expand(num_h, num_w, rs, rs, self.num_specialists)
+        probs = probs.permute(0, 2, 1, 3, 4).reshape(hp, wp, self.num_specialists)
+        return probs[:h, :w]
+    def _pad_regions(
+        self,
+        region_embeddings_list: List[torch.Tensor],
+        routing_logits_list: List[torch.Tensor],
+        routing_mask_list: List[torch.Tensor],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        if not region_embeddings_list:
+            return None, None, None
+        max_regions = max((regions.size(0) for regions in region_embeddings_list), default=0)
+        if max_regions == 0:
+            batch_size = len(region_embeddings_list)
+            return (
+                torch.zeros(batch_size, 0, self.retrieval_dim, device=device, dtype=dtype),
+                torch.zeros(batch_size, 0, self.num_specialists, device=device, dtype=dtype),
+                torch.zeros(batch_size, 0, device=device, dtype=torch.bool),
+            )
+        batch_size = len(region_embeddings_list)
+        padded_regions = torch.zeros(batch_size, max_regions, self.retrieval_dim, device=device, dtype=dtype)
+        padded_logits = torch.zeros(batch_size, max_regions, self.num_specialists, device=device, dtype=dtype)
+        padded_mask = torch.zeros(batch_size, max_regions, device=device, dtype=torch.bool)
+        for idx, (regions, logits, mask) in enumerate(zip(region_embeddings_list, routing_logits_list, routing_mask_list)):
+            if regions.numel() == 0:
+                continue
+            count = regions.size(0)
+            padded_regions[idx, :count] = regions.to(dtype)
+            padded_logits[idx, : logits.size(0)] = logits.to(dtype)
+            padded_mask[idx, : mask.numel()] = mask.to(torch.bool)
+        return padded_regions, padded_logits, padded_mask
+    # ----------------------------------------------------------------- #
+    # User-facing helpers
+    # ----------------------------------------------------------------- #
+    @torch.inference_mode()
+    def encode_queries(
+        self,
+        processor,
+        queries: List[str],
+        batch_size: int = 8,
+        max_length: Optional[int] = None,
+    ) -> List[torch.Tensor]:
+        """Encode a list of query strings into multi-vector embeddings.
+        Returns one tensor per query, since queries may have different lengths.
+        Run this on-GPU for speed; the returned tensors are moved to CPU for
+        the caller to manage batching.
+        """
+        device = next(self.parameters()).device
+        out: List[torch.Tensor] = []
+        for i in range(0, len(queries), batch_size):
+            batch = processor.process_texts(queries[i : i + batch_size], max_length=max_length).to(device)
+            emb = self(**batch).embeddings.cpu()
+            out.extend(list(torch.unbind(emb)))
+        return out
+    @torch.inference_mode()
+    def encode_images(self, processor, images, batch_size: int = 2) -> List[torch.Tensor]:
+        """Encode a list of PIL images into multi-vector embeddings."""
+        device = next(self.parameters()).device
+        out: List[torch.Tensor] = []
+        for i in range(0, len(images), batch_size):
+            batch = processor.process_images(images[i : i + batch_size]).to(device)
+            emb = self(**batch).embeddings.cpu()
+            out.extend(list(torch.unbind(emb)))
+        return out
+    @staticmethod
+    def score(
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        batch_size: int = 32,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """MaxSim scoring: for each (q_i, p_j) pair, compute
+        ``sum_t max_p <q_i_t, p_j_p>``. Returns a [N_q, N_p] matrix.
+        This reproduces ``processor.score_multi_vector`` but lives on the
+        model so users can compute relevance without touching the processor.
+        """
+        dev = torch.device(device) if device is not None else torch.device("cpu")
+        n_q, n_p = len(qs), len(ps)
+        scores = torch.zeros(n_q, n_p, device=dev)
+        for qi in range(0, n_q, batch_size):
+            q_slice = qs[qi : qi + batch_size]
+            q_len = max(x.size(0) for x in q_slice)
+            q_pad = torch.zeros(len(q_slice), q_len, q_slice[0].size(-1), device=dev)
+            q_mask = torch.zeros(len(q_slice), q_len, device=dev, dtype=torch.bool)
+            for i, t in enumerate(q_slice):
+                q_pad[i, : t.size(0)] = t.to(dev)
+                q_mask[i, : t.size(0)] = t.abs().sum(dim=-1) > 0
+            for pi in range(0, n_p, batch_size):
+                p_slice = ps[pi : pi + batch_size]
+                p_len = max(x.size(0) for x in p_slice)
+                p_pad = torch.zeros(len(p_slice), p_len, p_slice[0].size(-1), device=dev)
+                for j, t in enumerate(p_slice):
+                    p_pad[j, : t.size(0)] = t.to(dev)
+                sim = torch.einsum("qld,pkd->qplk", q_pad, p_pad)
+                maxsim = sim.max(dim=-1).values
+                maxsim = (maxsim * q_mask.unsqueeze(1).to(maxsim.dtype)).sum(dim=-1)
+                scores[qi : qi + len(q_slice), pi : pi + len(p_slice)] = maxsim
+        return scores
+__all__ = ["ArgusForRetrieval", "ArgusOutput"]

processing_argus.py ADDED Viewed

	@@ -0,0 +1,260 @@

+"""Argus: Region-Aware Query-Conditioned Mixture of Experts for Visual Document Retrieval.
+Self-contained processor for Argus-Colqwen3.5-9B. Wraps the Qwen3-VL processor
+(image processor + Qwen2 tokenizer + optional video processor) and adds ColPali-
+style ``process_images`` / ``process_texts`` / ``score_multi_vector`` helpers.
+"""
+from __future__ import annotations
+from pathlib import Path
+from typing import ClassVar, List, Optional, Tuple, Union
+import torch
+from PIL import Image
+from transformers import BatchEncoding, BatchFeature
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+from transformers.models.qwen3_vl import Qwen3VLProcessor
+class ArgusProcessor(Qwen3VLProcessor):
+    """Processor for Argus-Colqwen3.5-9B.
+    Subclasses ``Qwen3VLProcessor`` (the Qwen3.5-9B hub repo ships that
+    processor class even though the LLM is Qwen3.5). Adds:
+    - ``process_images``: batch-encode PIL images into the exact dict the
+      retriever forward expects (``pixel_values``, ``image_grid_thw``,
+      ``input_ids``, ``attention_mask``).
+    - ``process_texts``: batch-encode query strings.
+    - ``score`` / ``score_multi_vector``: MaxSim scoring helper.
+    - ``max_num_visual_tokens`` knob: caps the longest-edge pixel budget per
+      image so long documents don't blow up the vision encoder.
+    """
+    visual_prompt_prefix: ClassVar[str] = (
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>"
+    )
+    query_augmentation_token: ClassVar[str] = "<|endoftext|>"
+    query_prefix: ClassVar[str] = ""
+    image_token: ClassVar[str] = "<|image_pad|>"
+    # Number of <|endoftext|> tokens appended to every query — matches the
+    # training-time collator (``colpali_novel/data/layout_collator.py``).
+    # Removing or changing this number measurably hurts retrieval scores.
+    n_query_augmentation_tokens: ClassVar[int] = 10
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        video_processor=None,
+        chat_template=None,
+        **kwargs,
+    ):
+        # Explicit signature matters for ``ProcessorMixin``: it inspects
+        # __init__.__code__ to decide which modality attributes to set. A
+        # *args,**kwargs signature silently drops tokenizer/image_processor.
+        super().__init__(
+            image_processor=image_processor,
+            tokenizer=tokenizer,
+            video_processor=video_processor,
+            chat_template=chat_template,
+            **kwargs,
+        )
+        if getattr(self, "tokenizer", None) is not None:
+            self.tokenizer.padding_side = "left"
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        device_map: Optional[str] = None,
+        max_num_visual_tokens: Optional[int] = None,
+        **kwargs,
+    ):
+        """Load the processor from a local folder or HF repo id.
+        The Qwen3.5-9B hub repo declares ``processor_class=Qwen3VLProcessor``
+        but ``tokenizer_class=Qwen2Tokenizer``. The stock ``Qwen3VLProcessor
+        .from_pretrained`` returns ``tokenizer=None`` in that case and then
+        crashes on ``tokenizer.convert_tokens_to_ids(self.image_token)``.
+        We load tokenizer + image processor via the Auto* registry
+        explicitly so both are real objects before ``__init__`` runs.
+        """
+        from transformers import AutoImageProcessor, AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        image_processor = AutoImageProcessor.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        video_processor = None
+        try:
+            from transformers import AutoVideoProcessor
+            video_processor = AutoVideoProcessor.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        except Exception:  # noqa: BLE001 — video processing is optional
+            video_processor = None
+        chat_template = None
+        try:
+            candidate = Path(str(pretrained_model_name_or_path)) / "chat_template.jinja"
+            if candidate.is_file():
+                chat_template = candidate.read_text()
+        except Exception:  # noqa: BLE001
+            chat_template = None
+        instance = cls(
+            image_processor=image_processor,
+            tokenizer=tokenizer,
+            video_processor=video_processor,
+            chat_template=chat_template,
+        )
+        if max_num_visual_tokens is not None:
+            patch_size = getattr(instance.image_processor, "patch_size", None)
+            merge_size = getattr(instance.image_processor, "merge_size", None)
+            if patch_size is None or merge_size is None:
+                raise ValueError("Argus image processor missing patch_size or merge_size.")
+            tile = patch_size * merge_size
+            instance.image_processor.max_pixels = max_num_visual_tokens * tile * tile
+            instance.image_processor.size["longest_edge"] = instance.image_processor.max_pixels
+        return instance
+    # ------------------------------------------------------------------ #
+    # Encoding
+    # ------------------------------------------------------------------ #
+    def process_images(self, images: List[Image.Image]) -> Union[BatchFeature, BatchEncoding]:
+        """Encode PIL images into the backbone's expected input dict."""
+        images = [img.convert("RGB") for img in images]
+        batch_doc = self(
+            text=[self.visual_prompt_prefix] * len(images),
+            images=images,
+            padding="longest",
+            return_tensors="pt",
+        )
+        # Pack pixel_values so the forward can scatter them per image via
+        # image_grid_thw offsets. This mirrors the training-time collator.
+        offsets = batch_doc["image_grid_thw"][:, 1] * batch_doc["image_grid_thw"][:, 2]
+        pixel_values = list(torch.split(batch_doc["pixel_values"], offsets.tolist()))
+        batch_doc["pixel_values"] = torch.nn.utils.rnn.pad_sequence(pixel_values, batch_first=True)
+        return batch_doc
+    def process_texts(
+        self,
+        texts: List[str],
+        max_length: Optional[int] = None,
+    ) -> Union[BatchFeature, BatchEncoding]:
+        """Encode query strings into the backbone's expected input dict."""
+        kwargs = {"text": texts, "return_tensors": "pt", "padding": "longest"}
+        if max_length is not None:
+            kwargs["max_length"] = max_length
+            kwargs["truncation"] = True
+        return self(**kwargs)
+    def process_queries(
+        self,
+        queries: Optional[List[str]] = None,
+        texts: Optional[List[str]] = None,
+        max_length: Optional[int] = None,
+        suffix: Optional[str] = None,
+    ) -> Union[BatchFeature, BatchEncoding]:
+        """Encode queries with the training-time augmentation:
+        ``query_prefix + query + query_augmentation_token * n_query_augmentation_tokens``.
+        Mirrors ``colpali_engine.utils.processing_utils.BaseVisualRetrieverProcessor
+        .process_queries`` and the Argus training collator. The default 10 trailing
+        ``<|endoftext|>`` tokens are not optional — without them, MaxSim scoring
+        drops several nDCG points because the query has fewer active multi-vectors.
+        """
+        if texts is not None and queries is not None:
+            raise ValueError("Only one of 'texts' or 'queries' should be provided.")
+        if queries is None:
+            queries = texts
+        if queries is None:
+            raise ValueError("No queries provided.")
+        if suffix is None:
+            suffix = self.query_augmentation_token * self.n_query_augmentation_tokens
+        wrapped = [self.query_prefix + q + suffix for q in queries]
+        return self.process_texts(wrapped, max_length=max_length)
+    # ------------------------------------------------------------------ #
+    # Scoring
+    # ------------------------------------------------------------------ #
+    def score(
+        self,
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Alias for ``score_multi_vector`` (MaxSim over multi-vectors)."""
+        return self.score_multi_vector(qs, ps, device=device, **kwargs)
+    def score_multi_vector(
+        self,
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        batch_size: int = 128,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """Compute an [N_q, N_p] score matrix via MaxSim (ColBERT scoring).
+        For each (q, p) pair: ``sum_t max_p <q_t, p_p>``. Inputs are the raw
+        (potentially ragged) per-sample multi-vector tensors returned by
+        :meth:`encode_queries` / :meth:`encode_images`.
+        """
+        dev = torch.device(device) if device is not None else torch.device("cpu")
+        n_q, n_p = len(qs), len(ps)
+        scores = torch.zeros(n_q, n_p, device=dev)
+        for qi in range(0, n_q, batch_size):
+            q_slice = qs[qi : qi + batch_size]
+            q_len = max(x.size(0) for x in q_slice)
+            q_pad = torch.zeros(len(q_slice), q_len, q_slice[0].size(-1), device=dev)
+            q_mask = torch.zeros(len(q_slice), q_len, device=dev, dtype=torch.bool)
+            for i, t in enumerate(q_slice):
+                q_pad[i, : t.size(0)] = t.to(dev)
+                q_mask[i, : t.size(0)] = t.abs().sum(dim=-1) > 0
+            for pi in range(0, n_p, batch_size):
+                p_slice = ps[pi : pi + batch_size]
+                p_len = max(x.size(0) for x in p_slice)
+                p_pad = torch.zeros(len(p_slice), p_len, p_slice[0].size(-1), device=dev)
+                for j, t in enumerate(p_slice):
+                    p_pad[j, : t.size(0)] = t.to(dev)
+                sim = torch.einsum("qld,pkd->qplk", q_pad, p_pad)
+                maxsim = sim.max(dim=-1).values
+                maxsim = (maxsim * q_mask.unsqueeze(1).to(maxsim.dtype)).sum(dim=-1)
+                scores[qi : qi + len(q_slice), pi : pi + len(p_slice)] = maxsim
+        return scores
+    # ------------------------------------------------------------------ #
+    # Misc helpers (match colpali-engine BaseVisualRetrieverProcessor API)
+    # ------------------------------------------------------------------ #
+    def get_n_patches(
+        self,
+        image_size: Tuple[int, int],
+        spatial_merge_size: int,
+    ) -> Tuple[int, int]:
+        patch_size = self.image_processor.patch_size
+        height_new, width_new = smart_resize(
+            width=image_size[0],
+            height=image_size[1],
+            factor=patch_size * self.image_processor.merge_size,
+            min_pixels=self.image_processor.size["shortest_edge"],
+            max_pixels=self.image_processor.size["longest_edge"],
+        )
+        n_patches_x = width_new // patch_size // spatial_merge_size
+        n_patches_y = height_new // patch_size // spatial_merge_size
+        return n_patches_x, n_patches_y
+    def get_image_mask(self, batch_images: BatchFeature) -> torch.Tensor:
+        return batch_images.input_ids == self.image_token_id
+__all__ = ["ArgusProcessor"]

processor_config.json ADDED Viewed

	@@ -0,0 +1,64 @@

+{
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_processor_type": "Qwen2VLImageProcessor",
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "max_pixels": 2097152,
+    "merge_size": 2,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "longest_edge": 2097152,
+      "shortest_edge": 65536
+    },
+    "temporal_patch_size": 2
+  },
+  "processor_class": "ArgusProcessor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "fps": 2,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "max_frames": 768,
+    "merge_size": 2,
+    "min_frames": 4,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "size": {
+      "longest_edge": 25165824,
+      "shortest_edge": 4096
+    },
+    "temporal_patch_size": 2,
+    "video_processor_type": "Qwen3VLVideoProcessor"
+  },
+  "auto_map": {
+    "AutoProcessor": "processing_argus.ArgusProcessor"
+  }
+}

tokenizer.json ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523
+size 19989325

tokenizer_config.json ADDED Viewed

	@@ -0,0 +1,33 @@

+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": true,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "processor_class": "ArgusProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}