Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / omni / metadata.py: 83%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Metadata for whole, already-transformed Omni image-text conversations.""" 

16 

17from collections.abc import Mapping 

18from typing import Any 

19 

20from hyper_parallel.distributed_data.schema import SampleMetadata, WorkloadCost 

21 

22 

23def vlm_sample_metadata(sample: Mapping[str, Any]) -> SampleMetadata: 

24 """Estimate visual work without changing native VLM sample fields. 

25 

26 Args: 

27 sample: One CPU Omni transform output before collation. 

28 

29 Returns: 

30 Physical padded sequence width for capacity and raw image patch count 

31 for encoder cost. LLM cost uses the padded width: the native dense 

32 text computation does not disappear at masked padding positions. 

33 

34 Raises: 

35 ValueError: If sequence or image tensors do not follow the native schema. 

36 

37 Note: 

38 These are workload proxies, not calibrated time estimates. Image patches 

39 come from the processor's actual grid, not a hard-coded patch/merge size. 

40 Decode/transform has already run on the Reader and is not balanced here. 

41 """ 

42 input_ids = sample.get("input_ids") 

43 grid = sample.get("image_grid_thw") 

44 pixels = sample.get("pixel_values") 

45 if input_ids is None or input_ids.ndim != 1 or input_ids.shape[0] == 0: 

46 raise ValueError("VLM metadata requires one non-empty 1-D input_ids tensor before collation") 

47 if grid is None or grid.ndim != 2 or grid.shape[1] != 3: 

48 raise ValueError("VLM metadata requires image_grid_thw with shape [num_images, 3]") 

49 if grid.device.type != "cpu": 

50 raise ValueError("VLM metadata must be extracted on CPU before device transfer") 

51 image_patches = 0 

52 for dimensions in grid.tolist(): 

53 if any(not isinstance(dimension, int) or isinstance(dimension, bool) for dimension in dimensions): 

54 raise ValueError("VLM image_grid_thw must contain integer patch dimensions") 

55 if any(dimension <= 0 for dimension in dimensions): 

56 raise ValueError("VLM image_grid_thw dimensions must be positive") 

57 image_patches += dimensions[0] * dimensions[1] * dimensions[2] 

58 if pixels is None or pixels.ndim != 2 or pixels.shape[0] != image_patches: 

59 raise ValueError("VLM pixel_values rows must equal the raw patch count from image_grid_thw") 

60 seq_len = int(input_ids.shape[0]) 

61 return SampleMetadata( 

62 pack_tokens=seq_len, 

63 features={"P": seq_len, "D": 0, "image_patches": image_patches}, 

64 cost=WorkloadCost(encoder=float(image_patches), llm=float(seq_len)), 

65 )