Spaces:
Sleeping
Sleeping
| """Tokenize free-text annotation columns (e.g. `annotation.substask`) into | |
| fixed-length token tensors keyed by `annotation_tokens__<field>` / | |
| `annotation_mask__<field>`. | |
| The model-side CE loss branch consumes these keys; see | |
| `modeling_labvla._compute_annotation_ce`. | |
| Placed late in the transform chain but BEFORE the final | |
| `UnifyLabVLAInputsTransformFn` so that: | |
| - the adapter has already put the raw annotation string into the sample dict | |
| - the downstream Unify step can pass `annotation_tokens__*` through to the | |
| collated batch | |
| This transform is a no-op when `annotation_specs` is empty, so datasets | |
| without annotation_losses incur zero cost. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from dataclasses import dataclass, field, replace | |
| from typing import Any | |
| import torch | |
| from src.schema.annotation_loss import AnnotationLossSpec | |
| from src.transforms.core import DataDict, DataTransformFn | |
| class AnnotationTokenizeTransformFn(DataTransformFn): | |
| """Tokenize each configured annotation field with the VLM tokenizer. | |
| Args: | |
| tokenizer_path: | |
| HF model id or local path used to instantiate a fresh tokenizer. | |
| Should match the VLM tokenizer used by | |
| `Qwen3_VLProcessorTransformFn` to ensure token-id consistency. | |
| annotation_specs: | |
| Tuple of AnnotationLossSpec; one tokenization pass per entry. | |
| """ | |
| tokenizer_path: str = "Qwen/Qwen3-VL-4B-Instruct" | |
| # Populated by hydrate_all from schema.annotation_losses. Default empty | |
| # tuple → this transform is a cheap passthrough (no-op). | |
| annotation_specs: tuple[AnnotationLossSpec, ...] = () | |
| # When true, emit tensors at the effective per-sample cap instead of the | |
| # schema union cap. The LabVLA collator pads variable 1D token tensors back | |
| # to the batch-local max, preserving collatability. | |
| dynamic_shape: bool = False | |
| # Instantiated lazily in __post_init__; not a dataclass init arg. | |
| tokenizer: Any = field(default=None, init=False, repr=False) | |
| def __post_init__(self) -> None: | |
| if not self.annotation_specs: | |
| # No-op mode: skip tokenizer load entirely. Keeps OXE-only runs | |
| # (no annotations) free of unused HF downloads / CPU cost. | |
| self.tokenizer = None | |
| return | |
| from transformers import AutoTokenizer | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path) | |
| def __call__(self, data: DataDict) -> DataDict: | |
| if not self.annotation_specs or self.tokenizer is None: | |
| return data | |
| for spec in self.annotation_specs: | |
| raw = data.get(spec.field, "") | |
| if not isinstance(raw, str): | |
| # Some annotations are stored as bytes or json strings; | |
| # coerce to plain text. Non-coercible → empty string (fully | |
| # masked below). | |
| try: | |
| raw = str(raw) if raw is not None else "" | |
| except Exception: | |
| raw = "" | |
| # Per-sample max_length override. Adapters that | |
| # know a tighter cap for the current sample (e.g. RoboInter VQA's | |
| # task-family budgets: understanding=8, planning=64) can attach | |
| # `f"{spec.field}_max_length"` to the sample dict. We always clamp | |
| # against the schema-level `spec.max_length` (the schema value is | |
| # the union upper bound; per-sample can only TIGHTEN it). | |
| # | |
| # Tensor SHAPE is preserved at `(spec.max_length,)` so the default | |
| # collator can stack heterogeneous task-family samples into one | |
| # batch. The cap manifests as: real tokens are truncated to | |
| # `effective_max_length`; mask positions in | |
| # `[effective_max_length, spec.max_length)` are forced to 0; ID | |
| # positions in that range are zero-filled (so even if the LM ever | |
| # ran on them, CE would be ignored via mask). | |
| per_sample_key = f"{spec.field}_max_length" | |
| per_sample_cap = data.get(per_sample_key, None) | |
| if isinstance(per_sample_cap, int) and per_sample_cap > 0: | |
| effective_max_length = min(per_sample_cap, spec.max_length) | |
| else: | |
| effective_max_length = spec.max_length | |
| tensor_max_length = effective_max_length if self.dynamic_shape else spec.max_length | |
| enc = self.tokenizer( | |
| raw if raw else "", | |
| max_length=tensor_max_length, | |
| padding="max_length", | |
| truncation=True, | |
| return_tensors="pt", # direct tensor, skip list round-trip | |
| add_special_tokens=False, # never pollute annotation with BOS/EOS | |
| ) | |
| # return_tensors="pt" yields shape (1, L); squeeze to (L,). | |
| ids = enc["input_ids"][0].long() | |
| mask = enc["attention_mask"][0].bool() | |
| # Samples with empty annotation string still get non-zero length | |
| # (BOS/EOS depending on tokenizer) but attention_mask zeros out | |
| # every position — guarantees zero CE contribution. | |
| if not raw: | |
| mask = torch.zeros_like(mask) | |
| # Apply per-sample cap by zeroing IDs and mask beyond | |
| # `effective_max_length`. Done in-place on the freshly-allocated | |
| # squeezed tensors (no aliasing concern). When per-sample cap == | |
| # spec.max_length (the default no-override path), this is a no-op. | |
| if not self.dynamic_shape and effective_max_length < spec.max_length: | |
| ids[effective_max_length:] = 0 | |
| mask[effective_max_length:] = False | |
| data[spec.batch_key_tokens()] = ids | |
| data[spec.batch_key_mask()] = mask | |
| # Per-field weight carried inside the batch so the model can | |
| # compose the loss without a cross-cutting schema reference. All | |
| # samples from the same dataset carry the same scalar — this is | |
| # a per-sample-shape (not per-value) carrier for a schema-level | |
| # constant. | |
| data[spec.batch_key_weight()] = torch.tensor( | |
| float(spec.weight), dtype=torch.float32 | |
| ) | |
| return data | |
| def hydrate(self, ctx) -> "AnnotationTokenizeTransformFn": | |
| # Inject this dataset's annotation specs. Empty tuple → no-op path, | |
| # zero tokenizer load, zero overhead for OXE-style datasets. | |
| t = replace(self, annotation_specs=tuple(ctx.schema.annotation_losses)) | |
| logging.info( | |
| f"Hydrated {t.__class__.__name__} with " | |
| f"{len(t.annotation_specs)} annotation_specs " | |
| f"dynamic_shape={getattr(t, 'dynamic_shape', False)} " | |
| f"({ctx.schema.schema_id})" | |
| ) | |
| return t | |