> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.qwen3_reranker.collator

Context-aware collator for the Qwen3 reranker.

Owned by this model package rather than the generic dataset collators because every
constant it carries is part of the Qwen3-Reranker prompt contract: the ChatML markers,
the empty `<think>` block that makes the final tokens a yes/no next-token prediction,
the model card's system message, and the yes/no label semantics themselves.

It stands alone on `DataCollatorWithPadding` and imports neither
`model.py` nor the generic retrieval collator module, so training data assembly and
model definition stay independently importable.

## Module Contents

### Classes

| Name                                                                                                                               | Description                                                          |
| ---------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`Qwen3ContextAwareRerankerCollator`](#nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator) | Collate query-document pairs using the Qwen3 reranker chat template. |

### API

```python
class nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator(
    rerank_max_length: int,
    args: typing.Any = (),
    instructions: dict[str | tuple[str, ...] | frozenset[str], str] | None = None,
    system_message: str | None = None,
    reasoning_drop_prob: float = 0.5,
    global_query_drop_prob: float = 0.5,
    drop_seed: int = 42,
    global_query_max_length: int | None = None,
    sub_query_max_length: int | None = None,
    reasoning_max_length: int | None = None,
    passage_max_length: int | None = None,
    prefix_template: str | None = None,
    suffix_template: str | None = None,
    kwargs: typing.Any = {}
)
```

\<Indent>

**Bases:** `DataCollatorWithPadding`

Collate query-document pairs using the Qwen3 reranker chat template.

Wraps each (query, document) pair in the instruction-aware chat format used by
`Qwen/Qwen3-Reranker-*` so that the final tokens are the assistant think-prefix
and the next-token prediction is "yes"/"no". The middle (instruct/query/document)
is truncated to fit, then the fixed prefix/suffix token ids are concatenated so the
template markers are never truncated away.

Optional context fields `reasoning` and `global_query` are supported. When
present they are embedded inside the `&lt;Query&gt;:` block as `#`-prefixed markers.
The instruction is selected from `instructions` by which fields SURVIVE the drop,
so it always describes the context the prompt actually carries.

Templates per mode::

# base -- byte-identical to the out-of-the-box Qwen3-Reranker prompt

\<Instruct>: \{instruction}
\<Query>: \{query}
\<Document>: \{document}

# reasoning

\<Instruct>: \{instruction}
\<Query>: #Reasoning Trace: \{reasoning}
\#Query: \{query}
\<Document>: \{document}

# globalq

\<Instruct>: \{instruction}
\<Query>: #Original Question (global query): \{global\_query}
\#Query: \{query}
\<Document>: \{document}

# full

\<Instruct>: \{instruction}
\<Query>: #Original Question (global query): \{global\_query}
\#Reasoning Trace: \{reasoning}
\#Query: \{query}
\<Document>: \{document}

Context dropout (multimode training)

````
``reasoning_drop_prob`` and ``global_query_drop_prob`` (both 0.5 by default) drop each
field independently per query, so a row carrying both yields all four modes above. A
drop REMOVES the field -- no placeholder is substituted, so the prompt shape changes
with it and the instruction follows. The draw is a hash of
(drop_seed, epoch, field, query): identical across runs, workers and ranks, and
redrawn per epoch once ``set_epoch`` is called. Set both to 0.0 at
eval time to force full mode.


<ParamField path="DEFAULT_INSTRUCTIONS" type="dict[frozenset[str], str]">
</ParamField>

<ParamField path="DEFAULT_PREFIX_TEMPLATE">
</ParamField>

<ParamField path="DEFAULT_SUFFIX_TEMPLATE">
</ParamField>

<ParamField path="DEFAULT_SYSTEM">
</ParamField>

<ParamField path="_epoch">
</ParamField>

<ParamField path="instructions" type="= self._normalize_instructions(raw)">
</ParamField>

<ParamField path="prefix_ids">
</ParamField>

<ParamField path="suffix_ids">
</ParamField>
<Anchor id="nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator-__call__">

<CodeBlock showLineNumbers={false} wordWrap={true}>

```python
nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator.__call__(
    features: list[dict[str, typing.Any]]
) -> transformers.BatchEncoding
```

</CodeBlock>
</Anchor>

<Indent>

Tokenize flattened query-document rows into a padded batch.

**Parameters:**

<ParamField path="features" type="list[dict[str, Any]]">
One mapping per pair with question and doc_text strings, optional
reasoning/global_query strings, and num_labels giving the number of queries.
</ParamField>

**Returns:** `BatchEncoding`

BatchEncoding with input_ids and attention_mask of shape</Indent>
<Anchor id="nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator-_format_one">

<CodeBlock showLineNumbers={false} wordWrap={true}>

```python
nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator._format_one(
    query: str,
    doc: str,
    reasoning: str | None = None,
    global_query: str | None = None
) -> str
```

</CodeBlock>
</Anchor>

<Indent>

Build the user-turn text for a single (query, doc) pair.

Order: decide which context fields survive the drop draws, cap each surviving
item on its own token budget, pick the instruction matching what survived, then
assemble with the context sub-fields embedded inside ``&lt;Query&gt;:``.

Capping precedes assembly so each item is guaranteed its share. The document is
normally left uncapped and takes the remainder of ``rerank_max_length``, absorbing
all overflow on its own because it is assembled last.</Indent>
<Anchor id="nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator-_keep_field">

<CodeBlock showLineNumbers={false} wordWrap={true}>

```python
nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator._keep_field(
    kind: str,
    query: str,
    prob: float
) -> bool
```

</CodeBlock>
</Anchor>

<Indent>

Whether to KEEP a context field for this query, deterministically.

Keyed on the QUERY, not the (query, document) pair: the dataset repeats the
context fields across every passage of a group, and a listwise group is scored as
one softmax over 1 positive + n negatives. If passages within a group disagreed
about which context was present, the comparison would be between different
prompts rather than between documents.

Hashed rather than sampled from a shared RNG so the draw depends only on
(drop_seed, epoch, field, query) -- identical across runs, workers and ranks, and
independent of batch order. hashlib rather than hash(): PYTHONHASHSEED randomises
str hashing per process, which would make runs unreproducible.</Indent>
<Anchor id="nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator-_normalize_instructions">

<CodeBlock showLineNumbers={false} wordWrap={true}>

```python
nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator._normalize_instructions(
    raw: dict[str | tuple[str, ...] | frozenset[str], str]
) -> dict[frozenset[str], str]
```

</CodeBlock>
</Anchor>

<Indent>

<Badge>staticmethod</Badge>

Normalize instruction keys to frozensets of field name strings.

Accepts three key formats so callers can use whichever is most natural:
- ``frozenset`` — used directly (Python API).
- ``tuple`` of strings — converted to frozenset (Python API).
- ``str`` — comma-separated field names (YAML-friendly); empty string maps to
  the no-context frozenset. Examples::

      ""                    → frozenset()
      "reasoning"           → frozenset(&#123;"reasoning"&#125;)
      "global_query,reasoning" → frozenset(&#123;"global_query","reasoning"&#125;)</Indent>
<Anchor id="nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator-_truncate_tokens">

<CodeBlock showLineNumbers={false} wordWrap={true}>

```python
nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator._truncate_tokens(
    text: str,
    limit: int | None = None
) -> str
```

</CodeBlock>
</Anchor>

<Indent>

Cut ``text`` to at most ``limit`` tokens. No-op when limit is None.</Indent>
<Anchor id="nemo_automodel-components-models-qwen3_reranker-collator-Qwen3ContextAwareRerankerCollator-set_epoch">

<CodeBlock showLineNumbers={false} wordWrap={true}>

```python
nemo_automodel.components.models.qwen3_reranker.collator.Qwen3ContextAwareRerankerCollator.set_epoch(
    epoch: int
) -> None
```

</CodeBlock>
</Anchor>

<Indent>

Set the epoch used for deterministic context dropout.

Called by ``StepScheduler.set_epoch`` on the training dataloader's collate
function, so drops are redrawn each epoch. A validation collator is never given an
epoch and stays at 0, keeping its prompt mix fixed -- otherwise a val_loss that
moved because the sampled modes changed would be indistinguishable from one that
moved because the model did.

**Parameters:**

<ParamField path="epoch" type="int">
Zero-based index of the epoch about to run.
</ParamField></Indent></Indent>

<style>{`
  .light .fern-code-block,
  .light .fern-prose code:not(.code-block) {
    background-color: var(--nv-color-bg-alt, #f7f7f7) !important;
  }

  .dark .fern-code-block,
  .dark .fern-prose code:not(.code-block) {
    background-color: #1f1f1f !important;
  }
`}</style>
````