Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,235 @@
|
||||
"""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",
|
||||
]
|
||||
Reference in New Issue
Block a user