236 lines
8.5 KiB
Python
236 lines
8.5 KiB
Python
|
"""Multimodal text preprocessing helpers.
|
||
|
|
|
||
|
|
This file is generated by ``scripts/export_opensource.py``. It keeps only the
|
||
|
|
processor wrapper needed by harrix inference and uses a generic system prompt.
|
||
|
|
Internal robot-id, dataset, camera, and frequency maps are not bundled here.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from transformers import BatchFeature
|
||
|
|
from transformers.tokenization_utils_base import BatchEncoding
|
||
|
|
|
||
|
|
|
||
|
|
def pad_text_input_to_target_length(
|
||
|
|
text_inputs, target_length, pad_token_id=151643, padding_side="right"
|
||
|
|
):
|
||
|
|
"""Pad or truncate tokenized text to ``target_length``."""
|
||
|
|
batch_size, current_length = text_inputs.input_ids.shape
|
||
|
|
if current_length < target_length:
|
||
|
|
padding_size = target_length - current_length
|
||
|
|
padding = torch.full(
|
||
|
|
(batch_size, padding_size),
|
||
|
|
pad_token_id,
|
||
|
|
dtype=text_inputs.input_ids.dtype,
|
||
|
|
device=text_inputs.input_ids.device,
|
||
|
|
)
|
||
|
|
attention_padding = torch.zeros(
|
||
|
|
(batch_size, padding_size),
|
||
|
|
dtype=text_inputs.attention_mask.dtype,
|
||
|
|
device=text_inputs.attention_mask.device,
|
||
|
|
)
|
||
|
|
if padding_side == "right":
|
||
|
|
text_inputs["input_ids"] = torch.cat([text_inputs.input_ids, padding], dim=1)
|
||
|
|
text_inputs["attention_mask"] = torch.cat(
|
||
|
|
[text_inputs.attention_mask, attention_padding], dim=1
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
text_inputs["input_ids"] = torch.cat([padding, text_inputs.input_ids], dim=1)
|
||
|
|
text_inputs["attention_mask"] = torch.cat(
|
||
|
|
[attention_padding, text_inputs.attention_mask], dim=1
|
||
|
|
)
|
||
|
|
elif current_length > target_length:
|
||
|
|
if padding_side == "right":
|
||
|
|
text_inputs["input_ids"] = text_inputs.input_ids[:, :target_length]
|
||
|
|
text_inputs["attention_mask"] = text_inputs.attention_mask[:, :target_length]
|
||
|
|
else:
|
||
|
|
text_inputs["input_ids"] = text_inputs.input_ids[:, -target_length:]
|
||
|
|
text_inputs["attention_mask"] = text_inputs.attention_mask[:, -target_length:]
|
||
|
|
return text_inputs
|
||
|
|
|
||
|
|
|
||
|
|
def _replace_media_placeholders(text, grid_thw, token, merge_length):
|
||
|
|
if grid_thw is None:
|
||
|
|
return text
|
||
|
|
index = 0
|
||
|
|
for i in range(len(text)):
|
||
|
|
while token in text[i]:
|
||
|
|
if index >= len(grid_thw):
|
||
|
|
raise ValueError(
|
||
|
|
f"More {token} placeholders than media tensors in sample {i}"
|
||
|
|
)
|
||
|
|
token_count = int(grid_thw[index].prod() // merge_length)
|
||
|
|
text[i] = text[i].replace(token, "<|placeholder|>" * token_count, 1)
|
||
|
|
index += 1
|
||
|
|
text[i] = text[i].replace("<|placeholder|>", token)
|
||
|
|
return text
|
||
|
|
|
||
|
|
|
||
|
|
_PUBLIC_CAMERA_LABELS = {
|
||
|
|
"face_view": "front view",
|
||
|
|
"right_wrist_view": "right wrist view",
|
||
|
|
"left_wrist_view": "left wrist view",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _camera_label(cam_name):
|
||
|
|
return _PUBLIC_CAMERA_LABELS.get(str(cam_name), str(cam_name).replace("_", " "))
|
||
|
|
|
||
|
|
|
||
|
|
def preprocesser_call(
|
||
|
|
processor,
|
||
|
|
norm_state=None,
|
||
|
|
agent_pos_mask=None,
|
||
|
|
images=None,
|
||
|
|
prefix_text=None,
|
||
|
|
postfix_text=None,
|
||
|
|
videos=None,
|
||
|
|
padding=False,
|
||
|
|
padding_side="left",
|
||
|
|
truncation=None,
|
||
|
|
max_length=None,
|
||
|
|
return_tensors="pt",
|
||
|
|
pad_prefix_to_same_length=False,
|
||
|
|
pad_to_128_multiple=True,
|
||
|
|
state_augmentation_prob=0.0,
|
||
|
|
state_augmentation_ratio=0.0,
|
||
|
|
state_bins=256,
|
||
|
|
inference_mode=False,
|
||
|
|
**_,
|
||
|
|
):
|
||
|
|
"""Build a ``BatchFeature`` for Wall-X VLA inference.
|
||
|
|
|
||
|
|
This is the inference subset of the internal preprocessing helper: text,
|
||
|
|
image/video placeholder expansion, optional discretized proprioception
|
||
|
|
strings, padding, and labels=None for inference.
|
||
|
|
"""
|
||
|
|
if prefix_text is None:
|
||
|
|
raise ValueError("prefix_text is required")
|
||
|
|
if postfix_text is None:
|
||
|
|
postfix_text = [""] * len(prefix_text)
|
||
|
|
if not isinstance(prefix_text, list):
|
||
|
|
prefix_text = [prefix_text]
|
||
|
|
if not isinstance(postfix_text, list):
|
||
|
|
postfix_text = [postfix_text]
|
||
|
|
batch_size = len(prefix_text)
|
||
|
|
|
||
|
|
if images is not None and len(images) > 0:
|
||
|
|
image_inputs = processor.image_processor(images=images, return_tensors=return_tensors)
|
||
|
|
image_grid_thw = image_inputs["image_grid_thw"]
|
||
|
|
else:
|
||
|
|
image_inputs = {}
|
||
|
|
image_grid_thw = None
|
||
|
|
|
||
|
|
if videos is not None:
|
||
|
|
if hasattr(processor, "video_processor"):
|
||
|
|
videos_inputs = processor.video_processor(videos=videos, return_tensors=return_tensors)
|
||
|
|
else:
|
||
|
|
videos_inputs = processor.image_processor(
|
||
|
|
images=None, videos=videos, return_tensors=return_tensors
|
||
|
|
)
|
||
|
|
video_grid_thw = videos_inputs["video_grid_thw"]
|
||
|
|
else:
|
||
|
|
videos_inputs = {}
|
||
|
|
video_grid_thw = None
|
||
|
|
|
||
|
|
merge_length = processor.image_processor.merge_size**2
|
||
|
|
prefix_text = _replace_media_placeholders(
|
||
|
|
list(prefix_text), image_grid_thw, "<|image_pad|>", merge_length
|
||
|
|
)
|
||
|
|
prefix_text = _replace_media_placeholders(
|
||
|
|
prefix_text, video_grid_thw, "<|video_pad|>", merge_length
|
||
|
|
)
|
||
|
|
|
||
|
|
if norm_state is not None:
|
||
|
|
norm_state = norm_state.cpu().numpy() if isinstance(norm_state, torch.Tensor) else norm_state
|
||
|
|
agent_pos_mask = (
|
||
|
|
agent_pos_mask[:, 0, :].cpu().numpy().astype(bool)
|
||
|
|
if isinstance(agent_pos_mask, torch.Tensor)
|
||
|
|
else agent_pos_mask[:, 0, :].astype(bool)
|
||
|
|
)
|
||
|
|
discretized = np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1
|
||
|
|
discretized = discretized[:, 0, :]
|
||
|
|
for i in range(batch_size):
|
||
|
|
if "<|propri|>" not in prefix_text[i]:
|
||
|
|
continue
|
||
|
|
state_str = " ".join(map(str, discretized[i, agent_pos_mask[i]]))
|
||
|
|
prefix_text[i] = prefix_text[i].replace("<|propri|>", state_str)
|
||
|
|
|
||
|
|
if not pad_prefix_to_same_length:
|
||
|
|
text = [pre + post for pre, post in zip(prefix_text, postfix_text)]
|
||
|
|
text_inputs = processor.tokenizer(
|
||
|
|
text,
|
||
|
|
return_tensors=return_tensors,
|
||
|
|
padding=padding,
|
||
|
|
padding_side=padding_side,
|
||
|
|
truncation=truncation,
|
||
|
|
max_length=max_length,
|
||
|
|
)
|
||
|
|
text_inputs["prefix_length"] = None
|
||
|
|
else:
|
||
|
|
prefix_inputs = processor.tokenizer(
|
||
|
|
prefix_text,
|
||
|
|
return_tensors=return_tensors,
|
||
|
|
padding=padding,
|
||
|
|
padding_side="left",
|
||
|
|
truncation=truncation,
|
||
|
|
max_length=max_length,
|
||
|
|
)
|
||
|
|
postfix_inputs = processor.tokenizer(
|
||
|
|
postfix_text,
|
||
|
|
return_tensors=return_tensors,
|
||
|
|
padding=padding,
|
||
|
|
padding_side="right",
|
||
|
|
truncation=truncation,
|
||
|
|
max_length=max_length,
|
||
|
|
)
|
||
|
|
text_inputs = BatchEncoding(
|
||
|
|
data={
|
||
|
|
"input_ids": torch.cat([prefix_inputs.input_ids, postfix_inputs.input_ids], dim=1),
|
||
|
|
"attention_mask": torch.cat(
|
||
|
|
[prefix_inputs.attention_mask, postfix_inputs.attention_mask], dim=1
|
||
|
|
),
|
||
|
|
"prefix_length": prefix_inputs.input_ids.shape[1],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
pad_token_id = processor.tokenizer.pad_token_id
|
||
|
|
if pad_token_id is None:
|
||
|
|
pad_token_id = processor.tokenizer.eos_token_id
|
||
|
|
if pad_to_128_multiple:
|
||
|
|
target_length = 128 * ((max(len(t) for t in text_inputs.input_ids) + 127) // 128)
|
||
|
|
text_inputs = pad_text_input_to_target_length(
|
||
|
|
text_inputs, target_length, pad_token_id=pad_token_id, padding_side=padding_side
|
||
|
|
)
|
||
|
|
|
||
|
|
text_inputs["labels"] = None if inference_mode else None
|
||
|
|
return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})
|
||
|
|
|
||
|
|
|
||
|
|
def get_prologue_with_embodied_information(dataset_name, cam_mapping, robot_id, uid, config):
|
||
|
|
"""Return a generic VLA system prompt without private robot maps."""
|
||
|
|
role_start = "<|im_start|>"
|
||
|
|
role_end = "<|im_end|>"
|
||
|
|
prologue = (
|
||
|
|
f"{role_start}system\n"
|
||
|
|
"You are an embodied vision-language-action model controlling a robot "
|
||
|
|
"with language instructions."
|
||
|
|
)
|
||
|
|
if cam_mapping:
|
||
|
|
cameras = ", ".join(_camera_label(name) for name in cam_mapping.values())
|
||
|
|
prologue += f"\nCamera Setup: {cameras}"
|
||
|
|
if not getattr(config, "use_relative_action", False):
|
||
|
|
prologue += "\nAction Space: Abs EEF"
|
||
|
|
else:
|
||
|
|
prologue += "\nAction Space: Rel EEF"
|
||
|
|
return f"{prologue}\n{role_end}\n"
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"preprocesser_call",
|
||
|
|
"get_prologue_with_embodied_information",
|
||
|
|
"pad_text_input_to_target_length",
|
||
|
|
]
|