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

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"""Omni transform lifecycle and the default AutoProcessor implementation.""" 

16 

17import copy 

18import os 

19from collections.abc import Callable, Mapping 

20from dataclasses import dataclass 

21from typing import Any 

22 

23from transformers import AutoProcessor 

24 

25from hyper_parallel.data.constants import IGNORE_INDEX, OMNI_CP_TOKEN_FIELDS, ONLINE_SOURCE_PATH_KEY 

26 

27 

28def build_auto_processor( 

29 *, 

30 pretrained_model_name_or_path: str, 

31 trust_remote_code: bool = True, 

32 **processor_kwargs: Any, 

33) -> Any: 

34 """Build the Transformers processor used by the default Omni transform. 

35 

36 Args: 

37 pretrained_model_name_or_path: Model directory or Hub identifier. 

38 trust_remote_code: Whether Transformers may load model-owned code. 

39 **processor_kwargs: Additional ``AutoProcessor.from_pretrained`` options. 

40 

41 Returns: 

42 Processor with an attached tokenizer and resolved chat template. 

43 """ 

44 processor = AutoProcessor.from_pretrained( 

45 pretrained_model_name_or_path, 

46 trust_remote_code=trust_remote_code, 

47 **processor_kwargs, 

48 ) 

49 tokenizer = getattr(processor, "tokenizer", None) 

50 if tokenizer is None: 

51 raise ValueError("AutoProcessor must provide a tokenizer") 

52 chat_template = getattr(processor, "chat_template", None) 

53 if chat_template is None: 

54 chat_template = getattr(tokenizer, "chat_template", None) 

55 if chat_template is not None: 

56 processor.chat_template = chat_template 

57 

58 return processor 

59 

60 

61class OmniDataTransform: 

62 """Define model-specific sample and batch encoding hooks. 

63 

64 Implement ``encode_sample`` for eager encoding, ``preencode_sample`` and 

65 ``postencode_sample`` for deferred encoding, or only ``postencode_sample`` 

66 when source samples already contain offline packing metadata. The data 

67 pipeline detects the implemented hook combination automatically. 

68 

69 Packing selection and concatenation belong to the optional packing path 

70 rather than this transform. 

71 """ 

72 

73 def __init__(self, *, max_seq_len: int, processor: Any) -> None: 

74 """Store the sequence limit and externally built processor.""" 

75 if max_seq_len <= 0: 

76 raise ValueError("max_seq_len must be positive") 

77 if processor is None: 

78 raise ValueError("processor is required") 

79 self.max_seq_len = max_seq_len 

80 self.processor = processor 

81 

82 @staticmethod 

83 def is_valid_sample(sample: Mapping[str, Any]) -> bool: 

84 """Return whether one source record contains conversation messages.""" 

85 messages = sample.get("messages") 

86 if not isinstance(messages, list): 

87 raise ValueError("Omni sample messages must be a list") 

88 return bool(messages) 

89 

90 def encode_sample(self, sample: dict[str, Any]) -> dict[str, Any]: 

91 """Encode one source sample into one model-ready sample.""" 

92 raise NotImplementedError(f"{type(self).__name__} does not implement encode_sample") 

93 

94 def preencode_sample(self, sample: dict[str, Any]) -> dict[str, Any]: 

95 """Prepare metadata used by a packing selector without full encoding.""" 

96 raise NotImplementedError(f"{type(self).__name__} does not implement preencode_sample") 

97 

98 def postencode_sample(self, sample: dict[str, Any]) -> dict[str, Any]: 

99 """Encode selected online or offline metadata into a model-ready sample.""" 

100 raise NotImplementedError(f"{type(self).__name__} does not implement postencode_sample") 

101 

102 def encode_batch(self, batch: Any) -> Any: 

103 """Apply an optional model-specific batch encoding stage.""" 

104 return batch 

105 

106 def prepare_messages(self, sample: dict[str, Any]) -> list[dict[str, Any]]: 

107 """Copy messages and resolve media paths relative to their JSONL file.""" 

108 messages = sample.get("messages") 

109 if not isinstance(messages, list) or not messages: 

110 raise ValueError("Omni sample must contain non-empty messages") 

111 

112 prepared_messages = copy.deepcopy(messages) 

113 source_path = sample.get(ONLINE_SOURCE_PATH_KEY) 

114 if source_path is None: 

115 return prepared_messages 

116 

117 source_directory = os.path.dirname(os.path.abspath(source_path)) 

118 for message in prepared_messages: 

119 content = message.get("content") 

120 if not isinstance(content, list): 

121 continue 

122 for content_block in content: 

123 self._resolve_media_block(content_block, source_directory) 

124 return prepared_messages 

125 

126 @classmethod 

127 def _resolve_media_block(cls, content_block: Any, source_directory: str) -> None: 

128 """Resolve path-bearing fields without changing the processor's message format.""" 

129 if not isinstance(content_block, dict): 

130 return 

131 

132 block_type = content_block.get("type") 

133 if block_type == "image_url": 

134 image_url = content_block.get("image_url") 

135 if isinstance(image_url, str): 

136 content_block["image_url"] = cls._resolve_media_path(image_url, source_directory) 

137 elif isinstance(image_url, dict) and isinstance(image_url.get("url"), str): 

138 image_url["url"] = cls._resolve_media_path(image_url["url"], source_directory) 

139 return 

140 

141 media_field = content_block.get(block_type) 

142 if block_type in ("image", "video", "audio") and isinstance(media_field, str): 

143 content_block[block_type] = cls._resolve_media_path(media_field, source_directory) 

144 

145 @staticmethod 

146 def _resolve_media_path(media_path: str, source_directory: str) -> str: 

