| from __future__ import annotations |
|
|
| import torch |
| from .prompting import PromptParts, tokenize_without_special_tokens |
|
|
|
|
| def build_inference_batch(tokenizer, image_processor, image, instruction, config, device): |
| prompt = PromptParts() |
| header = tokenize_without_special_tokens(tokenizer, prompt.header) |
| instruction_ids = tokenize_without_special_tokens( |
| tokenizer, prompt.instruction_prefix + instruction + prompt.instruction_suffix |
| ) |
| placeholder = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id |
| processed = image_processor(image, return_tensors="pt") |
| grid = processed["grid_sizes"][0].long() |
| merge = int(processed["merge_sizes"][0]) |
| if int(grid[0]) != 1 or int(grid[1]) % merge or int(grid[2]) % merge: |
| raise ValueError(f"Unsupported inference grid: {grid.tolist()}, merge={merge}") |
| grid_height = int(grid[1]) // merge |
| grid_width = int(grid[2]) // merge |
| visual_tokens = grid_height * grid_width |
| input_ids = torch.tensor( |
| [header + [placeholder] * visual_tokens + instruction_ids], |
| dtype=torch.long, |
| device=device, |
| ) |
| return { |
| "input_ids": input_ids, |
| "header_length": len(header), |
| "instruction_length": len(instruction_ids), |
| "visual_token_counts": torch.tensor([visual_tokens], device=device), |
| "grid_heights": torch.tensor([grid_height], device=device), |
| "grid_widths": torch.tensor([grid_width], device=device), |
| "pixel_values": processed["pixel_values"].to(device), |
| "grid_sizes": processed["grid_sizes"].to(device), |
| "merge_sizes": processed["merge_sizes"].to(device), |
| } |
|
|
