"""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", ]