147 """Keep remote and absolute references, and anchor local relative paths.""" 

148 if os.path.isabs(media_path) or media_path.startswith(("http://", "https://", "data:")): 

149 return media_path 

150 

151 resolved_path = os.path.normpath(os.path.join(source_directory, media_path)) 

152 return resolved_path 

153 

154 

155class AutoProcessorTransform(OmniDataTransform): 

156 """Encode a conversation and supervise only its final assistant response.""" 

157 

158 def encode_sample(self, sample: dict[str, Any]) -> dict[str, Any]: 

159 """Build model inputs and assistant-only labels from processor messages.""" 

160 messages = self.prepare_messages(sample) 

161 if len(messages) < 2 or messages[-1].get("role") != "assistant": 

162 raise ValueError("Omni training sample must end with a final assistant message") 

163 

164 encoded_sample = dict(self.processor.apply_chat_template( 

165 messages, 

166 tokenize=True, 

167 return_dict=True, 

168 return_tensors="pt", 

169 add_generation_prompt=False, 

170 )) 

171 prompt_sample = self.processor.apply_chat_template( 

172 messages[:-1], 

173 tokenize=True, 

174 return_dict=True, 

175 return_tensors="pt", 

176 add_generation_prompt=True, 

177 ) 

178 

179 input_ids = encoded_sample["input_ids"].squeeze(0) 

180 prompt_ids = prompt_sample["input_ids"].squeeze(0) 

181 prompt_length = len(prompt_ids) 

182 if prompt_length >= len(input_ids) or not input_ids[:prompt_length].equal(prompt_ids): 

183 raise ValueError("Assistant prompt tokens must be a prefix of the full conversation") 

184 

185 labels = input_ids.clone() 

186 labels[:prompt_length] = IGNORE_INDEX 

187 if "mm_token_type_ids" in encoded_sample: 

188 labels[encoded_sample["mm_token_type_ids"].squeeze(0) != 0] = IGNORE_INDEX 

189 

190 for field in OMNI_CP_TOKEN_FIELDS.intersection(encoded_sample): 

191 encoded_sample[field] = encoded_sample[field].squeeze(0) 

192 encoded_sample["labels"] = labels 

193 return encoded_sample 

194 

195 

196SampleEncoder = Callable[[Any], Any] 

197 

198 

199@dataclass(frozen=True) 

200class _OmniTransformStrategy: 

201 """Resolve the encoding stages implemented by an Omni transform.""" 

202 

203 sample_encoder: SampleEncoder | None 

204 selected_sample_encoder: SampleEncoder | None 

205 

206 @classmethod 

207 def from_transform(cls, transform: OmniDataTransform) -> "_OmniTransformStrategy": 

208 """Build one strategy from Energon-style overridden hooks.""" 

209 encode_sample = cls._is_overridden(transform, "encode_sample") 

210 preencode_sample = cls._is_overridden(transform, "preencode_sample") 

211 postencode_sample = cls._is_overridden(transform, "postencode_sample") 

212 

213 if encode_sample and (preencode_sample or postencode_sample): 

214 raise TypeError( 

215 f"{type(transform).__name__} cannot combine encode_sample with " 

216 "preencode_sample or postencode_sample" 

217 ) 

218 if preencode_sample and not postencode_sample: 

219 raise TypeError( 

220 f"{type(transform).__name__} must pair preencode_sample with postencode_sample" 

221 ) 

222 if not encode_sample and not postencode_sample: 

223 raise TypeError( 

224 f"{type(transform).__name__} must implement encode_sample or postencode_sample" 

225 ) 

226 

227 # encode_sample -> encode before candidate selection 

228 # preencode_sample + postencode_sample -> prepare metadata, then defer full encoding 

229 # postencode_sample only -> preserve offline metadata, then defer encoding 

230 sample_encoder = None 

231 if preencode_sample: 

232 sample_encoder = transform.preencode_sample 

233 elif encode_sample: 

234 sample_encoder = transform.encode_sample 

235 

236 selected_sample_encoder = None 

237 if postencode_sample: 

238 selected_sample_encoder = transform.postencode_sample 

239 

240 transform_strategy = cls( 

241 sample_encoder=sample_encoder, 

242 selected_sample_encoder=selected_sample_encoder, 

243 ) 

244 return transform_strategy 

245 

246 @staticmethod 

247 def _is_overridden(transform: OmniDataTransform, method_name: str) -> bool: 

248 """Return whether a transform hook differs from its base implementation.""" 

249 transform_method = getattr(type(transform), method_name) 

250 base_method = getattr(OmniDataTransform, method_name) 

251 method_overridden = transform_method is not base_method 

252 return method_overridden 

253 

254 def prepare_sample(self, sample: Any) -> Any: 

255 """Encode a sample or attach metadata before packing selection.""" 

256 prepared_sample = sample 

257 if self.sample_encoder is not None: 

258 prepared_sample = self.sample_encoder(sample) 

259 return prepared_sample 

260 

261 def encode_selected_sample(self, sample: Any) -> Any: 

262 """Finish encoding a sample after packing selection.""" 

263 encoded_sample = sample 

264 if self.selected_sample_encoder is not None: 

265 encoded_sample = self.selected_sample_encoder(sample) 

266 return encoded_sample 

267 

268 def prepare_samples(self, raw_sample: Any) -> list[Any]: 

269 """Return encoded, preencoded, or unchanged samples for candidate selection.""" 

270 prepared_sample = self.prepare_sample(raw_sample) 

271 prepared_samples = [prepared_sample] 

272 return prepared_samples 

273 

274 

275__all__ = [ 

276 "build_auto_processor", 

277 "OmniDataTransform", 

278 "AutoProcessorTransform", 

279]