[Bugfix] 修复多模态tokenizer特殊token拆分逻辑，避免恶意token导致服务异常-MindIE-LLM-AtomGit

Mexamples/atb_models/atb_llm/models/base/router.py+139-120

	@@ -28,8 +28,13 @@ from ...utils.log.error_code import ErrorCode
28	from ...utils.configuration_utils import LLMConfig, GraphType	28	from ...utils.configuration_utils import LLMConfig, GraphType
29	from ...utils import file_utils	29	from ...utils import file_utils
30	from ...utils.parameter_validators import (	30	from ...utils.parameter_validators import (
31	- IntParameterValidator, DictionaryParameterValidator, BooleanParameterValidator,	31	+ IntParameterValidator,
32	- RangeParamaterValidator, FileParameterValidator, StringParameterValidator, Field	32	+ DictionaryParameterValidator,
		33	+ BooleanParameterValidator,
		34	+ RangeParamaterValidator,
		35	+ FileParameterValidator,
		36	+ StringParameterValidator,
		37	+ Field,
33	)	38	)
34		39
35		40
	@@ -40,8 +45,8 @@ QWEN2_MOE = "qwen2_moe"
40	QWEN3_MOE = "qwen3_moe"	45	QWEN3_MOE = "qwen3_moe"
41	DEEPSEEKV2 = "deepseekv2"	46	DEEPSEEKV2 = "deepseekv2"
42	LLAMA = "llama"	47	LLAMA = "llama"
43	-EOS_TOKEN_ID = 'eos_token_id'	48	+EOS_TOKEN_ID = "eos_token_id"
44	-FA_QUANT_TYPE = 'fa_quant_type'	49	+FA_QUANT_TYPE = "fa_quant_type"
45		50
46		51
47	def remove_part_of_generation_config(generation_config: dict) -> dict:	52	def remove_part_of_generation_config(generation_config: dict) -> dict:
	@@ -84,6 +89,7 @@ class BaseRouter:
84	This class should be inherited by the corresponding router subclasses of specified models. A specified model can use	89	This class should be inherited by the corresponding router subclasses of specified models. A specified model can use
85	a subclass router to find its custom properties.	90	a subclass router to find its custom properties.
86	"""	91	"""
		92	+
87	model_name_or_path: str = ""	93	model_name_or_path: str = ""
88		94
89	config_dict: Any = None	95	config_dict: Any = None
	@@ -98,6 +104,7 @@ class BaseRouter:
98	llm_config: LLMConfig = None	104	llm_config: LLMConfig = None
99		105
100	sub_model_path: str = ""	106	sub_model_path: str = ""
		107	+ is_multimodal: bool = False
101	prealloc_weight_mem_on_npu: bool = False	108	prealloc_weight_mem_on_npu: bool = False
102		109
103	# 初始化默认读取的autoconfig，各个模型可能会自定义，self.config会返回后续使用的config，注意不要循环依赖	110	# 初始化默认读取的autoconfig，各个模型可能会自定义，self.config会返回后续使用的config，注意不要循环依赖
	@@ -114,7 +121,7 @@ class BaseRouter:
114	_reasoning_parser: Any = None	121	_reasoning_parser: Any = None
115		122
116	def __post_init__(self):	123	def __post_init__(self):
117	- self.model_type = self.config_dict['model_type']	124	+ self.model_type = self.config_dict["model_type"]
118	if self.model_type == "glm4v":	125	if self.model_type == "glm4v":
119	self.model_type = "glm41v"	126	self.model_type = "glm41v"
120	if self.model_type == "chatglm" and "vision_config" in self.config_dict:	127	if self.model_type == "chatglm" and "vision_config" in self.config_dict:
	@@ -131,7 +138,7 @@ class BaseRouter:
131	self.model_type = "deepseekv2"	138	self.model_type = "deepseekv2"
132	if self.model_type == "minicpmv" and "MiniCPM-V-2_6" in self.model_name_or_path:	139	if self.model_type == "minicpmv" and "MiniCPM-V-2_6" in self.model_name_or_path:
133	self.model_type = "minicpm_qwen2_v2"	140	self.model_type = "minicpm_qwen2_v2"
134	- if self.model_type == "multi_modality" and 'aligner_config' in self.config_dict:	141	+ if self.model_type == "multi_modality" and "aligner_config" in self.config_dict:
135	self.model_type = "janus"	142	self.model_type = "janus"
136	if self.model_type == "qwen2_5_vl":	143	if self.model_type == "qwen2_5_vl":
137	self.model_type = "qwen2_vl"	144	self.model_type = "qwen2_vl"
	@@ -139,17 +146,17 @@ class BaseRouter:
139	self.model_type = "ernie_moe"	146	self.model_type = "ernie_moe"
140	self.model_type_cap = self.model_type.capitalize()	147	self.model_type_cap = self.model_type.capitalize()
141	if self.model_type_cap == "Qwen2_moe":	148	if self.model_type_cap == "Qwen2_moe":
142	- self.model_type_cap = self.model_type_cap.replace('_', '')	149	+ self.model_type_cap = self.model_type_cap.replace("_", "")
143	if self.model_type_cap == "Qwen2_audio":	150	if self.model_type_cap == "Qwen2_audio":
144	- self.model_type_cap = self.model_type_cap.replace('_', '')	151	+ self.model_type_cap = self.model_type_cap.replace("_", "")
145	if self.model_type_cap == "Qwen2_vl":	152	if self.model_type_cap == "Qwen2_vl":
146	- self.model_type_cap = self.model_type_cap.replace('_', '')	153	+ self.model_type_cap = self.model_type_cap.replace("_", "")
147	if self.model_type_cap == "Minicpm_qwen2_v2":	154	if self.model_type_cap == "Minicpm_qwen2_v2":
148	- self.model_type_cap = self.model_type_cap.replace('_', '')	155	+ self.model_type_cap = self.model_type_cap.replace("_", "")
149	if self.model_type_cap == "Ernie_moe":	156	if self.model_type_cap == "Ernie_moe":
150	- self.model_type_cap = self.model_type_cap.replace('_', '')	157	+ self.model_type_cap = self.model_type_cap.replace("_", "")
151	if self.model_type_cap == "Glm4_moe":	158	if self.model_type_cap == "Glm4_moe":
152	- self.model_type_cap = self.model_type_cap.replace('_', '')	159	+ self.model_type_cap = self.model_type_cap.replace("_", "")
153	if not self.tokenizer_path:	160	if not self.tokenizer_path:
154	self.tokenizer_path = self.model_name_or_path	161	self.tokenizer_path = self.model_name_or_path
155		162
	@@ -175,10 +182,12 @@ class BaseRouter:
175	"And the $llm_config.llm.ccl.enable_mc2 is %s in LLM's config. "	182	"And the $llm_config.llm.ccl.enable_mc2 is %s in LLM's config. "
176	"If one of them is False, the LCOC will not be enabled. "	183	"If one of them is False, the LCOC will not be enabled. "
177	"NOTE: The $ATB_LLM_LCOC_ENABLE environment variable will be deprecated after 2026/07/31.",	184	"NOTE: The $ATB_LLM_LCOC_ENABLE environment variable will be deprecated after 2026/07/31.",
178	- atb_llm_lcoc_enable, self.llm_config.llm.ccl.enable_mc2)	185	+ atb_llm_lcoc_enable,
		186	+ self.llm_config.llm.ccl.enable_mc2,
		187	+ )
179		188
180	self.process_tool_call_options()	189	self.process_tool_call_options()
181	-	190	+
182	@property	191	@property
183	def config(self):	192	def config(self):
184	"""The config property, which should not be overridden.	193	"""The config property, which should not be overridden.
	@@ -195,13 +204,13 @@ class BaseRouter:
195	self._generation_config = remove_part_of_generation_config(self._generation_config)	204	self._generation_config = remove_part_of_generation_config(self._generation_config)
196	self.config_dict.update(self._generation_config)	205	self.config_dict.update(self._generation_config)
197	self.config_dict["generation_config"] = self._generation_config	206	self.config_dict["generation_config"] = self._generation_config
198	- filename = os.path.join(self.model_name_or_path, 'quant_model_description.json')	207	+ filename = os.path.join(self.model_name_or_path, "quant_model_description.json")
199	if not file_utils.is_path_exists(filename):	208	if not file_utils.is_path_exists(filename):
200	- filename = os.path.join(self.model_name_or_path, self.sub_model_path, 'quant_model_description.json')	209	+ filename = os.path.join(self.model_name_or_path, self.sub_model_path, "quant_model_description.json")
201	if not file_utils.is_path_exists(filename):	210	if not file_utils.is_path_exists(filename):
202	quantize_type = self.config_dict.get(quantize, None)	211	quantize_type = self.config_dict.get(quantize, None)
203	if quantize_type:	212	if quantize_type:
204	- old_file = f'quant_model_description_{quantize_type.lower()}.json'	213	+ old_file = f"quant_model_description_{quantize_type.lower()}.json"
205	old_path = os.path.join(self.model_name_or_path, old_file)	214	old_path = os.path.join(self.model_name_or_path, old_file)
206	if file_utils.is_path_exists(old_path):	215	if file_utils.is_path_exists(old_path):
207	filename = old_path	216	filename = old_path
	@@ -212,7 +221,7 @@ class BaseRouter:
212	if not file_utils.is_path_exists(filename):	221	if not file_utils.is_path_exists(filename):
213	filename = None	222	filename = None
214	if filename is not None:	223	if filename is not None:
215	- with file_utils.safe_open(filename, 'r', check_link=False) as f:	224	+ with file_utils.safe_open(filename, "r", check_link=False) as f:
216	quant_descs = json.load(f)	225	quant_descs = json.load(f)
217	if quant_descs.get("model_quant_type") is not None and self.config_dict.get("quantize") is None:	226	if quant_descs.get("model_quant_type") is not None and self.config_dict.get("quantize") is None:
218	self.config_dict[quantize] = quant_descs.get("model_quant_type").lower()	227	self.config_dict[quantize] = quant_descs.get("model_quant_type").lower()
	@@ -238,7 +247,7 @@ class BaseRouter:
238	if not hasattr(self._config, quantize):	247	if not hasattr(self._config, quantize):
239	setattr(self._config, quantize, None)	248	setattr(self._config, quantize, None)
240	if self.max_position_embeddings is not None:	249	if self.max_position_embeddings is not None:
241	- setattr(self._config, 'max_position_embeddings', self.max_position_embeddings)	250	+ setattr(self._config, "max_position_embeddings", self.max_position_embeddings)
242	return self._config	251	return self._config
243		252
244	@property	253	@property
	@@ -296,9 +305,11 @@ class BaseRouter:
296	@property	305	@property
297	def custom_chat_template(self):	306	def custom_chat_template(self):
298	"""The custom chat template property, which should not be overridden."""	307	"""The custom chat template property, which should not be overridden."""
299	- if self._custom_chat_template is None \	308	+ if (
300	- and self.llm_config is not None \	309	+ self._custom_chat_template is None
301	- and getattr(self.llm_config.llm, "chat_template", None):	310	+ and self.llm_config is not None
		311	+ and getattr(self.llm_config.llm, "chat_template", None)
		312	+ ):
302	self._custom_chat_template = self.get_custom_chat_template()	313	self._custom_chat_template = self.get_custom_chat_template()
303	return self._custom_chat_template	314	return self._custom_chat_template
304		315
	@@ -321,45 +332,45 @@ class BaseRouter:
321	"""The validation of values in config."""	332	"""The validation of values in config."""
322		333
323	vocab_size = 0	334	vocab_size = 0
324	- vocab_size_field = 'vocab_size'	335	+ vocab_size_field = "vocab_size"
325	if hasattr(config, vocab_size_field):	336	if hasattr(config, vocab_size_field):
326	vocab_size = getattr(config, vocab_size_field)	337	vocab_size = getattr(config, vocab_size_field)
327	attribute_ranges = {	338	attribute_ranges = {
328	vocab_size_field: (1, 2147483647),	339	vocab_size_field: (1, 2147483647),
329	- 'max_position_embeddings': (1, 2147483647),	340	+ "max_position_embeddings": (1, 2147483647),
330	- 'hidden_size': (1, 2147483647),	341	+ "hidden_size": (1, 2147483647),
331	- 'intermediate_size': (1, 2147483647),	342	+ "intermediate_size": (1, 2147483647),
332	- 'num_hidden_layers': (1, 1000),	343	+ "num_hidden_layers": (1, 1000),
333	- 'num_attention_heads': (1, 10000),	344	+ "num_attention_heads": (1, 10000),
334	- 'initializer_range': (0, 2147483647),	345	+ "initializer_range": (0, 2147483647),
335	- 'rms_norm_eps': (0, 1),	346	+ "rms_norm_eps": (0, 1),
336	- 'pad_token_id': (-1, vocab_size),	347	+ "pad_token_id": (-1, vocab_size),
337	- 'bos_token_id': (0, vocab_size - 1),	348	+ "bos_token_id": (0, vocab_size - 1),
338	EOS_TOKEN_ID: (0, vocab_size - 1),	349	EOS_TOKEN_ID: (0, vocab_size - 1),
339	- 'temperature': (0, 2),	350	+ "temperature": (0, 2),
340	- 'top_k': (-1, vocab_size),	351	+ "top_k": (-1, vocab_size),
341	- 'top_p': (0, 1),	352	+ "top_p": (0, 1),
342	- 'repetition_penalty': (0, 2),	353	+ "repetition_penalty": (0, 2),
343	- 'frequency_penalty': (-2, 2),	354	+ "frequency_penalty": (-2, 2),
344	- 'presence_penalty': (-2, 2),	355	+ "presence_penalty": (-2, 2),
345	- 'max_new_tokens': (0, 10000000),	356	+ "max_new_tokens": (0, 10000000),
346	- 'max_length': (0, 2147483647),	357	+ "max_length": (0, 2147483647),
347	- 'user_token_id': (0, vocab_size-1),	358	+ "user_token_id": (0, vocab_size - 1),
348	- 'assistant_token_id': (0, vocab_size-1),	359	+ "assistant_token_id": (0, vocab_size - 1),
349	- 'system_token_id': (0, vocab_size-1),	360	+ "system_token_id": (0, vocab_size - 1),
350	- 'bot_token_id': (0, vocab_size-1)	361	+ "bot_token_id": (0, vocab_size - 1),
351	}	362	}
352	if config.is_reasoning_model:	363	if config.is_reasoning_model:
353	attribute_ranges["start_reasoning_token_id"] = (0, vocab_size - 1)	364	attribute_ranges["start_reasoning_token_id"] = (0, vocab_size - 1)
354	attribute_ranges["end_reasoning_token_id"] = (0, vocab_size - 1)	365	attribute_ranges["end_reasoning_token_id"] = (0, vocab_size - 1)
355	if hasattr(config, "head_dim"):	366	if hasattr(config, "head_dim"):
356	- attribute_ranges['head_dim'] = (1, 1000)	367	+ attribute_ranges["head_dim"] = (1, 1000)
357	if hasattr(config, "num_key_value_heads"):	368	if hasattr(config, "num_key_value_heads"):
358	- attribute_ranges['num_key_value_heads'] = (1, 10000)	369	+ attribute_ranges["num_key_value_heads"] = (1, 10000)
359	-	370	+
360	gen_config = config.generation_config	371	gen_config = config.generation_config
361	gen_config_dict = vars(gen_config)	372	gen_config_dict = vars(gen_config)
362	-	373	+
363	for attr, value in gen_config_dict.items():	374	for attr, value in gen_config_dict.items():
364	if value is not None:	375	if value is not None:
365	min_val, max_val = attribute_ranges.get(attr, (None, None))	376	min_val, max_val = attribute_ranges.get(attr, (None, None))
	@@ -380,51 +391,51 @@ class BaseRouter:
380	continue	391	continue
381	check_value(attr, value, min_val, max_val)	392	check_value(attr, value, min_val, max_val)
382		393
383	- if getattr(config, 'repetition_penalty', None) == 0:	394	+ if getattr(config, "repetition_penalty", None) == 0:
384	raise ValueError("repetition_penalty should not be 0.")	395	raise ValueError("repetition_penalty should not be 0.")
385	- if not isinstance(getattr(config, 'do_sample', None), bool):	396	+ if not isinstance(getattr(config, "do_sample", None), bool):
386	raise ValueError("do_sample must be bool.")	397	raise ValueError("do_sample must be bool.")
387		398
388	@classmethod	399	@classmethod
389	def get_llm_config_validators(cls):	400	def get_llm_config_validators(cls):
390	llm_config_validators = {"models": {}}	401	llm_config_validators = {"models": {}}
391	- llm_validators = DictionaryParameterValidator({	402	+ llm_validators = DictionaryParameterValidator(
392	- "ccl": DictionaryParameterValidator({	403	+ {
393	- "enable_mc2": BooleanParameterValidator()	404	+ "ccl": DictionaryParameterValidator({"enable_mc2": BooleanParameterValidator()}),
394	- }),	405	+ "stream_options": DictionaryParameterValidator({"micro_batch": BooleanParameterValidator()}),
395	- "stream_options": DictionaryParameterValidator({	406	+ "engine": DictionaryParameterValidator(
396	- "micro_batch": BooleanParameterValidator()	407	+ {"graph": RangeParamaterValidator(range_list=[GraphType.CPP, GraphType.PYTHON])}
397	- }),	408	+ ),
398	- "engine": DictionaryParameterValidator({	409	+ "parallel_options": DictionaryParameterValidator(
399	- "graph": RangeParamaterValidator(range_list=[GraphType.CPP, GraphType.PYTHON])	410	+ {
400	- }),	411	+ "o_proj_local_tp": IntParameterValidator(Field(ge=1, le=16), special_values=[-1]),
401	- "parallel_options": DictionaryParameterValidator({	412	+ "lm_head_local_tp": IntParameterValidator(Field(ge=1, le=16), special_values=[-1]),
402	- "o_proj_local_tp": IntParameterValidator(Field(ge=1, le=16), special_values=[-1]),	413	+ "hccl_buffer": IntParameterValidator(Field(ge=1)),
403	- "lm_head_local_tp": IntParameterValidator(Field(ge=1, le=16), special_values=[-1]),	414	+ "hccl_moe_ep_buffer": IntParameterValidator(Field(ge=512)),
404	- "hccl_buffer": IntParameterValidator(Field(ge=1)),	415	+ "hccl_moe_tp_buffer": IntParameterValidator(Field(ge=64)),
405	- "hccl_moe_ep_buffer": IntParameterValidator(Field(ge=512)),	416	+ }
406	- "hccl_moe_tp_buffer": IntParameterValidator(Field(ge=64)),	417	+ ),
407	- }),	418	+ "pmcc_obfuscation_options": DictionaryParameterValidator(
408	- "pmcc_obfuscation_options": DictionaryParameterValidator({	419	+ {
409	- "enable_model_obfuscation": BooleanParameterValidator(),	420	+ "enable_model_obfuscation": BooleanParameterValidator(),
410	- "data_obfuscation_ca_dir": FileParameterValidator(allow_none=True),	421	+ "data_obfuscation_ca_dir": FileParameterValidator(allow_none=True),
411	- "kms_agent_port": IntParameterValidator(Field(ge=1, le=65535)),	422	+ "kms_agent_port": IntParameterValidator(Field(ge=1, le=65535)),
412	- }),	423	+ }
413	- "kv_cache_options": DictionaryParameterValidator({	424	+ ),
414	- "enable_nz": BooleanParameterValidator()	425	+ "kv_cache_options": DictionaryParameterValidator({"enable_nz": BooleanParameterValidator()}),
415	- }),	426	+ "weights_options": DictionaryParameterValidator(
416	- "weights_options": DictionaryParameterValidator(	427	+ element_validator_mapping={"low_cpu_memory_mode": BooleanParameterValidator()},
417	- element_validator_mapping={	428	+ allow_addition_key=False,
418	- "low_cpu_memory_mode": BooleanParameterValidator()	429	+ ),
419	- },	430	+ "enable_reasoning": BooleanParameterValidator(),
420	- allow_addition_key=False	431	+ "tool_call_options": DictionaryParameterValidator(
421	- ),	432	+ {
422	- "enable_reasoning": BooleanParameterValidator(),	433	+ "tool_call_parser": StringParameterValidator(Field(max_length=1024), allow_none=True),
423	- "tool_call_options": DictionaryParameterValidator({	434	+ }
424	- "tool_call_parser": StringParameterValidator(Field(max_length=1024), allow_none=True),	435	+ ),
425	- }),	436	+ "chat_template": FileParameterValidator(allow_none=True),
426	- "chat_template": FileParameterValidator(allow_none=True)	437	+ }
427	- })	438	+ )
428	llm_config_validators["llm"] = llm_validators	439	llm_config_validators["llm"] = llm_validators
429	return llm_config_validators	440	return llm_config_validators
430		441
	@@ -440,9 +451,10 @@ class BaseRouter:
440	tool_call_parser_name = getattr(tool_call_options, "tool_call_parser", None)	451	tool_call_parser_name = getattr(tool_call_options, "tool_call_parser", None)
441	if tool_call_parser_name:	452	if tool_call_parser_name:
442	if tool_call_parser_name not in valid_tool_call_processors:	453	if tool_call_parser_name not in valid_tool_call_processors:
443	- warn_msg = \	454	+ warn_msg = (
444	- f"Invalid tool_call_parser: {tool_call_parser_name}. " \	455	+ f"Invalid tool_call_parser: {tool_call_parser_name}. "
445	f"Please chose from {valid_tool_call_processors.keys()}"	456	f"Please chose from {valid_tool_call_processors.keys()}"
		457	+ )
446	logger.warning(message_filter(warn_msg))	458	logger.warning(message_filter(warn_msg))
447	logger.warning("Will use default tool_call_parser of model.")	459	logger.warning("Will use default tool_call_parser of model.")
448	return	460	return
	@@ -497,9 +509,8 @@ class BaseRouter:
497	"""The default method to get config class."""	509	"""The default method to get config class."""
498	model_file_dir_name = f"atb_llm.models.{self.model_type}."	510	model_file_dir_name = f"atb_llm.models.{self.model_type}."
499	if self.model_version:	511	if self.model_version:
500	- model_file_dir_name = model_file_dir_name + \	512	+ model_file_dir_name = model_file_dir_name + f"{self.model_version}."
501	- f"{self.model_version}."	513	+ config_file_name = f"config_{self.model_type}"
502	- config_file_name = f'config_{self.model_type}'
503	module_path = f"{model_file_dir_name}{config_file_name}"	514	module_path = f"{model_file_dir_name}{config_file_name}"
504	module = importlib.import_module(module_path)	515	module = importlib.import_module(module_path)
505	config_cls_name = f"{self.model_type_cap}Config"	516	config_cls_name = f"{self.model_type_cap}Config"
	@@ -538,9 +549,8 @@ class BaseRouter:
538	if self.llm_config is not None and self.llm_config.llm.engine.graph == GraphType.PYTHON:	549	if self.llm_config is not None and self.llm_config.llm.engine.graph == GraphType.PYTHON:
539	enable_v3 = True	550	enable_v3 = True
540	if self.model_version:	551	if self.model_version:
541	- model_file_dir_name = model_file_dir_name + \	552	+ model_file_dir_name = model_file_dir_name + f"{self.model_version}."
542	- f"{self.model_version}."	553	+ model_file_name = "flash_causal" if self.is_flash_causal_lm else "causal"
543	- model_file_name = 'flash_causal' if self.is_flash_causal_lm else 'causal'
544	if self.embedding_model_name: # for embedding model, example: gte-qwen2	554	if self.embedding_model_name: # for embedding model, example: gte-qwen2
545	module_path = f"{model_file_dir_name}{model_file_name}_{self.model_type}_{self.embedding_model_name}"	555	module_path = f"{model_file_dir_name}{model_file_name}_{self.model_type}_{self.embedding_model_name}"
546	else:	556	else:
	@@ -577,7 +587,7 @@ class BaseRouter:
577	padding_side="left",	587	padding_side="left",
578	truncation_side="left",	588	truncation_side="left",
579	trust_remote_code=self.trust_remote_code,	589	trust_remote_code=self.trust_remote_code,
580	- use_fast=True	590	+ use_fast=True,
581	)	591	)
582		592
583	def get_reasoning_parser(self):	593	def get_reasoning_parser(self):
	@@ -596,57 +606,66 @@ class BaseRouter:
596	try:	606	try:
597	tools_call_processor = ToolParserManager.get_tool_call_processor(self.tool_call_parser)(self.tokenizer)	607	tools_call_processor = ToolParserManager.get_tool_call_processor(self.tool_call_parser)(self.tokenizer)
598	except KeyError as e:	608	except KeyError as e:
599	- warn_msg = \	609	+ warn_msg = (
600	- f"Cannot use {self.tool_call_parser} ToolsCallProcessor: {e}. " \	610	+ f"Cannot use {self.tool_call_parser} ToolsCallProcessor: {e}. Will use default ToolsCallProcessor."
601	- f"Will use default ToolsCallProcessor."	611	+ )
602	logger.warning(message_filter(warn_msg))	612	logger.warning(message_filter(warn_msg))
603	tools_call_processor = ToolsCallProcessor(self.model_version)	613	tools_call_processor = ToolsCallProcessor(self.model_version)
604	return tools_call_processor	614	return tools_call_processor
605	-	615	+
606	def _check_llm_config_model_type(self, model_type):	616	def _check_llm_config_model_type(self, model_type):
607	"""check the compatibility between llm_config and model_type"""	617	"""check the compatibility between llm_config and model_type"""
608	if self.llm_config.llm.stream_options.micro_batch and model_type not in [QWEN2, QWEN3, DEEPSEEKV2, GLM4_1V]:	618	if self.llm_config.llm.stream_options.micro_batch and model_type not in [QWEN2, QWEN3, DEEPSEEKV2, GLM4_1V]:
609	- error_msg = f"stream_options.micro_batch only supports [qwen2, qwen3, deepseekv2, glm4.1v], " \	619	+ error_msg = (
610	- f"but model_type is {model_type}, please refer to the MindIE official document" \	620	+ f"stream_options.micro_batch only supports [qwen2, qwen3, deepseekv2, glm4.1v], "
611	- f"and modify config.json."	621	+ f"but model_type is {model_type}, please refer to the MindIE official document"
		622	+ f"and modify config.json."
		623	+ )
612	logger.error(error_msg, ErrorCode.ATB_MODELS_PARAM_INVALID)	624	logger.error(error_msg, ErrorCode.ATB_MODELS_PARAM_INVALID)
613	raise ValueError(error_msg)	625	raise ValueError(error_msg)
614	-	626	+
615	if self.llm_config.llm.engine.graph == GraphType.PYTHON and model_type not in [LLAMA, QWEN2, QWEN3]:	627	if self.llm_config.llm.engine.graph == GraphType.PYTHON and model_type not in [LLAMA, QWEN2, QWEN3]:
616	- error_msg = f"engine.graph only supports [llama, qwen2, qwen3] when set to \"python\", " \	628	+ error_msg = (
617	- f"but model_type is {model_type}, please refer to the MindIE official document" \	629	+ f'engine.graph only supports [llama, qwen2, qwen3] when set to "python", '
618	- f"and modify config.json."	630	+ f"but model_type is {model_type}, please refer to the MindIE official document"
		631	+ f"and modify config.json."
		632	+ )
619	logger.error(error_msg, ErrorCode.ATB_MODELS_PARAM_INVALID)	633	logger.error(error_msg, ErrorCode.ATB_MODELS_PARAM_INVALID)
620	raise ValueError(error_msg)	634	raise ValueError(error_msg)
621	-	635	+
622	if self.llm_config.llm.kv_cache_options.enable_nz and model_type not in [DEEPSEEKV2]:	636	if self.llm_config.llm.kv_cache_options.enable_nz and model_type not in [DEEPSEEKV2]:
623	- error_msg = f"kv_cache_options.enable_nz only supports [deepseekv2] when set to true, " \	637	+ error_msg = (
624	- f"but model_type is {model_type}, please refer to the MindIE official document" \	638	+ f"kv_cache_options.enable_nz only supports [deepseekv2] when set to true, "
625	- f"and modify config.json."	639	+ f"but model_type is {model_type}, please refer to the MindIE official document"
		640	+ f"and modify config.json."
		641	+ )
626	logger.error(error_msg, ErrorCode.ATB_MODELS_PARAM_INVALID)	642	logger.error(error_msg, ErrorCode.ATB_MODELS_PARAM_INVALID)
627	raise ValueError(error_msg)	643	raise ValueError(error_msg)
628		644
629	# The following features are configured but are not applicable to some models.	645	# The following features are configured but are not applicable to some models.
630	if self.llm_config.llm.parallel_options.o_proj_local_tp != -1 and model_type != DEEPSEEKV2:	646	if self.llm_config.llm.parallel_options.o_proj_local_tp != -1 and model_type != DEEPSEEKV2:
631	logger.warning(f"parallel_options.o_proj_local_tp only supports deepseek, {model_type} won't be affected.")	647	logger.warning(f"parallel_options.o_proj_local_tp only supports deepseek, {model_type} won't be affected.")
632	-	648	+
633	if self.llm_config.llm.parallel_options.lm_head_local_tp != -1 and model_type != DEEPSEEKV2:	649	if self.llm_config.llm.parallel_options.lm_head_local_tp != -1 and model_type != DEEPSEEKV2:
634	- logger.warning(f"parallel_options.lm_head_local_tp only supports deepseek, " \	650	+ logger.warning(f"parallel_options.lm_head_local_tp only supports deepseek, {model_type} won't be affected.")
635	- f"{model_type} won't be affected.")
636		651
637	weights_options = self.llm_config.get("llm").get("weights_options")	652	weights_options = self.llm_config.get("llm").get("weights_options")
638	if weights_options is not None and weights_options.low_cpu_memory_mode:	653	if weights_options is not None and weights_options.low_cpu_memory_mode:
639	engine_graph = self.llm_config.get("llm").get("engine").get("graph")	654	engine_graph = self.llm_config.get("llm").get("engine").get("graph")
640	if engine_graph != GraphType.PYTHON or model_type not in [QWEN2, QWEN3]:	655	if engine_graph != GraphType.PYTHON or model_type not in [QWEN2, QWEN3]:
641	- logger.warning(f"low_cpu_memory_mode is only effective for [qwen2, qwen3] " \	656	+ logger.warning(
642	- f"when engine graph is set to \"python\". " \	657	+ f"low_cpu_memory_mode is only effective for [qwen2, qwen3] "
643	- f"For {model_type} with engine graph {engine_graph}, this mode will not work.")	658	+ f'when engine graph is set to "python". '
		659	+ f"For {model_type} with engine graph {engine_graph}, this mode will not work."
		660	+ )
644		661
645	if self.llm_config.llm.enable_reasoning and model_type not in [QWEN3, QWEN3_MOE, DEEPSEEKV2]:	662	if self.llm_config.llm.enable_reasoning and model_type not in [QWEN3, QWEN3_MOE, DEEPSEEKV2]:
646	- logger.warning(f"enable_reasoning only supports [qwen3, qwen3_moe, deepseek] when set to true, " \	663	+ logger.warning(
647	- f"{model_type} won't be affected.")	664	+ f"enable_reasoning only supports [qwen3, qwen3_moe, deepseek] when set to true, "
		665	+ f"{model_type} won't be affected."
		666	+ )
648		667
649	def _check_llm_config_orthogonality(self):	668	def _check_llm_config_orthogonality(self):
650	"""check the orthogonality within llm_config"""	669	"""check the orthogonality within llm_config"""
651	if self.llm_config.llm.stream_options.micro_batch and self.llm_config.llm.ccl.enable_mc2:	670	if self.llm_config.llm.stream_options.micro_batch and self.llm_config.llm.ccl.enable_mc2:
652	- raise ValueError("ccl.enable_mc2 and stream_options.micro_batch can't be both true.")	671	+ raise ValueError("ccl.enable_mc2 and stream_options.micro_batch can't be both true.")

Mexamples/atb_models/atb_llm/models/glm41v/router_glm41v.py+8-14

	@@ -53,6 +53,7 @@ messages_template = [
53		53
54	@dataclass	54	@dataclass
55	class Glm41vRouter(BaseRouter):	55	class Glm41vRouter(BaseRouter):
		56	+ is_multimodal: bool = True
56	_processor: Any = None	57	_processor: Any = None
57		58
58	@property	59	@property
	@@ -62,11 +63,11 @@ class Glm41vRouter(BaseRouter):
62	elif self._processor is None:	63	elif self._processor is None:
63	self._processor = self.get_processor()	64	self._processor = self.get_processor()
64	return self._processor	65	return self._processor
65	-	66	+
66	def tokenize(self, inputs: List[Dict], **kwargs):	67	def tokenize(self, inputs: List[Dict], **kwargs):
67	text = ""	68	text = ""
68	message_list = []	69	message_list = []
69	- shm_name_save_path = kwargs.get('shm_name_save_path', None)	70	+ shm_name_save_path = kwargs.get("shm_name_save_path", None)
70	self.input_builder.check_image_and_video_concurrency(inputs)	71	self.input_builder.check_image_and_video_concurrency(inputs)
71	for item in inputs:	72	for item in inputs:
72	if item.get(IMAGE, None):	73	if item.get(IMAGE, None):
	@@ -82,18 +83,11 @@ class Glm41vRouter(BaseRouter):
82	shm_name_save_path = self.input_builder.get_shm_name_save_path(video_path)	83	shm_name_save_path = self.input_builder.get_shm_name_save_path(video_path)
83	if item.get(TEXT, None):	84	if item.get(TEXT, None):
84	text = item[TEXT]	85	text = item[TEXT]
85	- message_list.append({	86	+ message_list.append({TYPE: TEXT, TEXT: text})
86	- TYPE: TEXT,
87	- TEXT: text
88	- })
89	messages = copy.deepcopy(messages_template)	87	messages = copy.deepcopy(messages_template)
90	messages[0]["content"] = message_list	88	messages[0]["content"] = message_list
91	inputs = self.processor.apply_chat_template(	89	inputs = self.processor.apply_chat_template(
92	- messages,	90	+ messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
93	- tokenize=True,
94	- add_generation_prompt=True,
95	- return_dict=True,
96	- return_tensors='pt'
97	)	91	)
98	shm_info = process_shared_memory(inputs, shm_name_save_path)	92	shm_info = process_shared_memory(inputs, shm_name_save_path)
99	input_ids = self.input_builder.update_token_id(inputs, shm_info)	93	input_ids = self.input_builder.update_token_id(inputs, shm_info)
	@@ -106,9 +100,9 @@ class Glm41vRouter(BaseRouter):
106	config_cls = self.get_config_cls()	100	config_cls = self.get_config_cls()
107	config = config_cls.from_dict(self.config_dict)	101	config = config_cls.from_dict(self.config_dict)
108	config.model_name_or_path = self.model_name_or_path	102	config.model_name_or_path = self.model_name_or_path
109	- if config.text_config.dtype == 'bfloat16':	103	+ if config.text_config.dtype == "bfloat16":
110	config.torch_dtype = torch.bfloat16	104	config.torch_dtype = torch.bfloat16
111	- elif config.text_config.dtype == 'float16':	105	+ elif config.text_config.dtype == "float16":
112	config.torch_dtype = torch.float16	106	config.torch_dtype = torch.float16
113	else:	107	else:
114	err_msg = "`torch_dtype` is only supported for type `float16` and `bfloat16`"	108	err_msg = "`torch_dtype` is only supported for type `float16` and `bfloat16`"
	@@ -121,4 +115,4 @@ class Glm41vRouter(BaseRouter):
121	return self.processor.tokenizer	115	return self.processor.tokenizer
122		116
123	def get_processor(self):	117	def get_processor(self):
124	- return safe_from_pretrained(AutoProcessor, self.model_name_or_path)	118	+ return safe_from_pretrained(AutoProcessor, self.model_name_or_path)

Mexamples/atb_models/atb_llm/models/glm4v/router_glm4v.py+42-37

	@@ -29,6 +29,8 @@ _IMG_TOKEN_LEN = 1600
29		29
30	@dataclass	30	@dataclass
31	class Glm4vRouter(BaseRouter):	31	class Glm4vRouter(BaseRouter):
		32	+ is_multimodal: bool = True
		33	+
32	def tokenize(self, inputs: List[Dict], **kwargs):	34	def tokenize(self, inputs: List[Dict], **kwargs):
33	pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else _PAD_TOKEN_ID	35	pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else _PAD_TOKEN_ID
34	boi_token_id = self.config.boi_token_id if self.config.boi_token_id is not None else _BOI_TOKEN_ID	36	boi_token_id = self.config.boi_token_id if self.config.boi_token_id is not None else _BOI_TOKEN_ID
	@@ -45,21 +47,25 @@ class Glm4vRouter(BaseRouter):
45	text = ele["text"]	47	text = ele["text"]
46	if text == "":	48	if text == "":
47	msg = "Input text is empty, please check."	49	msg = "Input text is empty, please check."
48	- logger.error(msg,	50	+ logger.error(msg, ErrorCode.ATB_MODELS_PARAM_INVALID)
49	- ErrorCode.ATB_MODELS_PARAM_INVALID)
50	raise ValueError(msg)	51	raise ValueError(msg)
51	-	52	+
52	# image preprocess	53	# image preprocess
53	from PIL import Image	54	from PIL import Image
		55	+
54	image_pix = safe_open_image(Image, image_path).convert("RGB")	56	image_pix = safe_open_image(Image, image_path).convert("RGB")
55	- inputs = self.tokenizer.apply_chat_template([{"role": "user", "image": image_pix, "content": text}],	57	+ inputs = self.tokenizer.apply_chat_template(
56	- add_generation_prompt=True, tokenize=True, return_tensors="pt",	58	+ [{"role": "user", "image": image_pix, "content": text}],
57	- return_dict=True)	59	+ add_generation_prompt=True,
		60	+ tokenize=True,
		61	+ return_tensors="pt",
		62	+ return_dict=True,
		63	+ )
58	input_ids = inputs["input_ids"].flatten()	64	input_ids = inputs["input_ids"].flatten()
59	input_image = inputs["images"]	65	input_image = inputs["images"]
60		66
61	# save image to shared memory	67	# save image to shared memory
62	- shm_name_save_path = kwargs.get('shm_name_save_path', None)	68	+ shm_name_save_path = kwargs.get("shm_name_save_path", None)
63	if shm_name_save_path is None:	69	if shm_name_save_path is None:
64	shm_name_save_dir = os.path.dirname(os.path.dirname(image_path))	70	shm_name_save_dir = os.path.dirname(os.path.dirname(image_path))
65	shm_name_save_path = os.path.join(shm_name_save_dir, "shm_name.txt")	71	shm_name_save_path = os.path.join(shm_name_save_dir, "shm_name.txt")
	@@ -71,28 +77,28 @@ class Glm4vRouter(BaseRouter):
71	shape_value = encode_shape_to_int64(input_image.shape)	77	shape_value = encode_shape_to_int64(input_image.shape)
72		78
73	# prepare input_ids with encoded shm info	79	# prepare input_ids with encoded shm info
74	- img_padding = torch.full((_IMG_TOKEN_LEN, ), pad_token_id, dtype=input_ids.dtype)	80	+ img_padding = torch.full((_IMG_TOKEN_LEN,), pad_token_id, dtype=input_ids.dtype)
75	new_input_ids = []	81	new_input_ids = []
76	boi_pos = torch.where(torch.eq(input_ids, boi_token_id))[0]	82	boi_pos = torch.where(torch.eq(input_ids, boi_token_id))[0]
77	eoi_pos = torch.where(torch.eq(input_ids, eoi_token_id))[0]	83	eoi_pos = torch.where(torch.eq(input_ids, eoi_token_id))[0]
78	- new_input_ids.append(torch.cat((input_ids[:boi_pos + 1], img_padding, input_ids[eoi_pos:])))	84	+ new_input_ids.append(torch.cat((input_ids[: boi_pos + 1], img_padding, input_ids[eoi_pos:])))
79	new_input_ids = torch.cat(new_input_ids)	85	new_input_ids = torch.cat(new_input_ids)
80	new_input_ids[boi_pos + 1] = shm_name	86	new_input_ids[boi_pos + 1] = shm_name
81	new_input_ids[boi_pos + 2] = shape_value	87	new_input_ids[boi_pos + 2] = shape_value
82	return new_input_ids	88	return new_input_ids
83	-	89	+
84	def get_input_builder(self):	90	def get_input_builder(self):
85	if hasattr(self.config, "max_position_embeddings") and self.config.max_position_embeddings:	91	if hasattr(self.config, "max_position_embeddings") and self.config.max_position_embeddings:
86	return Glm4vInputBuilder(self.tokenizer, self.config, max_length=self.config.max_position_embeddings)	92	return Glm4vInputBuilder(self.tokenizer, self.config, max_length=self.config.max_position_embeddings)
87	return Glm4vInputBuilder(self.tokenizer, self.config)	93	return Glm4vInputBuilder(self.tokenizer, self.config)
88	-	94	+
89	def get_config(self):	95	def get_config(self):
90	config_cls = self.get_config_cls()	96	config_cls = self.get_config_cls()
91	config = config_cls.from_dict(self.config_dict)	97	config = config_cls.from_dict(self.config_dict)
92	config.model_name_or_path = self.model_name_or_path	98	config.model_name_or_path = self.model_name_or_path
93	self.check_config(config)	99	self.check_config(config)
94	return config	100	return config
95	-	101	+
96	def check_config(self, config):	102	def check_config(self, config):
97	vocab_size = 0	103	vocab_size = 0
98	vocab_size_string = "vocab_size"	104	vocab_size_string = "vocab_size"
	@@ -102,30 +108,30 @@ class Glm4vRouter(BaseRouter):
102	if hasattr(config, padded_vocab_size_string):	108	if hasattr(config, padded_vocab_size_string):
103	vocab_size = getattr(config, padded_vocab_size_string)	109	vocab_size = getattr(config, padded_vocab_size_string)
104	attribute_ranges = {	110	attribute_ranges = {
105	- 'kv_channels': (1, 2147483647),	111	+ "kv_channels": (1, 2147483647),
106	- 'multi_query_group_num': (1, 2147483647),	112	+ "multi_query_group_num": (1, 2147483647),
107	- 'num_key_value_heads': (1, 2147483647),	113	+ "num_key_value_heads": (1, 2147483647),
108	- 'layernorm_epsilon': (0, 1),	114	+ "layernorm_epsilon": (0, 1),
109	- 'attention_dropout': (0, 1),	115	+ "attention_dropout": (0, 1),
110	- 'ffn_hidden_size': (1, 2147483647),	116	+ "ffn_hidden_size": (1, 2147483647),
111	- 'hidden_dropout': (0, 1),	117	+ "hidden_dropout": (0, 1),
112	padded_vocab_size_string: (1, 2147483647),	118	padded_vocab_size_string: (1, 2147483647),
113	- 'rope_ratio': (0, 2147483647),	119	+ "rope_ratio": (0, 2147483647),
114	- 'seq_length': (1, 2147483647),	120	+ "seq_length": (1, 2147483647),
115	vocab_size_string: (1, 2147483647),	121	vocab_size_string: (1, 2147483647),
116	- 'max_position_embeddings': (1, 2147483647),	122	+ "max_position_embeddings": (1, 2147483647),
117	- 'hidden_size': (1, 2147483647),	123	+ "hidden_size": (1, 2147483647),
118	- 'intermediate_size': (1, 2147483647),	124	+ "intermediate_size": (1, 2147483647),
119	- 'num_hidden_layers': (1, 1000),	125	+ "num_hidden_layers": (1, 1000),
120	- 'num_attention_heads': (1, 10000),	126	+ "num_attention_heads": (1, 10000),
121	- 'initializer_range': (0, 2147483647),	127	+ "initializer_range": (0, 2147483647),
122	- 'rms_norm_eps': (0, 1),	128	+ "rms_norm_eps": (0, 1),
123	- 'pad_token_id': (0, vocab_size),	129	+ "pad_token_id": (0, vocab_size),
124	- 'bos_token_id': (0, vocab_size),	130	+ "bos_token_id": (0, vocab_size),
125	- 'eos_token_id': (0, vocab_size)	131	+ "eos_token_id": (0, vocab_size),
126	}	132	}
127	if hasattr(config, "head_dim"):	133	if hasattr(config, "head_dim"):
128	- attribute_ranges['head_dim'] = (1, 1000)	134	+ attribute_ranges["head_dim"] = (1, 1000)
129	for attr, (min_val, max_val) in attribute_ranges.items():	135	for attr, (min_val, max_val) in attribute_ranges.items():
130	if not hasattr(config, attr) or getattr(config, attr) is None:	136	if not hasattr(config, attr) or getattr(config, attr) is None:
131	continue	137	continue
	@@ -134,13 +140,12 @@ class Glm4vRouter(BaseRouter):
134	value = max(value)	140	value = max(value)
135	if value < min_val or value > max_val:	141	if value < min_val or value > max_val:
136	msg = f"The {attr} value in config must be between {min_val} and {max_val}."	142	msg = f"The {attr} value in config must be between {min_val} and {max_val}."
137	- logger.error(msg,	143	+ logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
138	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
139	raise ValueError(msg)	144	raise ValueError(msg)
140		145
141	def get_config_cls(self):	146	def get_config_cls(self):
142	model_file_dir_name = f"atb_llm.models.{self.model_type}."	147	model_file_dir_name = f"atb_llm.models.{self.model_type}."
143	- config_file_name = 'config'	148	+ config_file_name = "config"
144	module_path = f"{model_file_dir_name}{config_file_name}_{self.model_type}"	149	module_path = f"{model_file_dir_name}{config_file_name}_{self.model_type}"
145	module = importlib.import_module(module_path)	150	module = importlib.import_module(module_path)
146	config_cls_name = f"{self.model_type_cap}Config"	151	config_cls_name = f"{self.model_type_cap}Config"
	@@ -151,10 +156,10 @@ class Glm4vRouter(BaseRouter):
151	get_model_cls	156	get_model_cls
152	"""	157	"""
153	model_file_dir_name = f"atb_llm.models.{self.model_type}."	158	model_file_dir_name = f"atb_llm.models.{self.model_type}."
154	- model_file_name = 'flash_causal' if self.is_flash_causal_lm else 'causal'	159	+ model_file_name = "flash_causal" if self.is_flash_causal_lm else "causal"
155	module_path = f"{model_file_dir_name}{model_file_name}_{self.model_type}"	160	module_path = f"{model_file_dir_name}{model_file_name}_{self.model_type}"
156	module = importlib.import_module(module_path)	161	module = importlib.import_module(module_path)
157	model_cls_name = f"{self.model_type_cap}ForCausalLM"	162	model_cls_name = f"{self.model_type_cap}ForCausalLM"
158	if self.is_flash_causal_lm:	163	if self.is_flash_causal_lm:
159	model_cls_name = "Flash" + model_cls_name	164	model_cls_name = "Flash" + model_cls_name
160	- return getattr(module, model_cls_name)	165	+ return getattr(module, model_cls_name)

Mexamples/atb_models/atb_llm/models/internlmxcomposer2/router_internlmxcomposer2.py+4-2

	@@ -16,12 +16,14 @@ from ..base.model_utils import safe_get_tokenizer_from_pretrained
16		16
17	@dataclass	17	@dataclass
18	class Internlmxcomposer2Router(BaseRouter):	18	class Internlmxcomposer2Router(BaseRouter):
		19	+ is_multimodal: bool = True
		20	+
19	def get_tokenizer(self):	21	def get_tokenizer(self):
20	tokenizer = safe_get_tokenizer_from_pretrained(	22	tokenizer = safe_get_tokenizer_from_pretrained(
21	self.tokenizer_path,	23	self.tokenizer_path,
22	revision=self.revision,	24	revision=self.revision,
23	padding_side="left",	25	padding_side="left",
24	trust_remote_code=self.trust_remote_code,	26	trust_remote_code=self.trust_remote_code,
25	- use_fast=False	27	+ use_fast=False,
26	)	28	)
27	- return tokenizer	29	+ return tokenizer

Mexamples/atb_models/atb_llm/models/internvl/router_internvl.py+31-35

	@@ -27,12 +27,14 @@ _VIDEO = "video"
27	_TEXT = "text"	27	_TEXT = "text"
28		28
29	# 服务化支持torch分布式所需环境变量	29	# 服务化支持torch分布式所需环境变量
30	-os.environ.setdefault('MASTER_ADDR', 'localhost')	30	+os.environ.setdefault("MASTER_ADDR", "localhost")
31	-os.environ.setdefault('MASTER_PORT', '5678')	31	+os.environ.setdefault("MASTER_PORT", "5678")
32		32
33		33
34	@dataclass	34	@dataclass
35	class InternvlRouter(BaseRouter):	35	class InternvlRouter(BaseRouter):
		36	+ is_multimodal: bool = True
		37	+
36	def get_config(self):	38	def get_config(self):
37	config = InternvlConfig.from_pretrained(self.model_name_or_path)	39	config = InternvlConfig.from_pretrained(self.model_name_or_path)
38	if self.max_position_embeddings:	40	if self.max_position_embeddings:
	@@ -43,16 +45,17 @@ class InternvlRouter(BaseRouter):
43		45
44	def get_tokenizer(self):	46	def get_tokenizer(self):
45	try:	47	try:
46	- llm_model_architectures = self.config_dict['llm_config']['architectures'][0]	48	+ llm_model_architectures = self.config_dict["llm_config"]["architectures"][0]
47	except KeyError as e:	49	except KeyError as e:
48	- logger.error("`llm_config.architectures` does not exist! Check `config.json`.",	50	+ logger.error(
49	- ErrorCode.ATB_MODELS_MODEL_PARAM_JSON_INVALID)	51	+ "`llm_config.architectures` does not exist! Check `config.json`.",
		52	+ ErrorCode.ATB_MODELS_MODEL_PARAM_JSON_INVALID,
		53	+ )
50	raise ValueError("`llm_config.architectures` does not exist! Check `config.json`.") from e	54	raise ValueError("`llm_config.architectures` does not exist! Check `config.json`.") from e
51	-	55	+
52	if llm_model_architectures == INTERNLM2_ARCHITECTURE:	56	if llm_model_architectures == INTERNLM2_ARCHITECTURE:
53	tokenizer = safe_get_tokenizer_from_pretrained(	57	tokenizer = safe_get_tokenizer_from_pretrained(
54	- self.model_name_or_path,	58	+ self.model_name_or_path, trust_remote_code=self.trust_remote_code
55	- trust_remote_code=self.trust_remote_code
56	)	59	)
57	elif llm_model_architectures == LLAMA_ARCHITECTURE:	60	elif llm_model_architectures == LLAMA_ARCHITECTURE:
58	tokenizer = safe_get_tokenizer_from_pretrained(	61	tokenizer = safe_get_tokenizer_from_pretrained(
	@@ -60,7 +63,7 @@ class InternvlRouter(BaseRouter):
60	revision=self.revision,	63	revision=self.revision,
61	padding_side="left",	64	padding_side="left",
62	trust_remote_code=self.trust_remote_code,	65	trust_remote_code=self.trust_remote_code,
63	- use_fast=False	66	+ use_fast=False,
64	)	67	)
65	elif llm_model_architectures == QWEN2_ARCHITECTURE:	68	elif llm_model_architectures == QWEN2_ARCHITECTURE:
66	tokenizer = safe_get_tokenizer_from_pretrained(	69	tokenizer = safe_get_tokenizer_from_pretrained(
	@@ -73,11 +76,13 @@ class InternvlRouter(BaseRouter):
73	"`llm_config.architectures` must in "	76	"`llm_config.architectures` must in "
74	f"[{LLAMA_ARCHITECTURE}, {INTERNLM2_ARCHITECTURE}, {QWEN2_ARCHITECTURE}], "	77	f"[{LLAMA_ARCHITECTURE}, {INTERNLM2_ARCHITECTURE}, {QWEN2_ARCHITECTURE}], "
75	f"got {llm_model_architectures}.",	78	f"got {llm_model_architectures}.",
76	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	79	+ ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE,
		80	+ )
77	raise ValueError(	81	raise ValueError(
78	"`llm_config.architectures` must in "	82	"`llm_config.architectures` must in "
79	f"[{LLAMA_ARCHITECTURE}, {INTERNLM2_ARCHITECTURE}, {QWEN2_ARCHITECTURE}], "	83	f"[{LLAMA_ARCHITECTURE}, {INTERNLM2_ARCHITECTURE}, {QWEN2_ARCHITECTURE}], "
80	- f"got {llm_model_architectures}.")	84	+ f"got {llm_model_architectures}."
		85	+ )
81	return tokenizer	86	return tokenizer
82		87
83	def get_input_builder(self):	88	def get_input_builder(self):
	@@ -91,16 +96,14 @@ class InternvlRouter(BaseRouter):
91	image_size = self.config.force_image_size or self.config.vision_config.image_size	96	image_size = self.config.force_image_size or self.config.vision_config.image_size
92	patch_size = self.config.vision_config.patch_size	97	patch_size = self.config.vision_config.patch_size
93	if patch_size == 0:	98	if patch_size == 0:
94	- logger.error('The vision `patch_size` of config can not be 0.',	99	+ logger.error("The vision `patch_size` of config can not be 0.", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
95	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	100	+ raise ValueError("The vision `patch_size` of config can not be 0.")
96	- raise ValueError('The vision `patch_size` of config can not be 0.')	101	+ num_image_token = int((image_size // patch_size) ** 2 * (self.config.downsample_ratio**2))
97	- num_image_token = int((image_size // patch_size) ** 2 * (self.config.downsample_ratio ** 2))	102	+
98	-
99	use_dynamic_prepro = False if self.config.ps_version == "v1" else True	103	use_dynamic_prepro = False if self.config.ps_version == "v1" else True
100		104
101	system_prompt = INTERNVL_SYSTEM_PROMPTS[self.config.ps_version][self.config.template]	105	system_prompt = INTERNVL_SYSTEM_PROMPTS[self.config.ps_version][self.config.template]
102	- query = ('<\|im_start\|>system\n'	106	+ query = f"<\|im_start\|>system\n{system_prompt}<\|im_end\|><\|im_start\|>user\n"
103	- f'{system_prompt}<\|im_end\|><\|im_start\|>user\n')
104		107
105	text = ""	108	text = ""
106	image_index = 1	109	image_index = 1
	@@ -113,12 +116,7 @@ class InternvlRouter(BaseRouter):
113	continue	116	continue
114	if _IMAGE in single_input:	117	if _IMAGE in single_input:
115	current_query, shm_name_value, shape_value = process_image_input(	118	current_query, shm_name_value, shape_value = process_image_input(
116	- single_input,	119	+ single_input, image_num, image_index, use_dynamic_prepro, num_image_token, shm_name_save_path
117	- image_num,
118	- image_index,
119	- use_dynamic_prepro,
120	- num_image_token,
121	- shm_name_save_path
122	)	120	)
123	query += current_query	121	query += current_query
124	image_index += 1	122	image_index += 1
	@@ -126,28 +124,26 @@ class InternvlRouter(BaseRouter):
126	shape_value_list.append(shape_value)	124	shape_value_list.append(shape_value)
127	elif _VIDEO in single_input:	125	elif _VIDEO in single_input:
128	current_query, shm_name_value, shape_value = process_video_input(	126	current_query, shm_name_value, shape_value = process_video_input(
129	- single_input,	127	+ single_input, use_dynamic_prepro, num_image_token, shm_name_save_path
130	- use_dynamic_prepro,
131	- num_image_token,
132	- shm_name_save_path
133	)	128	)
134	query += current_query	129	query += current_query
135	shm_name_list += shm_name_value	130	shm_name_list += shm_name_value
136	shape_value_list += shape_value	131	shape_value_list += shape_value
137	else:	132	else:
138	- logger.error("Unsupported media type, it should be a video or image, please check your input.",	133	+ logger.error(
139	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	134	+ "Unsupported media type, it should be a video or image, please check your input.",
		135	+ ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE,
		136	+ )
140	raise KeyError("Unsupported media type, it should be a video or image, please check your input.")	137	raise KeyError("Unsupported media type, it should be a video or image, please check your input.")
141	- query += f'{text}<\|im_end\|><\|im_start\|>assistant\n'	138	+ query += f"{text}<\|im_end\|><\|im_start\|>assistant\n"
142	query_ids = torch.tensor(self.tokenizer.encode(query))	139	query_ids = torch.tensor(self.tokenizer.encode(query))
143	bos_pos_set = torch.nonzero(query_ids == img_begin_id).view(-1)	140	bos_pos_set = torch.nonzero(query_ids == img_begin_id).view(-1)
144	eos_pos_set = torch.nonzero(query_ids == img_end_id).view(-1)	141	eos_pos_set = torch.nonzero(query_ids == img_end_id).view(-1)
145	for i, (bos_pos, eos_pos) in enumerate(zip(bos_pos_set, eos_pos_set)):	142	for i, (bos_pos, eos_pos) in enumerate(zip(bos_pos_set, eos_pos_set)):
146	if eos_pos - bos_pos < 3:	143	if eos_pos - bos_pos < 3:
147	- logger.error("tokenize input error.",	144	+ logger.error("tokenize input error.", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
148	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
149	raise ValueError("tokenize input error.")	145	raise ValueError("tokenize input error.")
150	query_ids[bos_pos + 1] = shm_name_list[i]	146	query_ids[bos_pos + 1] = shm_name_list[i]
151	query_ids[bos_pos + 2] = shape_value_list[i]	147	query_ids[bos_pos + 2] = shape_value_list[i]
152	-	148	+
153	- return query_ids	149	+ return query_ids

Mexamples/atb_models/atb_llm/models/janus/router_janus.py+4-3

	@@ -14,13 +14,14 @@ from ..base.router import BaseRouter
14		14
15	@dataclass	15	@dataclass
16	class JanusRouter(BaseRouter):	16	class JanusRouter(BaseRouter):
17	-	17	+ is_multimodal: bool = True
		18	+
18	def check_config_janus(self, config):	19	def check_config_janus(self, config):
19	super().check_config(config)	20	super().check_config(config)
20	-	21	+
21	def get_config(self):	22	def get_config(self):
22	config_cls = self.get_config_cls()	23	config_cls = self.get_config_cls()
23	config = config_cls.from_dict(self.config_dict)	24	config = config_cls.from_dict(self.config_dict)
24	config.model_name_or_path = self.model_name_or_path	25	config.model_name_or_path = self.model_name_or_path
25	self.check_config(config)	26	self.check_config(config)
26	- return config	27	+ return config

Mexamples/atb_models/atb_llm/models/llava/router_llava.py+21-34

	@@ -36,33 +36,26 @@ def from_list_format(list_format: List[Dict], image_token_str: str):
36	text += ele["text"]	36	text += ele["text"]
37	else:	37	else:
38	msg = "Unsupported element: " + str(ele)	38	msg = "Unsupported element: " + str(ele)
39	- logger.error(	39	+ logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
40	- msg,
41	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE
42	- )
43	raise KeyError(msg)	40	raise KeyError(msg)
44	return text	41	return text
45		42
46		43
47	@dataclass	44	@dataclass
48	class LlavaRouter(BaseRouter):	45	class LlavaRouter(BaseRouter):
		46	+ is_multimodal: bool = True
49		47
50	def tokenize(self, inputs, **kwargs):	48	def tokenize(self, inputs, **kwargs):
51	- pad_token_id = self.config.pad_token_id \	49	+ pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else _PAD_TOKEN_ID
52	- if self.config.pad_token_id is not None else _PAD_TOKEN_ID	50	+ image_token_id = self.config.image_token_index if self.config.image_token_index is not None else _IMAGE_TOKEN_ID
53	- image_token_id = self.config.image_token_index \
54	- if self.config.image_token_index is not None else _IMAGE_TOKEN_ID
55	image_token_str = self.tokenizer.decode(image_token_id)	51	image_token_str = self.tokenizer.decode(image_token_id)
56	prompt = from_list_format(inputs, image_token_str)	52	prompt = from_list_format(inputs, image_token_str)
57	input_ids = self.tokenizer(prompt, return_tensors="pt")["input_ids"].flatten()	53	input_ids = self.tokenizer(prompt, return_tensors="pt")["input_ids"].flatten()
58	new_input_ids = input_ids	54	new_input_ids = input_ids
59	- bos_pos = torch.where(torch.eq(input_ids, image_token_id))[0]	55	+ bos_pos = torch.where(torch.eq(input_ids, image_token_id))[0]
60	if bos_pos.size(0) % 2 != 0:	56	if bos_pos.size(0) % 2 != 0:
61	msg = "Ensure that your input_text does not contain the '<image>' character!"	57	msg = "Ensure that your input_text does not contain the '<image>' character!"
62	- logger.error(	58	+ logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
63	- msg,
64	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE
65	- )
66	raise ValueError(msg)	59	raise ValueError(msg)
67	image_num = bos_pos.size(0) // 2	60	image_num = bos_pos.size(0) // 2
68	expand_token_ids = []	61	expand_token_ids = []
	@@ -72,40 +65,37 @@ class LlavaRouter(BaseRouter):
72	path_token_len = bos_pos[i + 1] - bos_pos[i]	65	path_token_len = bos_pos[i + 1] - bos_pos[i]
73	if path_token_len < 0 or path_token_len >= _IMAGE_FEATURE_WIDTH:	66	if path_token_len < 0 or path_token_len >= _IMAGE_FEATURE_WIDTH:
74	msg = f"Token length in path must be greater than 0 and less than {_IMAGE_FEATURE_WIDTH}."	67	msg = f"Token length in path must be greater than 0 and less than {_IMAGE_FEATURE_WIDTH}."
75	- logger.error(	68	+ logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
76	- msg,
77	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE
78	- )
79	raise ValueError(msg)	69	raise ValueError(msg)
80	- text_token = input_ids[pre: bos_pos[i]]	70	+ text_token = input_ids[pre : bos_pos[i]]
81	- path_token = input_ids[bos_pos[i] + 1: bos_pos[i + 1] + 1]	71	+ path_token = input_ids[bos_pos[i] + 1 : bos_pos[i + 1] + 1]
82	pre = bos_pos[i + 1] + 1	72	pre = bos_pos[i + 1] + 1
83	new_path_token = torch.cat(	73	new_path_token = torch.cat(
84	[	74	[
85	torch.tensor([image_token_id], dtype=path_token.dtype),	75	torch.tensor([image_token_id], dtype=path_token.dtype),
86	- torch.full((_IMAGE_FEATURE_WIDTH - path_token_len - 1, ), pad_token_id, dtype=path_token.dtype),	76	+ torch.full((_IMAGE_FEATURE_WIDTH - path_token_len - 1,), pad_token_id, dtype=path_token.dtype),
87	- path_token	77	+ path_token,
88	]	78	]
89	)	79	)
90	if text_token.size(0) != 0:	80	if text_token.size(0) != 0:
91	expand_token_ids.append(text_token)	81	expand_token_ids.append(text_token)
92	-	82	+
93	if new_path_token.size(0) != 0:	83	if new_path_token.size(0) != 0:
94	expand_token_ids.append(new_path_token)	84	expand_token_ids.append(new_path_token)
95	-	85	+
96	text_token = input_ids[pre:]	86	text_token = input_ids[pre:]
97	if text_token.size(0) != 0:	87	if text_token.size(0) != 0:
98	expand_token_ids.append(text_token)	88	expand_token_ids.append(text_token)
99	-	89	+
100	if expand_token_ids:	90	if expand_token_ids:
101	new_input_ids = torch.cat(expand_token_ids)	91	new_input_ids = torch.cat(expand_token_ids)
102	return new_input_ids	92	return new_input_ids
103	-	93	+
104	def check_config_llava(self, config):	94	def check_config_llava(self, config):
105	super().check_config(config)	95	super().check_config(config)
106	attribute_ranges = {	96	attribute_ranges = {
107	- 'mm_hidden_size': (1, 2147483647),	97	+ "mm_hidden_size": (1, 2147483647),
108	- 'num_key_value_heads': (1, 2147483647),	98	+ "num_key_value_heads": (1, 2147483647),
109	}	99	}
110	for attr, (min_val, max_val) in attribute_ranges.items():	100	for attr, (min_val, max_val) in attribute_ranges.items():
111	if not hasattr(config, attr) or getattr(config, attr) is None:	101	if not hasattr(config, attr) or getattr(config, attr) is None:
	@@ -113,16 +103,13 @@ class LlavaRouter(BaseRouter):
113	value = getattr(config, attr)	103	value = getattr(config, attr)
114	if value < min_val or value > max_val:	104	if value < min_val or value > max_val:
115	msg = f"The {attr} value in config must be between {min_val} and {max_val}."	105	msg = f"The {attr} value in config must be between {min_val} and {max_val}."
116	- logger.error(	106	+ logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
117	- msg,
118	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE
119	- )
120	raise ValueError(msg)	107	raise ValueError(msg)
121		108
122	def get_config(self):	109	def get_config(self):
123	config = LlavaConfig.from_pretrained(self.model_name_or_path)	110	config = LlavaConfig.from_pretrained(self.model_name_or_path)
124	config.text_config = LlamaConfig.from_dict(config.text_config)	111	config.text_config = LlamaConfig.from_dict(config.text_config)
125	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	112	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
126	if self.max_position_embeddings:	113	if self.max_position_embeddings:
127	config.max_position_embeddings = self.max_position_embeddings	114	config.max_position_embeddings = self.max_position_embeddings
128	self.check_config_llava(config)	115	self.check_config_llava(config)
	@@ -137,5 +124,5 @@ class LlavaRouter(BaseRouter):
137	padding_side="left",	124	padding_side="left",
138	truncation_side="left",	125	truncation_side="left",
139	trust_remote_code=self.trust_remote_code,	126	trust_remote_code=self.trust_remote_code,
140	- use_fast=use_fast	127	+ use_fast=use_fast,
141	- )	128	+ )

Mexamples/atb_models/atb_llm/models/llava_next/router_llava_next.py+27-29

	@@ -26,15 +26,11 @@
26	# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,	26	# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
27	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	27	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
28	# See the Mulan PSL v2 for more details.	28	# See the Mulan PSL v2 for more details.
29	-import os
30	from dataclasses import dataclass	29	from dataclasses import dataclass
31	from typing import Any, Dict, List	30	from typing import Any, Dict, List
32		31
33	-import torch
34	-import numpy as np
35	from transformers import AutoProcessor	32	from transformers import AutoProcessor
36		33
37	-from atb_llm.utils.shm_utils import encode_shm_name_to_int64, encode_shape_to_int64, create_shm
38	from atb_llm.utils.log import logger	34	from atb_llm.utils.log import logger
39	from atb_llm.utils.log.error_code import ErrorCode	35	from atb_llm.utils.log.error_code import ErrorCode
40	from ..base.router import BaseRouter	36	from ..base.router import BaseRouter
	@@ -42,7 +38,6 @@ from .flash_causal_llava_next import LlavaNextConfig
42	from ..llama.config_llama import LlamaConfig	38	from ..llama.config_llama import LlamaConfig
43	from ..base.config import QuantizationConfig	39	from ..base.config import QuantizationConfig
44	from ..base.model_utils import safe_get_tokenizer_from_pretrained, safe_from_pretrained	40	from ..base.model_utils import safe_get_tokenizer_from_pretrained, safe_from_pretrained
45	-from .data_preprocess_llava_next import data_prerocess_llava_next
46	from .input_builder_llava_next import LlavaNextInputBuilder, get_token_ids	41	from .input_builder_llava_next import LlavaNextInputBuilder, get_token_ids
47		42
48	_VIDEO_TOKEN_ID = -2	43	_VIDEO_TOKEN_ID = -2
	@@ -55,22 +50,19 @@ def from_list_format(list_format: List[Dict], image_token_str: str, video_token_
55	for ele in list_format:	50	for ele in list_format:
56	if "image" in ele:	51	if "image" in ele:
57	if image_token_str is None:	52	if image_token_str is None:
58	- logger.error(f"The model seems doesn't support image inputs.",	53	+ logger.error("The model seems doesn't support image inputs.", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
59	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	54	+ raise ValueError("The model seems doesn't support image inputs.")
60	- raise ValueError(f"The model seems doesn't support image inputs.")
61	image_or_video += image_token_str + "\n"	55	image_or_video += image_token_str + "\n"
62	elif "video" in ele:	56	elif "video" in ele:
63	if video_token_str is None:	57	if video_token_str is None:
64	- logger.error(f"The model seems doesn't support image inputs.",	58	+ logger.error("The model seems doesn't support video inputs.", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
65	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	59	+ raise ValueError("The model seems doesn't support video inputs.")
66	- raise ValueError(f"The model seems doesn't support video inputs.")
67	image_or_video += video_token_str + "\n"	60	image_or_video += video_token_str + "\n"
68	elif "text" in ele:	61	elif "text" in ele:
69	text += ele["text"]	62	text += ele["text"]
70	else:	63	else:
71	- logger.error("Unsupported element: " + str(ele) + '.',	64	+ logger.error("Unsupported element: " + str(ele) + ".", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
72	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	65	+ raise KeyError("Unsupported element: " + str(ele) + ".")
73	- raise KeyError("Unsupported element: " + str(ele) + '.')
74	if image_or_video != "":	66	if image_or_video != "":
75	result_promt += image_or_video + "\n"	67	result_promt += image_or_video + "\n"
76	result_promt += text + "ASSISTANT:"	68	result_promt += text + "ASSISTANT:"
	@@ -79,8 +71,9 @@ def from_list_format(list_format: List[Dict], image_token_str: str, video_token_
79		71
80	@dataclass	72	@dataclass
81	class Llava_nextRouter(BaseRouter):	73	class Llava_nextRouter(BaseRouter):
		74	+ is_multimodal: bool = True
82	_processor: Any = None	75	_processor: Any = None
83	-	76	+
84	@property	77	@property
85	def model_chat_template_version(self):	78	def model_chat_template_version(self):
86	if "mistral" in self.config.text_model_name:	79	if "mistral" in self.config.text_model_name:
	@@ -104,34 +97,37 @@ class Llava_nextRouter(BaseRouter):
104	]	97	]
105	image_token_id = self.config.image_token_index	98	image_token_id = self.config.image_token_index
106	video_token_id = self.config.video_token_index if self.config.video_token_index else _VIDEO_TOKEN_ID	99	video_token_id = self.config.video_token_index if self.config.video_token_index else _VIDEO_TOKEN_ID
107	-	100	+
108	image_token_str = self.tokenizer.decode(image_token_id)	101	image_token_str = self.tokenizer.decode(image_token_id)
109	video_token_str = self.tokenizer.decode(video_token_id) if self.config.video_token_index else None	102	video_token_str = self.tokenizer.decode(video_token_id) if self.config.video_token_index else None
110		103
111	query = from_list_format(inputs, image_token_str, video_token_str)	104	query = from_list_format(inputs, image_token_str, video_token_str)
112	- new_input_ids = get_token_ids(query, image_or_video_path_list, self.config,	105	+ new_input_ids = get_token_ids(
113	- self.processor(), shm_name_save_path)	106	+ query, image_or_video_path_list, self.config, self.processor(), shm_name_save_path
		107	+ )
114	return new_input_ids	108	return new_input_ids
115		109
116	def check_config_llava(self, config):	110	def check_config_llava(self, config):
117	super().check_config(config)	111	super().check_config(config)
118	attribute_ranges = {	112	attribute_ranges = {
119	- 'mm_hidden_size': (1, 2147483647),	113	+ "mm_hidden_size": (1, 2147483647),
120	- 'num_key_value_heads': (1, 2147483647),	114	+ "num_key_value_heads": (1, 2147483647),
121	}	115	}
122	for attr, (min_val, max_val) in attribute_ranges.items():	116	for attr, (min_val, max_val) in attribute_ranges.items():
123	if not hasattr(config, attr) or getattr(config, attr) is None:	117	if not hasattr(config, attr) or getattr(config, attr) is None:
124	continue	118	continue
125	value = getattr(config, attr)	119	value = getattr(config, attr)
126	if value < min_val or value > max_val:	120	if value < min_val or value > max_val:
127	- logger.error(f"`self._config.{attr}` must be between {min_val} and {max_val}.",	121	+ logger.error(
128	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	122	+ f"`self._config.{attr}` must be between {min_val} and {max_val}.",
		123	+ ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE,
		124	+ )
129	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")	125	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")
130		126
131	def get_config(self):	127	def get_config(self):
132	config = LlavaNextConfig.from_pretrained(self.model_name_or_path)	128	config = LlavaNextConfig.from_pretrained(self.model_name_or_path)
133	config.text_config = LlamaConfig.from_dict(config.text_config)	129	config.text_config = LlamaConfig.from_dict(config.text_config)
134	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	130	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
135	if self.max_position_embeddings:	131	if self.max_position_embeddings:
136	config.max_position_embeddings = self.max_position_embeddings	132	config.max_position_embeddings = self.max_position_embeddings
137	self.check_config_llava(config)	133	self.check_config_llava(config)
	@@ -141,13 +137,15 @@ class Llava_nextRouter(BaseRouter):
141	def processor(self):	137	def processor(self):
142	if self._processor is not None:	138	if self._processor is not None:
143	return self._processor	139	return self._processor
144	- if self.config.text_config.num_hidden_layers == 60 and \	140	+ if (
145	- self.config.architectures[0] == "LlavaNextForConditionalGeneration":	141	+ self.config.text_config.num_hidden_layers == 60
		142	+ and self.config.architectures[0] == "LlavaNextForConditionalGeneration"
		143	+ ):
146	self._processor = safe_from_pretrained(AutoProcessor, self.config.model_name_or_path, use_fast=False)	144	self._processor = safe_from_pretrained(AutoProcessor, self.config.model_name_or_path, use_fast=False)
147	else:	145	else:
148	self._processor = safe_from_pretrained(AutoProcessor, self.config.model_name_or_path)	146	self._processor = safe_from_pretrained(AutoProcessor, self.config.model_name_or_path)
149	return self._processor	147	return self._processor
150	-	148	+
151	def is_llava_next34b(self):	149	def is_llava_next34b(self):
152	config = self.get_config()	150	config = self.get_config()
153	num_hidden_layers = config.text_config.num_hidden_layers	151	num_hidden_layers = config.text_config.num_hidden_layers
	@@ -155,11 +153,11 @@ class Llava_nextRouter(BaseRouter):
155	if num_hidden_layers in [60] and architectures[0] == "LlavaNextForConditionalGeneration":	153	if num_hidden_layers in [60] and architectures[0] == "LlavaNextForConditionalGeneration":
156	return True	154	return True
157	return False	155	return False
158	-	156	+
159	def get_tokenizer(self):	157	def get_tokenizer(self):
160	use_fast = True	158	use_fast = True
161	if self.is_llava_next34b():	159	if self.is_llava_next34b():
162	- use_fast = False	160	+ use_fast = False
163	return safe_get_tokenizer_from_pretrained(	161	return safe_get_tokenizer_from_pretrained(
164	self.tokenizer_path,	162	self.tokenizer_path,
165	revision=self.revision,	163	revision=self.revision,
	@@ -170,4 +168,4 @@ class Llava_nextRouter(BaseRouter):
170	)	168	)
171		169
172	def get_input_builder(self):	170	def get_input_builder(self):
173	- return LlavaNextInputBuilder(self.processor(), self.model_chat_template_version, self.config)	171	+ return LlavaNextInputBuilder(self.processor(), self.model_chat_template_version, self.config)

Mexamples/atb_models/atb_llm/models/minicpm_qwen2_v2/router_minicpm_qwen2_v2.py+17-19

	@@ -25,35 +25,35 @@ from ..base.model_utils import safe_get_tokenizer_from_pretrained
25		25
26	@dataclass	26	@dataclass
27	class Minicpmqwen2v2Router(BaseRouter):	27	class Minicpmqwen2v2Router(BaseRouter):
		28	+ is_multimodal: bool = True
28		29
29	def __post_init__(self):	30	def __post_init__(self):
30	super().__post_init__()	31	super().__post_init__()
31	- self.processor = safe_from_pretrained(AutoProcessor, self.config.model_name_or_path,	32	+ self.processor = safe_from_pretrained(
32	- trust_remote_code=self.trust_remote_code)	33	+ AutoProcessor, self.config.model_name_or_path, trust_remote_code=self.trust_remote_code
33	-	34	+ )
		35	+
34	def tokenize(self, inputs, **kwargs):	36	def tokenize(self, inputs, **kwargs):
35	- return inner_tokenize(	37	+ return inner_tokenize(inputs=inputs, config=self.config, processor=self.processor, **kwargs)
36	- inputs=inputs,
37	- config=self.config,
38	- processor=self.processor,
39	- **kwargs)
40		38
41	def check_config_minicpmqwen2v2(self, config):	39	def check_config_minicpmqwen2v2(self, config):
42	super().check_config(config)	40	super().check_config(config)
43	- attribute_ranges = {'mm_hidden_size': (1, 2147483647), 'num_key_value_heads': (1, 2147483647)}	41	+ attribute_ranges = {"mm_hidden_size": (1, 2147483647), "num_key_value_heads": (1, 2147483647)}
44	for attr, (min_val, max_val) in attribute_ranges.items():	42	for attr, (min_val, max_val) in attribute_ranges.items():
45	if not hasattr(config, attr) or getattr(config, attr) is None:	43	if not hasattr(config, attr) or getattr(config, attr) is None:
46	continue	44	continue
47	value = getattr(config, attr)	45	value = getattr(config, attr)
48	if value < min_val or value > max_val:	46	if value < min_val or value > max_val:
49	- logger.error(f"`self._config.{attr}` must be between {min_val} and {max_val}.",	47	+ logger.error(
50	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	48	+ f"`self._config.{attr}` must be between {min_val} and {max_val}.",
		49	+ ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE,
		50	+ )
51	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")	51	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")
52		52
53	def get_config(self):	53	def get_config(self):
54	config = Minicpmqwen2v2Config.from_dict(self.config_dict)	54	config = Minicpmqwen2v2Config.from_dict(self.config_dict)
55	config.model_name_or_path = self.model_name_or_path	55	config.model_name_or_path = self.model_name_or_path
56	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	56	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
57	if self.max_position_embeddings:	57	if self.max_position_embeddings:
58	config.max_position_embeddings = self.max_position_embeddings	58	config.max_position_embeddings = self.max_position_embeddings
59	self.check_config_minicpmqwen2v2(config)	59	self.check_config_minicpmqwen2v2(config)
	@@ -67,18 +67,16 @@ class Minicpmqwen2v2Router(BaseRouter):
67	padding_side="left",	67	padding_side="left",
68	truncation_side="left",	68	truncation_side="left",
69	trust_remote_code=self.trust_remote_code,	69	trust_remote_code=self.trust_remote_code,
70	- use_fast=use_fast)	70	+ use_fast=use_fast,
		71	+ )
71		72
72	def get_generation_config(self):	73	def get_generation_config(self):
73	generation_config = super().get_generation_config()	74	generation_config = super().get_generation_config()
74	return generation_config	75	return generation_config
75		76
76	def get_input_builder(self, **kwargs):	77	def get_input_builder(self, **kwargs):
77	- return Minicpmqwen2v2InputBuilder(	78	+ return Minicpmqwen2v2InputBuilder(config=self.config, processor=self.processor, **kwargs)
78	- config=self.config,	79	+
79	- processor=self.processor,
80	- **kwargs)
81	-
82	def generate_position_ids(self, input_ids):	80	def generate_position_ids(self, input_ids):
83	position_ids = np.arange(len(input_ids), dtype=np.int64)	81	position_ids = np.arange(len(input_ids), dtype=np.int64)
84	- return position_ids	82	+ return position_ids

Mexamples/atb_models/atb_llm/models/mllama/router_mllama.py+2-1

	@@ -14,7 +14,8 @@ from ..base.router import BaseRouter
14		14
15	@dataclass	15	@dataclass
16	class MllamaRouter(BaseRouter):	16	class MllamaRouter(BaseRouter):
		17	+ is_multimodal: bool = True
17		18
18	@staticmethod	19	@staticmethod
19	def check_config(config):	20	def check_config(config):
20	- BaseRouter.check_config(config.text_config)	21	+ BaseRouter.check_config(config.text_config)

Mexamples/atb_models/atb_llm/models/qwen2_audio/router_qwen2_audio.py+20-23

	@@ -59,23 +59,21 @@ def get_feature_lens(config: Any, inputs: list, processor: Any):
59	image_or_video_token_id = config.audio_token_index	59	image_or_video_token_id = config.audio_token_index
60	audio = load_feature_by_torchaudio(image_or_video_path)	60	audio = load_feature_by_torchaudio(image_or_video_path)
61		61
62	- prompt_head = '<\|im_start\|>system\nYou are a helpful assistant.<\|im_end\|>\n<\|im_start\|>user\n'	62	+ prompt_head = "<\|im_start\|>system\nYou are a helpful assistant.<\|im_end\|>\n<\|im_start\|>user\n"
63	- audio_head = 'Audio 1: <\|audio_bos\|><\|AUDIO\|><\|audio_eos\|>\n'	63	+ audio_head = "Audio 1: <\|audio_bos\|><\|AUDIO\|><\|audio_eos\|>\n"
64	- prompt_tail = '<\|im_end\|>\n<\|im_start\|>assistant\n'	64	+ prompt_tail = "<\|im_end\|>\n<\|im_start\|>assistant\n"
65	text = prompt_head + audio_head + text + prompt_tail	65	text = prompt_head + audio_head + text + prompt_tail
66	new_inputs = processor(text=text, audios=audio, return_tensors="pt")	66	new_inputs = processor(text=text, audios=audio, return_tensors="pt")
67	feature_attention_mask = new_inputs.feature_attention_mask	67	feature_attention_mask = new_inputs.feature_attention_mask
68		68
69	- _, audio_output_lengths = get_feat_extract_output_lengths(	69	+ _, audio_output_lengths = get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
70	- feature_attention_mask.sum(-1)
71	- )
72	return [audio_output_lengths, image_or_video_token_id, [image_or_video_path], new_inputs]	70	return [audio_output_lengths, image_or_video_token_id, [image_or_video_path], new_inputs]
73		71
74		72
75	def get_conversation_feature_lens(config: Any, inputs: list, processor: Any, text: str):	73	def get_conversation_feature_lens(config: Any, inputs: list, processor: Any, text: str):
76	audios, audio_paths = [], []	74	audios, audio_paths = [], []
77	audio_types = "audio_type"	75	audio_types = "audio_type"
78	- audio_type = audio_types.split('_')[0]	76	+ audio_type = audio_types.split("_")[0]
79	for message in inputs:	77	for message in inputs:
80	if isinstance(message["content"], list):	78	if isinstance(message["content"], list):
81	for ele in message["content"]:	79	for ele in message["content"]:
	@@ -89,9 +87,7 @@ def get_conversation_feature_lens(config: Any, inputs: list, processor: Any, tex
89	new_inputs = processor(text=text, audios=audios, return_tensors="pt")	87	new_inputs = processor(text=text, audios=audios, return_tensors="pt")
90	feature_attention_mask = new_inputs.feature_attention_mask	88	feature_attention_mask = new_inputs.feature_attention_mask
91		89
92	- _, audio_output_lengths = get_feat_extract_output_lengths(	90	+ _, audio_output_lengths = get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
93	- feature_attention_mask.sum(-1)
94	- )
95	return [audio_output_lengths, image_or_video_token_id, audio_paths, new_inputs]	91	return [audio_output_lengths, image_or_video_token_id, audio_paths, new_inputs]
96		92
97		93
	@@ -104,6 +100,7 @@ class Qwen2AudioTextConfig(Qwen2Config):
104		100
105	@dataclass	101	@dataclass
106	class Qwen2audioRouter(BaseRouter):	102	class Qwen2audioRouter(BaseRouter):
		103	+ is_multimodal: bool = True
107		104
108	def __post_init__(self):	105	def __post_init__(self):
109	super().__post_init__()	106	super().__post_init__()
	@@ -116,22 +113,22 @@ class Qwen2audioRouter(BaseRouter):
116	def check_conversation(inputs):	113	def check_conversation(inputs):
117	conversation_state = False	114	conversation_state = False
118	for item in inputs:	115	for item in inputs:
119	- if 'role' in item:	116	+ if "role" in item:
120	conversation_state = True	117	conversation_state = True
121	return conversation_state	118	return conversation_state
122	-	119	+
123	@staticmethod	120	@staticmethod
124	def check_single_audio_file(inputs):	121	def check_single_audio_file(inputs):
125	for item in inputs:	122	for item in inputs:
126	- inputs_new = item['content'] if 'role' in item else inputs	123	+ inputs_new = item["content"] if "role" in item else inputs
127	audio_num = 0	124	audio_num = 0
128	for item in inputs_new:	125	for item in inputs_new:
129	- if 'audio' in item:	126	+ if "audio" in item:
130	audio_num += 1	127	audio_num += 1
131	if audio_num > 1:	128	if audio_num > 1:
132	return False, inputs_new	129	return False, inputs_new
133	return True, inputs_new	130	return True, inputs_new
134	-	131	+
135	@staticmethod	132	@staticmethod
136	def process_shm(image_pixel, shm_name_save_path, dtype=np.float32):	133	def process_shm(image_pixel, shm_name_save_path, dtype=np.float32):
137	shm = create_shm(image_pixel.nbytes, shm_name_save_path)	134	shm = create_shm(image_pixel.nbytes, shm_name_save_path)
	@@ -164,9 +161,9 @@ class Qwen2audioRouter(BaseRouter):
164	new_inputs.feature_attention_mask = torch.tensor(new_inputs.feature_attention_mask, dtype=torch.int64)	161	new_inputs.feature_attention_mask = torch.tensor(new_inputs.feature_attention_mask, dtype=torch.int64)
165	for feature_len in feature_lens:	162	for feature_len in feature_lens:
166	pad_token_length += feature_len - 1	163	pad_token_length += feature_len - 1
167	- new_prompt_token = torch.full((pad_token_length, ), pad_token_id, dtype=torch.int64)	164	+ new_prompt_token = torch.full((pad_token_length,), pad_token_id, dtype=torch.int64)
168		165
169	- shm_name_save_path = kwargs.get('shm_name_save_path', None)	166	+ shm_name_save_path = kwargs.get("shm_name_save_path", None)
170	shm_name_list, shape_value_list, audio_pixels = [], [], []	167	shm_name_list, shape_value_list, audio_pixels = [], [], []
171		168
172	audio_pixels.append(new_inputs.input_features)	169	audio_pixels.append(new_inputs.input_features)
	@@ -183,7 +180,7 @@ class Qwen2audioRouter(BaseRouter):
183	shm_name, shape_value = self.process_shm(audio_pixel, shm_name_save_path, dtype=np.int64)	180	shm_name, shape_value = self.process_shm(audio_pixel, shm_name_save_path, dtype=np.int64)
184	shm_name_list.append(shm_name)	181	shm_name_list.append(shm_name)
185	shape_value_list.append(shape_value)	182	shape_value_list.append(shape_value)
186	-	183	+
187	for i, _ in enumerate(shm_name_list):	184	for i, _ in enumerate(shm_name_list):
188	new_prompt_token[2 * i] = torch.tensor(shm_name_list[i], dtype=torch.int64)	185	new_prompt_token[2 * i] = torch.tensor(shm_name_list[i], dtype=torch.int64)
189	new_prompt_token[2 * i + 1] = torch.tensor(shape_value_list[i], dtype=torch.int64)	186	new_prompt_token[2 * i + 1] = torch.tensor(shape_value_list[i], dtype=torch.int64)
	@@ -195,7 +192,7 @@ class Qwen2audioRouter(BaseRouter):
195	config = Qwen2AudioConfig.from_dict(self.config_dict)	192	config = Qwen2AudioConfig.from_dict(self.config_dict)
196	config.model_name_or_path = self.model_name_or_path	193	config.model_name_or_path = self.model_name_or_path
197	config.text_config = Qwen2AudioTextConfig.from_dict(config.text_config)	194	config.text_config = Qwen2AudioTextConfig.from_dict(config.text_config)
198	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	195	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
199	if self.max_position_embeddings:	196	if self.max_position_embeddings:
200	config.max_position_embeddings = self.max_position_embeddings	197	config.max_position_embeddings = self.max_position_embeddings
201	return config	198	return config
	@@ -216,14 +213,14 @@ class Qwen2audioRouter(BaseRouter):
216	return position_ids	213	return position_ids
217		214
218	def make_context(	215	def make_context(
219	- self,	216	+ self,
220	rank: int,	217	rank: int,
221	- conversation: List[Dict[str, List[Dict]]],	218	+ conversation: List[Dict[str, List[Dict]]],
222	system: str = "You are a helpful assistant.",	219	system: str = "You are a helpful assistant.",
223	- **kwargs):	220	+ **kwargs,
		221	+ ):
224	context_tokens = self.tokenize(conversation)	222	context_tokens = self.tokenize(conversation)
225	return context_tokens	223	return context_tokens
226		224
227	def get_input_builder(self):	225	def get_input_builder(self):
228	return self	226	return self
229	-

Mexamples/atb_models/atb_llm/models/qwen2_vl/router_qwen2_vl.py+47-42

	@@ -20,7 +20,7 @@ from atb_llm.models.qwen2_vl.config_qwen2_vl import Qwen2vlConfig
20	from atb_llm.models.qwen2_vl.data_preprocess_qwen2_vl import fetch_image, fetch_video, process_shared_memory	20	from atb_llm.models.qwen2_vl.data_preprocess_qwen2_vl import fetch_image, fetch_video, process_shared_memory
21	from atb_llm.models.qwen2_vl.input_builder_qwen2_vl import Qwen2vlInputBuilder	21	from atb_llm.models.qwen2_vl.input_builder_qwen2_vl import Qwen2vlInputBuilder
22	from atb_llm.utils.log.error_code import ErrorCode	22	from atb_llm.utils.log.error_code import ErrorCode
23	-from atb_llm.utils.log.logging import logger, print_log	23	+from atb_llm.utils.log.logging import logger
24		24
25	VISION_START_TOKEN_ID = 151652	25	VISION_START_TOKEN_ID = 151652
26	VISION_END_TOKEN_ID = 151653	26	VISION_END_TOKEN_ID = 151653
	@@ -42,6 +42,7 @@ messages_template = [
42		42
43	@dataclass	43	@dataclass
44	class Qwen2vlRouter(BaseRouter):	44	class Qwen2vlRouter(BaseRouter):
		45	+ is_multimodal: bool = True
45	_processor: Any = None	46	_processor: Any = None
46		47
47	@property	48	@property
	@@ -66,7 +67,7 @@ class Qwen2vlRouter(BaseRouter):
66		67
67	vision_info_list = []	68	vision_info_list = []
68	message_list = []	69	message_list = []
69	- shm_name_save_path = kwargs.get('shm_name_save_path', None)	70	+ shm_name_save_path = kwargs.get("shm_name_save_path", None)
70		71
71	for single_input in inputs:	72	for single_input in inputs:
72	if single_input.get(IMAGE, None):	73	if single_input.get(IMAGE, None):
	@@ -79,17 +80,17 @@ class Qwen2vlRouter(BaseRouter):
79	images_inputs, feature_lens = fetch_image(self.image_processor, single_input)	80	images_inputs, feature_lens = fetch_image(self.image_processor, single_input)
80		81
81	shared_memory_result = process_shared_memory(	82	shared_memory_result = process_shared_memory(
82	- images_inputs.pixel_values,	83	+ images_inputs.pixel_values, shm_name_save_path, images_inputs.image_grid_thw
83	- shm_name_save_path,	84	+ )
84	- images_inputs.image_grid_thw	85	+ vision_info_list.append(
		86	+ [
		87	+ shared_memory_result["pixel_values_shm_name"],
		88	+ shared_memory_result["pixel_values_shape_value"],
		89	+ shared_memory_result["thw_value"],
		90	+ feature_lens,
		91	+ image_token_id,
		92	+ ]
85	)	93	)
86	- vision_info_list.append([
87	- shared_memory_result['pixel_values_shm_name'],
88	- shared_memory_result['pixel_values_shape_value'],
89	- shared_memory_result['thw_value'],
90	- feature_lens,
91	- image_token_id
92	- ])
93		94
94	elif single_input.get(VIDEO, None):	95	elif single_input.get(VIDEO, None):
95	message_list.append(single_input)	96	message_list.append(single_input)
	@@ -100,43 +101,47 @@ class Qwen2vlRouter(BaseRouter):
100	# 默认 fps 为 2	101	# 默认 fps 为 2
101	video_single_message = {"type": "video", "video": single_input[VIDEO], "fps": 2.0}	102	video_single_message = {"type": "video", "video": single_input[VIDEO], "fps": 2.0}
102	video_inputs, feature_lens, second_per_grid_t = fetch_video(	103	video_inputs, feature_lens, second_per_grid_t = fetch_video(
103	- self.image_processor,	104	+ self.image_processor, video_single_message, IMAGE_FACTOR
104	- video_single_message,
105	- IMAGE_FACTOR
106	)	105	)
107	shared_memory_result = process_shared_memory(	106	shared_memory_result = process_shared_memory(
108	- video_inputs.pixel_values_videos,	107	+ video_inputs.pixel_values_videos, shm_name_save_path, video_inputs.video_grid_thw, second_per_grid_t
109	- shm_name_save_path,	108	+ )
110	- video_inputs.video_grid_thw,	109	+ vision_info_list.append(
111	- second_per_grid_t	110	+ [
		111	+ shared_memory_result["pixel_values_shm_name"],
		112	+ shared_memory_result["pixel_values_shape_value"],
		113	+ shared_memory_result["thw_value"],
		114	+ shared_memory_result["second_per_grid_t_shm_name"],
		115	+ shared_memory_result["second_per_grid_t_shape_value"],
		116	+ feature_lens,
		117	+ video_token_id,
		118	+ ]
112	)	119	)
113	- vision_info_list.append([
114	- shared_memory_result['pixel_values_shm_name'],
115	- shared_memory_result['pixel_values_shape_value'],
116	- shared_memory_result['thw_value'],
117	- shared_memory_result['second_per_grid_t_shm_name'],
118	- shared_memory_result['second_per_grid_t_shape_value'],
119	- feature_lens,
120	- video_token_id
121	- ])
122		120
123	elif single_input.get(TEXT, None):	121	elif single_input.get(TEXT, None):
124	message_list.append(single_input)	122	message_list.append(single_input)
125	else:	123	else:
126	- logger.error("The input field currently only supports 'image', 'video' or 'text'.",	124	+ logger.error(
127	- ErrorCode.ATB_MODELS_PARAM_INVALID)	125	+ "The input field currently only supports 'image', 'video' or 'text'.",
		126	+ ErrorCode.ATB_MODELS_PARAM_INVALID,
		127	+ )
128	raise TypeError("The input field currently only supports 'image', 'video' or 'text'.")	128	raise TypeError("The input field currently only supports 'image', 'video' or 'text'.")
129	return self.process_token(vision_info_list, message_list)	129	return self.process_token(vision_info_list, message_list)
130		130
131	def get_input_builder(self):	131	def get_input_builder(self):
132	if hasattr(self.config, "max_position_embeddings") and self.config.max_position_embeddings:	132	if hasattr(self.config, "max_position_embeddings") and self.config.max_position_embeddings:
133	- return Qwen2vlInputBuilder(self.tokenizer, self.image_processor, self.processor, self.config,	133	+ return Qwen2vlInputBuilder(
134	- max_length=self.config.max_position_embeddings)	134	+ self.tokenizer,
		135	+ self.image_processor,
		136	+ self.processor,
		137	+ self.config,
		138	+ max_length=self.config.max_position_embeddings,
		139	+ )
135	return Qwen2vlInputBuilder(self.tokenizer, self.image_processor, self.processor, self.config)	140	return Qwen2vlInputBuilder(self.tokenizer, self.image_processor, self.processor, self.config)
136		141
137	def get_config(self):	142	def get_config(self):
138	config = Qwen2vlConfig.from_dict(self.config_dict)	143	config = Qwen2vlConfig.from_dict(self.config_dict)
139	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	144	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
140	if self.max_position_embeddings:	145	if self.max_position_embeddings:
141	config.max_position_embeddings = self.max_position_embeddings	146	config.max_position_embeddings = self.max_position_embeddings
142	config.model_name_or_path = self.model_name_or_path	147	config.model_name_or_path = self.model_name_or_path
	@@ -155,16 +160,18 @@ class Qwen2vlRouter(BaseRouter):
155	def check_config_qwen2_vl(self, config):	160	def check_config_qwen2_vl(self, config):
156	super().check_config(config)	161	super().check_config(config)
157	attribute_ranges = {	162	attribute_ranges = {
158	- 'mm_hidden_size': (1, 2147483647),	163	+ "mm_hidden_size": (1, 2147483647),
159	- 'num_key_value_heads': (1, 2147483647),	164	+ "num_key_value_heads": (1, 2147483647),
160	}	165	}
161	for attr, (min_val, max_val) in attribute_ranges.items():	166	for attr, (min_val, max_val) in attribute_ranges.items():
162	if not hasattr(config, attr) or getattr(config, attr) is None:	167	if not hasattr(config, attr) or getattr(config, attr) is None:
163	continue	168	continue
164	value = getattr(config, attr)	169	value = getattr(config, attr)
165	if value < min_val or value > max_val:	170	if value < min_val or value > max_val:
166	- logger.error(f"`self._config.{attr}` must be between {min_val} and {max_val}.",	171	+ logger.error(
167	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	172	+ f"`self._config.{attr}` must be between {min_val} and {max_val}.",
		173	+ ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE,
		174	+ )
168	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")	175	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")
169		176
170	def process_token(self, vision_info_list, message_list):	177	def process_token(self, vision_info_list, message_list):
	@@ -172,9 +179,7 @@ class Qwen2vlRouter(BaseRouter):
172	messages[0]["content"] = message_list	179	messages[0]["content"] = message_list
173	prompt = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)	180	prompt = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
174	if not prompt:	181	if not prompt:
175	- prompt = self.processor.apply_chat_template(	182	+ prompt = self.processor.apply_chat_template(message_list, tokenize=False, add_generation_prompt=True)
176	- message_list, tokenize=False, add_generation_prompt=True
177	- )
178	input_ids = self.tokenizer(prompt, return_tensors="pt")["input_ids"].flatten()	183	input_ids = self.tokenizer(prompt, return_tensors="pt")["input_ids"].flatten()
179	new_input_ids = input_ids	184	new_input_ids = input_ids
180	vision_start_token_id = getattr(self.config, "vision_start_token_id", VISION_START_TOKEN_ID)	185	vision_start_token_id = getattr(self.config, "vision_start_token_id", VISION_START_TOKEN_ID)
	@@ -189,7 +194,7 @@ class Qwen2vlRouter(BaseRouter):
189	pre = 0	194	pre = 0
190	for i in range(0, vision_num, 1):	195	for i in range(0, vision_num, 1):
191	feature_lens = vision_info_list[i][-2]	196	feature_lens = vision_info_list[i][-2]
192	- text_token = input_ids[pre: bos_pos[i]]	197	+ text_token = input_ids[pre : bos_pos[i]]
193	pre = eos_pos[i] + 1	198	pre = eos_pos[i] + 1
194	if vision_info_list[i][-1] == image_token_id:	199	if vision_info_list[i][-1] == image_token_id:
195	vision_pad_token = torch.cat(	200	vision_pad_token = torch.cat(
	@@ -227,4 +232,4 @@ class Qwen2vlRouter(BaseRouter):
227		232
228	if expand_token_ids:	233	if expand_token_ids:
229	new_input_ids = torch.cat(expand_token_ids)	234	new_input_ids = torch.cat(expand_token_ids)
230	- return new_input_ids	235	+ return new_input_ids

Mexamples/atb_models/atb_llm/models/qwen3_vl/router_qwen3_vl.py+7-16

	@@ -8,7 +8,6 @@
8	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	8	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
9	# See the Mulan PSL v2 for more details.	9	# See the Mulan PSL v2 for more details.
10		10
11	-import os
12	import copy	11	import copy
13	from typing import Dict, List, Any	12	from typing import Dict, List, Any
14	from dataclasses import dataclass	13	from dataclasses import dataclass
	@@ -22,8 +21,6 @@ from atb_llm.utils.shm_utils import process_shared_memory
22	from atb_llm.utils.multimodal_utils import safe_open_image, check_video_path	21	from atb_llm.utils.multimodal_utils import safe_open_image, check_video_path
23		22
24	from .input_builder_qwen3_vl import Qwen3vlInputBuilder	23	from .input_builder_qwen3_vl import Qwen3vlInputBuilder
25	-from ...utils.log import logger
26	-from ...utils.log.error_code import ErrorCode
27		24
28		25
29	IMAGE = "image"	26	IMAGE = "image"
	@@ -41,6 +38,7 @@ messages_template = [
41		38
42	@dataclass	39	@dataclass
43	class Qwen3vlRouter(BaseRouter):	40	class Qwen3vlRouter(BaseRouter):
		41	+ is_multimodal: bool = True
44	_processor: Any = None	42	_processor: Any = None
45		43
46	def __post_init__(self):	44	def __post_init__(self):
	@@ -53,11 +51,11 @@ class Qwen3vlRouter(BaseRouter):
53	if not hasattr(self, "_processor") or self._processor is None:	51	if not hasattr(self, "_processor") or self._processor is None:
54	self._processor = self.get_processor()	52	self._processor = self.get_processor()
55	return self._processor	53	return self._processor
56	-	54	+
57	def tokenize(self, inputs: List[Dict], **kwargs):	55	def tokenize(self, inputs: List[Dict], **kwargs):
58	text = ""	56	text = ""
59	message_list = []	57	message_list = []
60	- shm_name_save_path = kwargs.get('shm_name_save_path', None)	58	+ shm_name_save_path = kwargs.get("shm_name_save_path", None)
61	for item in inputs:	59	for item in inputs:
62	if item.get(IMAGE, None):	60	if item.get(IMAGE, None):
63	img_fname = item[IMAGE]	61	img_fname = item[IMAGE]
	@@ -72,18 +70,11 @@ class Qwen3vlRouter(BaseRouter):
72	shm_name_save_path = self.input_builder.get_shm_name_save_path(video_path)	70	shm_name_save_path = self.input_builder.get_shm_name_save_path(video_path)
73	if item.get(TEXT, None):	71	if item.get(TEXT, None):
74	text = item[TEXT]	72	text = item[TEXT]
75	- message_list.append({	73	+ message_list.append({TYPE: TEXT, TEXT: text})
76	- TYPE: TEXT,
77	- TEXT: text
78	- })
79	messages = copy.deepcopy(messages_template)	74	messages = copy.deepcopy(messages_template)
80	messages[0]["content"] = message_list	75	messages[0]["content"] = message_list
81	inputs = self.processor.apply_chat_template(	76	inputs = self.processor.apply_chat_template(
82	- messages,	77	+ messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
83	- tokenize=True,
84	- add_generation_prompt=True,
85	- return_dict=True,
86	- return_tensors='pt'
87	)	78	)
88	shm_info = process_shared_memory(inputs, shm_name_save_path)	79	shm_info = process_shared_memory(inputs, shm_name_save_path)
89	input_ids = self.input_builder.update_token_id(inputs, shm_info)	80	input_ids = self.input_builder.update_token_id(inputs, shm_info)
	@@ -91,7 +82,7 @@ class Qwen3vlRouter(BaseRouter):
91		82
92	def get_input_builder(self):	83	def get_input_builder(self):
93	return Qwen3vlInputBuilder(self.tokenizer, self.config, self.processor)	84	return Qwen3vlInputBuilder(self.tokenizer, self.config, self.processor)
94	-	85	+
95	def get_config(self):	86	def get_config(self):
96	config_cls = self.get_config_cls()	87	config_cls = self.get_config_cls()
97	config = config_cls.from_dict(self.config_dict)	88	config = config_cls.from_dict(self.config_dict)
	@@ -105,7 +96,7 @@ class Qwen3vlRouter(BaseRouter):
105	raise NotImplementedError(err_msg)	96	raise NotImplementedError(err_msg)
106	super().check_config(config)	97	super().check_config(config)
107	return config	98	return config
108	-	99	+
109	def get_tokenizer(self):	100	def get_tokenizer(self):
110	return self.processor.tokenizer	101	return self.processor.tokenizer
111		102

Mexamples/atb_models/atb_llm/models/vita/router_vita.py+70-66

	@@ -59,30 +59,28 @@ def from_list_format(list_format: List[Dict]):
59	if "text" in ele:	59	if "text" in ele:
60	text = ele["text"]	60	text = ele["text"]
61	elif "image_url" in ele or "image" in ele:	61	elif "image_url" in ele or "image" in ele:
62	- image = ele.get('image') or ele.get('image_url')	62	+ image = ele.get("image") or ele.get("image_url")
63	elif "video_url" in ele or "video" in ele:	63	elif "video_url" in ele or "video" in ele:
64	- video = ele.get('video') or ele.get('video_url')	64	+ video = ele.get("video") or ele.get("video_url")
65	elif "audio_url" in ele or "audio" in ele:	65	elif "audio_url" in ele or "audio" in ele:
66	- audio = ele.get("audio") or ele.get('audio_url')	66	+ audio = ele.get("audio") or ele.get("audio_url")
67	else:	67	else:
68	- logger.error("Unsupported element: " + str(ele) + ".",	68	+ logger.error("Unsupported element: " + str(ele) + ".", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
69	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
70	raise KeyError("Unsupported element: " + str(ele) + ".")	69	raise KeyError("Unsupported element: " + str(ele) + ".")
71	return InputAttrs(text, image, video, audio)	70	return InputAttrs(text, image, video, audio)
72		71
73		72
74	@dataclass	73	@dataclass
75	class VitaRouter(BaseRouter):	74	class VitaRouter(BaseRouter):
		75	+ is_multimodal: bool = True
76		76
77	def process_images(self, images, model_cfg, image_processor):	77	def process_images(self, images, model_cfg, image_processor):
78	image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)	78	image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
79	-	79	+
80	new_images = []	80	new_images = []
81	if image_aspect_ratio == "pad":	81	if image_aspect_ratio == "pad":
82	for image in images:	82	for image in images:
83	- image = self.expand2square(	83	+ image = self.expand2square(image, tuple(int(x * 255) for x in image_processor.image_mean))
84	- image, tuple(int(x * 255) for x in image_processor.image_mean)
85	- )
86	image = image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]	84	image = image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
87	new_images.append(image)	85	new_images.append(image)
88	else:	86	else:
	@@ -90,7 +88,7 @@ class VitaRouter(BaseRouter):
90	if all(x.shape == new_images[0].shape for x in new_images):	88	if all(x.shape == new_images[0].shape for x in new_images):
91	new_images = torch.stack(new_images, dim=0)	89	new_images = torch.stack(new_images, dim=0)
92	return new_images	90	return new_images
93	-	91	+
94	def get_audio_processor(self):	92	def get_audio_processor(self):
95	audio_path = os.path.join(self.config.model_name_or_path, self.config.mm_audio_encoder)	93	audio_path = os.path.join(self.config.model_name_or_path, self.config.mm_audio_encoder)
96	with open(os.path.join(audio_path, "train.yaml"), "r") as fin:	94	with open(os.path.join(audio_path, "train.yaml"), "r") as fin:
	@@ -109,25 +107,21 @@ class VitaRouter(BaseRouter):
109	if "text" in ele:	107	if "text" in ele:
110	text = ele.get("text")	108	text = ele.get("text")
111	new_q = qs + text	109	new_q = qs + text
112	- qs = ''	110	+ qs = ""
113	conv.append_message(conv.roles[0], new_q)	111	conv.append_message(conv.roles[0], new_q)
114	elif "image_url" in ele or img in ele:	112	elif "image_url" in ele or img in ele:
115	image_path = ele.get(img) or ele.get("image_url")	113	image_path = ele.get(img) or ele.get("image_url")
116	image = multimodal_utils.safe_load_multimodal_source(Image.open, image_path)	114	image = multimodal_utils.safe_load_multimodal_source(Image.open, image_path)
117	- if image.mode != 'RGB':	115	+ if image.mode != "RGB":
118	- image = image.convert('RGB')	116	+ image = image.convert("RGB")
119	- image, p_num = dynamic_preprocess(	117	+ image, p_num = dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=True)
120	- image, min_num=1, max_num=12, image_size=448, use_thumbnail=True	118	+ image_tensor = self.process_images(image, self.config, image_processor).half()
121	- )
122	- image_tensor = self.process_images(image,
123	- self.config,
124	- image_processor).half()
125	qs += DEFAULT_IMAGE_TOKEN * p_num[0] + "\n"	119	qs += DEFAULT_IMAGE_TOKEN * p_num[0] + "\n"
126	modality = img	120	modality = img
127	maxlen_image_id = p_num[0] * 255	121	maxlen_image_id = p_num[0] * 255
128	media_tensor.append((image_path, image_tensor, maxlen_image_id))	122	media_tensor.append((image_path, image_tensor, maxlen_image_id))
129	elif "video_url" in ele or vid in ele:	123	elif "video_url" in ele or vid in ele:
130	- video = ele.get(vid) or ele.get('video_url')	124	+ video = ele.get(vid) or ele.get("video_url")
131	video_config = VideoConfig(	125	video_config = VideoConfig(
132	video_path=video,	126	video_path=video,
133	image_processor=image_processor,	127	image_processor=image_processor,
	@@ -142,9 +136,9 @@ class VitaRouter(BaseRouter):
142	maxlen_image_id = slice_len * 255	136	maxlen_image_id = slice_len * 255
143	media_tensor.append((video, image_tensor, maxlen_image_id))	137	media_tensor.append((video, image_tensor, maxlen_image_id))
144	elif "audio_url" in ele or "audio" in ele:	138	elif "audio_url" in ele or "audio" in ele:
145	- audio = ele.get("audio") or ele.get('audio_url')	139	+ audio = ele.get("audio") or ele.get("audio_url")
146	new_q = qs + DEFAULT_AUDIO_TOKEN	140	new_q = qs + DEFAULT_AUDIO_TOKEN
147	- qs = ''	141	+ qs = ""
148	conv.append_message(conv.roles[0], new_q)	142	conv.append_message(conv.roles[0], new_q)
149	audio_process = self.get_audio_processor()	143	audio_process = self.get_audio_processor()
150	audio_input, _ = audio_process.process(os.path.join(audio))	144	audio_input, _ = audio_process.process(os.path.join(audio))
	@@ -154,8 +148,7 @@ class VitaRouter(BaseRouter):
154	max_audio_id = math.ceil((((audio_length - 3) // 2 - 2) // 2 + 1) / 2) - 1	148	max_audio_id = math.ceil((((audio_length - 3) // 2 - 2) // 2 + 1) / 2) - 1
155	media_tensor.append((audio, audio_input, max_audio_id))	149	media_tensor.append((audio, audio_input, max_audio_id))
156	else:	150	else:
157	- logger.error("Unsupported element: " + str(ele) + ".",	151	+ logger.error("Unsupported element: " + str(ele) + ".", ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
158	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
159	raise KeyError("Unsupported element: " + str(ele) + ".")	152	raise KeyError("Unsupported element: " + str(ele) + ".")
160	return conv, media_tensor, modality	153	return conv, media_tensor, modality
161		154
	@@ -169,42 +162,45 @@ class VitaRouter(BaseRouter):
169	shm = create_shm(tensor.nbytes, shm_name_save_path)	162	shm = create_shm(tensor.nbytes, shm_name_save_path)
170	shared_array = np.ndarray(tensor.shape, dtype=tensor.dtype, buffer=shm.buf)	163	shared_array = np.ndarray(tensor.shape, dtype=tensor.dtype, buffer=shm.buf)
171	shared_array[:] = tensor	164	shared_array[:] = tensor
172	- shm_name = encode_shm_name_to_int64(shm.name)	165	+ shm_name = encode_shm_name_to_int64(shm.name)
173	shape_value = encode_shape_to_int64(tensor.shape)	166	shape_value = encode_shape_to_int64(tensor.shape)
174	intermediate_token = torch.tensor([shm_name, shape_value])	167	intermediate_token = torch.tensor([shm_name, shape_value])
175	if i == 0:	168	if i == 0:
176	- new_input_ids = torch.cat([	169	+ new_input_ids = torch.cat(
177	- input_ids[start: media_pos[i] + 1],	170	+ [
178	- intermediate_token,	171	+ input_ids[start : media_pos[i] + 1],
179	- torch.full((length - 2, ), IGNORE_INDEX, dtype=input_ids.dtype),	172	+ intermediate_token,
180	- ])	173	+ torch.full((length - 2,), IGNORE_INDEX, dtype=input_ids.dtype),
181	- else:	174	+ ]
182	- new_input_ids = torch.cat([	175	+ )
183	- new_input_ids,	176	+ else:
184	- input_ids[start: media_pos[i] + 1],	177	+ new_input_ids = torch.cat(
185	- intermediate_token,	178	+ [
186	- torch.full((length - 2, ), IGNORE_INDEX, dtype=input_ids.dtype),	179	+ new_input_ids,
187	- ])	180	+ input_ids[start : media_pos[i] + 1],
		181	+ intermediate_token,
		182	+ torch.full((length - 2,), IGNORE_INDEX, dtype=input_ids.dtype),
		183	+ ]
		184	+ )
188	start = media_pos[i] + 1	185	start = media_pos[i] + 1
189	return start, new_input_ids	186	return start, new_input_ids
190		187
191	-
192	def tokenize(self, inputs, **kwargs):	188	def tokenize(self, inputs, **kwargs):
193	tokenizer = safe_from_pretrained(AutoTokenizer, self.config.model_name_or_path, use_fast=True)	189	tokenizer = safe_from_pretrained(AutoTokenizer, self.config.model_name_or_path, use_fast=True)
194	- image_processor = safe_from_pretrained(CLIPImageProcessor,	190	+ image_processor = safe_from_pretrained(
195	- os.path.join(self.config.model_name_or_path, self.config.mm_vision_tower))	191	+ CLIPImageProcessor, os.path.join(self.config.model_name_or_path, self.config.mm_vision_tower)
196	- if self.config.model_type.lower() == 'vita-qwen2':	192	+ )
		193	+ if self.config.model_type.lower() == "vita-qwen2":
197	self.conv_mode = "qwen2p5_instruct"	194	self.conv_mode = "qwen2p5_instruct"
198	else:	195	else:
199	self.conv_mode = "mixtral_two"	196	self.conv_mode = "mixtral_two"
200	conv = conv_templates[self.conv_mode].copy()	197	conv = conv_templates[self.conv_mode].copy()
201	conv, media_tensor, modality = self.process_media(inputs, conv, image_processor)	198	conv, media_tensor, modality = self.process_media(inputs, conv, image_processor)
202	- conv.append_message(conv.roles[1], None)	199	+ conv.append_message(conv.roles[1], None)
203	prompt = conv.get_prompt(modality)	200	prompt = conv.get_prompt(modality)
204	-	201	+
205	input_ids = (	202	input_ids = (
206	- tokenizer_image_audio_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")	203	+ tokenizer_image_audio_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0)
207	- .unsqueeze(0)
208	).flatten()	204	).flatten()
209	image_pos = torch.where(torch.eq(input_ids, IMAGE_TOKEN_INDEX))[0]	205	image_pos = torch.where(torch.eq(input_ids, IMAGE_TOKEN_INDEX))[0]
210	audio_pos = torch.where(torch.eq(input_ids, AUDIO_TOKEN_INDEX))[0]	206	audio_pos = torch.where(torch.eq(input_ids, AUDIO_TOKEN_INDEX))[0]
	@@ -221,39 +217,46 @@ class VitaRouter(BaseRouter):
221	media_pos = image_final_pos	217	media_pos = image_final_pos
222	if len(media_tensor) != 0:	218	if len(media_tensor) != 0:
		ascend-robot4月21日代码逻辑和结构: `tokenize`方法中`start`变量的使用存在逻辑问题。第218行检查`if start == 0`，但`start`只在`build_input_ids`中被赋值，而`build_input_ids`只在第219行`if len(media_tensor) != 0`条件下调用。这意味着当`media_tensor`为空时，`start`变量未被定义，但代码会执行到第220行，可能导致`NameError`异常。问题类型: 代码逻辑和结构文件路径: `examples/atb_models/atb_llm/models/vita/router_vita.py` 行号: 218 问题代码: if len(media_tensor) != 0: start, new_input_ids = self.build_input_ids(media_tensor, input_ids, media_pos) if start == 0: new_input_ids = input_ids else: new_input_ids = torch.cat( [ new_input_ids, input_ids[start:], ] ) 修改建议: 需要确保`start`变量在所有代码路径中都有定义。建议修改为： ```python if len(media_tensor) != 0: start, new_input_ids = self.build_input_ids(media_tensor, input_ids, media_pos) else: start = 0 new_input_ids = input_ids 然后根据start的值决定是否需要拼接剩余部分 if start > 0: new_input_ids = torch.cat([new_input_ids, input_ids[start:]]) `这样逻辑更清晰，且避免了变量未定义的风险。</pre><div style='margin-top: 4px; font-size: 9px; color: #999;'>---<br>此评论由代码审查工具自动生成</div></div>`
223	start, new_input_ids = self.build_input_ids(media_tensor, input_ids, media_pos)	219	start, new_input_ids = self.build_input_ids(media_tensor, input_ids, media_pos)
		220	+ else:
		221	+ start = 0
		222	+ new_input_ids = input_ids
224	if start == 0:	223	if start == 0:
225	new_input_ids = input_ids	224	new_input_ids = input_ids
226	else:	225	else:
227	- new_input_ids = torch.cat([	226	+ new_input_ids = torch.cat(
228	- new_input_ids,	227	+ [
229	- input_ids[start:],	228	+ new_input_ids,
230	- ])	229	+ input_ids[start:],
		230	+ ]
		231	+ )
231	return new_input_ids	232	return new_input_ids
232		233
233	def check_config_vita(self, config):	234	def check_config_vita(self, config):
234	super().check_config(config)	235	super().check_config(config)
235	attribute_ranges = {	236	attribute_ranges = {
236	- 'mm_hidden_size': (1, 2147483647),	237	+ "mm_hidden_size": (1, 2147483647),
237	- 'num_key_value_heads': (1, 2147483647),	238	+ "num_key_value_heads": (1, 2147483647),
238	}	239	}
239	for attr, (min_val, max_val) in attribute_ranges.items():	240	for attr, (min_val, max_val) in attribute_ranges.items():
240	if not hasattr(config, attr) or getattr(config, attr) is None:	241	if not hasattr(config, attr) or getattr(config, attr) is None:
241	continue	242	continue
242	value = getattr(config, attr)	243	value = getattr(config, attr)
243	if value < min_val or value > max_val:	244	if value < min_val or value > max_val:
244	- logger.error(f"`self._config.{attr}` must be between {min_val} and {max_val}.",	245	+ logger.error(
245	- ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	246	+ f"`self._config.{attr}` must be between {min_val} and {max_val}.",
		247	+ ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE,
		248	+ )
246	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")	249	raise ValueError(f"`self._config.{attr}` must be between {min_val} and {max_val}.")
247		250
248	def get_config(self):	251	def get_config(self):
249	- model_type = self.config_dict.get('model_type').lower()	252	+ model_type = self.config_dict.get("model_type").lower()
250	config_cls = self.get_config_cls()	253	config_cls = self.get_config_cls()
251	- if 'mixtral' in model_type:	254	+ if "mixtral" in model_type:
252	config_cls.model_type = "vita-mixtral"	255	config_cls.model_type = "vita-mixtral"
253	- elif 'qwen' in model_type:	256	+ elif "qwen" in model_type:
254	config_cls.model_type = "vita-Qwen2"	257	config_cls.model_type = "vita-Qwen2"
255	config = config_cls.from_dict(self.config_dict)	258	config = config_cls.from_dict(self.config_dict)
256	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	259	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
257	if self.max_position_embeddings:	260	if self.max_position_embeddings:
258	config.max_position_embeddings = self.max_position_embeddings	261	config.max_position_embeddings = self.max_position_embeddings
259	config.torch_dtype = torch.float16	262	config.torch_dtype = torch.float16
	@@ -263,14 +266,15 @@ class VitaRouter(BaseRouter):
263		266
264	def get_tokenizer(self):	267	def get_tokenizer(self):
265	use_fast = True	268	use_fast = True
266	- return safe_from_pretrained(AutoTokenizer,	269	+ return safe_from_pretrained(
267	- self.tokenizer_path,	270	+ AutoTokenizer,
268	- revision=self.revision,	271	+ self.tokenizer_path,
269	- padding_side="left",	272	+ revision=self.revision,
270	- truncation_side="left",	273	+ padding_side="left",
271	- trust_remote_code=self.trust_remote_code,	274	+ truncation_side="left",
272	- use_fast=use_fast,	275	+ trust_remote_code=self.trust_remote_code,
273	- )	276	+ use_fast=use_fast,
274	-	277	+ )
		278	+
275	def get_input_builder(self):	279	def get_input_builder(self):
276	return VitaInputBuilder(self.tokenizer, self.config)	280	return VitaInputBuilder(self.tokenizer, self.config)

Mexamples/atb_models/atb_llm/models/yivl/router_yivl.py+27-29

	@@ -66,31 +66,31 @@ def parse_service_inputs(inputs: List[Dict]):
66	@dataclass	66	@dataclass
67	class YivlRouter(BaseRouter):	67	class YivlRouter(BaseRouter):
68	model_type = "yivl"	68	model_type = "yivl"
		69	+ is_multimodal: bool = True
69		70
70	def __post_init__(self):	71	def __post_init__(self):
71	super().__post_init__()	72	super().__post_init__()
72	- self.image_processor = DataProcessorYiVl(self.config.vison_tower_path,	73	+ self.image_processor = DataProcessorYiVl(self.config.vision_tower_path, self.trust_remote_code)
73	- self.trust_remote_code)
74		74
75	def check_config_yivl(self, config):	75	def check_config_yivl(self, config):
76	super().check_config(config)	76	super().check_config(config)
77	vision_attribute_ranges = {	77	vision_attribute_ranges = {
78	- 'hidden_size': (1, _INT32_MAX),	78	+ "hidden_size": (1, _INT32_MAX),
79	- 'image_size': (1, _INT32_MAX),	79	+ "image_size": (1, _INT32_MAX),
80	- 'intermediate_size': (1, _INT32_MAX),	80	+ "intermediate_size": (1, _INT32_MAX),
81	- 'num_attention_heads': (1, _INT32_MAX),	81	+ "num_attention_heads": (1, _INT32_MAX),
82	- 'num_channels': (1, _INT32_MAX),	82	+ "num_channels": (1, _INT32_MAX),
83	- 'num_hidden_layers': (1, _INT32_MAX),	83	+ "num_hidden_layers": (1, _INT32_MAX),
84	- 'patch_size': (1, config.vision_config.image_size),	84	+ "patch_size": (1, config.vision_config.image_size),
85	- 'projection_dim': (1, _INT32_MAX)	85	+ "projection_dim": (1, _INT32_MAX),
86	}	86	}
87	check_value_range(vision_attribute_ranges, config.vision_config)	87	check_value_range(vision_attribute_ranges, config.vision_config)
88		88
89	vision_layers = config.vision_config.num_hidden_layers	89	vision_layers = config.vision_config.num_hidden_layers
90	attribute_ranges = {	90	attribute_ranges = {
91	- 'mm_hidden_size': (1, _INT32_MAX),	91	+ "mm_hidden_size": (1, _INT32_MAX),
92	- 'num_key_value_heads': (1, _INT32_MAX),	92	+ "num_key_value_heads": (1, _INT32_MAX),
93	- 'mm_vision_select_layer': (-vision_layers, vision_layers),	93	+ "mm_vision_select_layer": (-vision_layers, vision_layers),
94	}	94	}
95	check_value_range(attribute_ranges, config)	95	check_value_range(attribute_ranges, config)
96		96
	@@ -100,10 +100,9 @@ class YivlRouter(BaseRouter):
100		100
101	if config.mm_vision_tower:	101	if config.mm_vision_tower:
102	from os.path import join	102	from os.path import join
103	- config.vison_tower_path = join(	103	+
104	- self.model_name_or_path, config.mm_vision_tower.replace("./", "")	104	+ config.vision_tower_path = join(self.model_name_or_path, config.mm_vision_tower.replace("./", ""))
105	- )	105	+ config.vision_config = safe_from_pretrained(CLIPVisionConfig, config.vision_tower_path)
106	- config.vision_config = safe_from_pretrained(CLIPVisionConfig, config.vison_tower_path)
107	else:	106	else:
108	msg = "Key 'mm_vision_tower' not found at config"	107	msg = "Key 'mm_vision_tower' not found at config"
109	logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	108	logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
	@@ -114,15 +113,15 @@ class YivlRouter(BaseRouter):
114	if self.max_position_embeddings:	113	if self.max_position_embeddings:
115	config.max_position_embeddings = self.max_position_embeddings	114	config.max_position_embeddings = self.max_position_embeddings
116	config.vocab_size = config.text_config.vocab_size	115	config.vocab_size = config.text_config.vocab_size
117	- setattr(config, 'quantization_config', QuantizationConfig(**{}))	116	+ setattr(config, "quantization_config", QuantizationConfig(**{}))
118	self.check_config_yivl(config)	117	self.check_config_yivl(config)
119	config.model_name_or_path = self.model_name_or_path	118	config.model_name_or_path = self.model_name_or_path
120	- setattr(config, 'img_token_id', IMAGE_TOKEN_INDEX)	119	+ setattr(config, "img_token_id", IMAGE_TOKEN_INDEX)
121	- setattr(config, 'eos_token_id', getattr(self.tokenizer, 'eos_token_id', _EOS_TOKEN_ID))	120	+ setattr(config, "eos_token_id", getattr(self.tokenizer, "eos_token_id", _EOS_TOKEN_ID))
122	- setattr(config, 'pad_token_id', getattr(self.tokenizer, 'pad_token_id', _PAD_TOKEN_ID))	121	+ setattr(config, "pad_token_id", getattr(self.tokenizer, "pad_token_id", _PAD_TOKEN_ID))
123	- setattr(config, 'num_img_patches', (config.vision_config.image_size // config.vision_config.patch_size) ** 2)	122	+ setattr(config, "num_img_patches", (config.vision_config.image_size // config.vision_config.patch_size) ** 2)
124	return config	123	return config
125	-	124	+
126	def get_tokenizer(self):	125	def get_tokenizer(self):
127	use_fast = True	126	use_fast = True
128	tokenizer = safe_get_tokenizer_from_pretrained(	127	tokenizer = safe_get_tokenizer_from_pretrained(
	@@ -131,13 +130,12 @@ class YivlRouter(BaseRouter):
131	padding_side="left",	130	padding_side="left",
132	truncation_side="left",	131	truncation_side="left",
133	trust_remote_code=self.trust_remote_code,	132	trust_remote_code=self.trust_remote_code,
134	- use_fast=use_fast	133	+ use_fast=use_fast,
135	)	134	)
136	tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(SEP_TOKEN)	135	tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(SEP_TOKEN)
137	tokenizer.eos_token = SEP_TOKEN	136	tokenizer.eos_token = SEP_TOKEN
138	return tokenizer	137	return tokenizer
139		138
140	-
141	def get_generation_config(self):	139	def get_generation_config(self):
142	generation_config = super().get_generation_config()	140	generation_config = super().get_generation_config()
143	generation_config["eos_token_id"] = self.tokenizer.convert_tokens_to_ids(SEP_TOKEN)	141	generation_config["eos_token_id"] = self.tokenizer.convert_tokens_to_ids(SEP_TOKEN)
	@@ -147,7 +145,7 @@ class YivlRouter(BaseRouter):
147	img_token_id = self.config.img_token_id	145	img_token_id = self.config.img_token_id
148	pad_token_id = self.config.pad_token_id	146	pad_token_id = self.config.pad_token_id
149	img_patch_num = self.config.num_img_patches	147	img_patch_num = self.config.num_img_patches
150	- shm_name_save_path = kwargs.get('shm_name_save_path', None)	148	+ shm_name_save_path = kwargs.get("shm_name_save_path", None)
151	image_path, question = parse_service_inputs(inputs)	149	image_path, question = parse_service_inputs(inputs)
152	input_ids = tokenize_text(question, self.tokenizer)	150	input_ids = tokenize_text(question, self.tokenizer)
153		151
	@@ -162,9 +160,9 @@ class YivlRouter(BaseRouter):
162	shared_array[:] = image_pixel	160	shared_array[:] = image_pixel
163	shm_name = encode_shm_name_to_int64(shm.name)	161	shm_name = encode_shm_name_to_int64(shm.name)
164	shape_value = encode_shape_to_int64(image_pixel.shape)	162	shape_value = encode_shape_to_int64(image_pixel.shape)
165	-	163	+
166	image_info_ids = [img_token_id, shm_name, shape_value]	164	image_info_ids = [img_token_id, shm_name, shape_value]
167	- if (img_patch_num < len(image_info_ids) + 1):	165	+ if img_patch_num < len(image_info_ids) + 1:
168	msg = f"Image patch num should be less than {len(image_info_ids) + 1}, please check model config."	166	msg = f"Image patch num should be less than {len(image_info_ids) + 1}, please check model config."
169	logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)	167	logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
170	raise ValueError(msg)	168	raise ValueError(msg)
	@@ -174,4 +172,4 @@ class YivlRouter(BaseRouter):
174	image_pos = input_ids.index(img_token_id)	172	image_pos = input_ids.index(img_token_id)
175	new_input_ids = input_ids[:image_pos] + image_info_ids + input_ids[image_pos:]	173	new_input_ids = input_ids[:image_pos] + image_info_ids + input_ids[image_pos:]
176		174
177	- return torch.tensor(new_input_ids)	175	+ return torch.tensor(new_input_ids)

Mexamples/atb_models/atb_llm/runner/tokenizer_wrapper.py+120-56

	@@ -42,13 +42,14 @@ class TokenizerWrapper:
42	def __init__(self, model_name_or_path: str, **kwargs) -> None:	42	def __init__(self, model_name_or_path: str, **kwargs) -> None:
43	model_dict_key = "models_dict"	43	model_dict_key = "models_dict"
44	if model_dict_key in kwargs.keys():	44	if model_dict_key in kwargs.keys():
45	- kwargs[model_dict_key] = None if kwargs[model_dict_key] == '' else json.loads(kwargs[model_dict_key])	45	+ kwargs[model_dict_key] = None if kwargs[model_dict_key] == "" else json.loads(kwargs[model_dict_key])
46	router_ins = get_model(model_name_or_path, **kwargs)	46	router_ins = get_model(model_name_or_path, **kwargs)
47	self.config = router_ins.config	47	self.config = router_ins.config
48	self.tokenizer = router_ins.tokenizer	48	self.tokenizer = router_ins.tokenizer
		49	+ self.is_multimodal = getattr(router_ins, "is_multimodal", False)
49	self.input_builder = router_ins.input_builder	50	self.input_builder = router_ins.input_builder
50	self.postprocessor = router_ins.postprocessor	51	self.postprocessor = router_ins.postprocessor
51	- self.tokenize = router_ins.tokenize # This is only used by multi-modal models	52	+ self.tokenize = router_ins.tokenize # This is only used by multi-modal models
52	self.tool_calls_parser = router_ins.toolscallprocessor	53	self.tool_calls_parser = router_ins.toolscallprocessor
53	self.reasoning_parser = router_ins.reasoning_parser	54	self.reasoning_parser = router_ins.reasoning_parser
54	self.llm_config = router_ins.llm_config	55	self.llm_config = router_ins.llm_config
	@@ -58,19 +59,29 @@ class TokenizerWrapper:
58	self.truncation = "truncation"	59	self.truncation = "truncation"
59	if self.llm_config.llm.pmcc_obfuscation_options.data_obfuscation_ca_dir is not None:	60	if self.llm_config.llm.pmcc_obfuscation_options.data_obfuscation_ca_dir is not None:
60	from ai_asset_obfuscate import data_asset_obfuscation	61	from ai_asset_obfuscate import data_asset_obfuscation
		62	+
61	self.obfuscation_func = data_asset_obfuscation.DataAssetObfuscation(self.config.vocab_size)	63	self.obfuscation_func = data_asset_obfuscation.DataAssetObfuscation(self.config.vocab_size)
62	self.enable_data_obfuscation = True	64	self.enable_data_obfuscation = True
63	- file_name_list = ["kms_ca.pem", "kms_cfs.pem", "kms_cfs.key", "kms_client_key_enc.txt", \	65	+ file_name_list = [
64	- "aiguard_psk", "aiguard_psk_enc.txt"]	66	+ "kms_ca.pem",
		67	+ "kms_cfs.pem",
		68	+ "kms_cfs.key",
		69	+ "kms_client_key_enc.txt",
		70	+ "aiguard_psk",
		71	+ "aiguard_psk_enc.txt",
		72	+ ]
65	file_list = []	73	file_list = []
66	for name in file_name_list:	74	for name in file_name_list:
67	file_path = standardize_path(	75	file_path = standardize_path(
68	- os.path.join(self.llm_config.llm.pmcc_obfuscation_options.data_obfuscation_ca_dir, name))	76	+ os.path.join(self.llm_config.llm.pmcc_obfuscation_options.data_obfuscation_ca_dir, name)
		77	+ )
69	check_file_safety(file_path)	78	check_file_safety(file_path)
70	file_list.append(file_path)	79	file_list.append(file_path)
71		80
72	tls_info = (	81	tls_info = (
73	- file_list[0], file_list[1], file_list[2],	82	+ file_list[0],
		83	+ file_list[1],
		84	+ file_list[2],
74	self.llm_config.llm.pmcc_obfuscation_options.kms_agent_port,	85	self.llm_config.llm.pmcc_obfuscation_options.kms_agent_port,
75	self.llm_config.llm.pmcc_obfuscation_options.data_obfuscation_ca_dir,	86	self.llm_config.llm.pmcc_obfuscation_options.data_obfuscation_ca_dir,
76	file_list[3],	87	file_list[3],
	@@ -82,9 +93,10 @@ class TokenizerWrapper:
82	)	93	)
83	res, msg = self.obfuscation_func.set_seed_safer(tls_info, psk_info)	94	res, msg = self.obfuscation_func.set_seed_safer(tls_info, psk_info)
84	if res != 0:	95	if res != 0:
85	- raise RuntimeError(f"Data obfuscation init failed, please refer to MindIE official document, " \	96	+ raise RuntimeError(
86	- f"error msg: {msg}")	97	+ f"Data obfuscation init failed, please refer to MindIE official document, error msg: {msg}"
87	-	98	+ )
		99	+
88	def encode(self, inputs, **kwargs):	100	def encode(self, inputs, **kwargs):
89	is_chatting = kwargs.pop("is_chatting", False)	101	is_chatting = kwargs.pop("is_chatting", False)
90	if is_chatting:	102	if is_chatting:
	@@ -94,11 +106,12 @@ class TokenizerWrapper:
94	if truncation_method == TruncationSide.Disable:	106	if truncation_method == TruncationSide.Disable:
95	kwargs[self.truncation] = False	107	kwargs[self.truncation] = False
96	else:	108	else:
97	- kwargs[self.truncation] = True	109	+ kwargs[self.truncation] = True
98	- if truncation_method == TruncationSide.RIGHT:	110	+ if truncation_method == TruncationSide.RIGHT:
99	self.tokenizer.truncation_side = "right"	111	self.tokenizer.truncation_side = "right"
100	else:	112	else:
101	self.tokenizer.truncation_side = "left"	113	self.tokenizer.truncation_side = "left"
		114	+ kwargs["split_special_tokens"] = self.is_multimodal
102	token_ids = self.tokenizer(inputs, **kwargs)["input_ids"][0].tolist()	115	token_ids = self.tokenizer(inputs, **kwargs)["input_ids"][0].tolist()
103	if self.enable_data_obfuscation:	116	if self.enable_data_obfuscation:
104	if len(token_ids) > 0 and isinstance(token_ids[0], list):	117	if len(token_ids) > 0 and isinstance(token_ids[0], list):
	@@ -107,10 +120,17 @@ class TokenizerWrapper:
107	token_ids = self.obfuscation_func.data_1d_obf(token_ids)	120	token_ids = self.obfuscation_func.data_1d_obf(token_ids)
108	return token_ids	121	return token_ids
109		122
110	- def decode(self, all_token_ids: list[int], skip_special_tokens: bool, use_tool_calls: bool, is_chat_req: bool,	123	+ def decode(
111	- stream: bool, **kwargs):	124	+ self,
		125	+ all_token_ids: list[int],
		126	+ skip_special_tokens: bool,
		127	+ use_tool_calls: bool,
		128	+ is_chat_req: bool,
		129	+ stream: bool,
		130	+ **kwargs,
		131	+ ):
112	metadata = kwargs.get(METADATA_KEY, {})	132	metadata = kwargs.get(METADATA_KEY, {})
113	- self.tool_calls_parser.tools = metadata.get('tools', None)	133	+ self.tool_calls_parser.tools = metadata.get("tools", None)
114	use_reasoning_parser = self._is_use_reasoning_parser(metadata)	134	use_reasoning_parser = self._is_use_reasoning_parser(metadata)
115	if not stream:	135	if not stream:
116	if use_reasoning_parser and use_tool_calls and is_chat_req:	136	if use_reasoning_parser and use_tool_calls and is_chat_req:
	@@ -121,12 +141,11 @@ class TokenizerWrapper:
121	elif use_reasoning_parser and is_chat_req:	141	elif use_reasoning_parser and is_chat_req:
122	result = self.extract_reasoning_content(all_token_ids, skip_special_tokens)	142	result = self.extract_reasoning_content(all_token_ids, skip_special_tokens)
123	elif use_tool_calls and is_chat_req:	143	elif use_tool_calls and is_chat_req:
124	- result = self.tool_calls_parser.decode(self._tokenizer_decode(	144	+ result = self.tool_calls_parser.decode(
125	- all_token_ids, skip_special_tokens=skip_special_tokens))	145	+ self._tokenizer_decode(all_token_ids, skip_special_tokens=skip_special_tokens)
		146	+ )
126	else:	147	else:
127	- result = {	148	+ result = {CONTENT_KEY: self._tokenizer_decode(all_token_ids, skip_special_tokens=skip_special_tokens)}
128	- CONTENT_KEY: self._tokenizer_decode(all_token_ids, skip_special_tokens=skip_special_tokens)
129	- }
130		149
131	if not self.config.is_reasoning_model:	150	if not self.config.is_reasoning_model:
132	return result	151	return result
	@@ -138,25 +157,29 @@ class TokenizerWrapper:
138	curr_decode_index = kwargs.get("curr_decode_index", -1)	157	curr_decode_index = kwargs.get("curr_decode_index", -1)
139	prev_decode_index = kwargs.get("prev_decode_index", -1)	158	prev_decode_index = kwargs.get("prev_decode_index", -1)
140	curr_and_prev_content = self._tokenizer_decode(	159	curr_and_prev_content = self._tokenizer_decode(
141	- all_token_ids[prev_decode_index:], skip_special_tokens=skip_special_tokens)	160	+ all_token_ids[prev_decode_index:], skip_special_tokens=skip_special_tokens
		161	+ )
142	pre_text = self._tokenizer_decode(	162	pre_text = self._tokenizer_decode(
143	- all_token_ids[prev_decode_index:curr_decode_index], skip_special_tokens=skip_special_tokens)	163	+ all_token_ids[prev_decode_index:curr_decode_index], skip_special_tokens=skip_special_tokens
		164	+ )
144	# No new content is added or the characters are incomplete.	165	# No new content is added or the characters are incomplete.
145	if len(curr_and_prev_content) <= len(pre_text) or curr_and_prev_content.endswith("�"):	166	if len(curr_and_prev_content) <= len(pre_text) or curr_and_prev_content.endswith("�"):
146	if not metadata.get("req_end_flag", False):	167	if not metadata.get("req_end_flag", False):
147	if use_tool_calls:	168	if use_tool_calls:
148	return {"update_index": False, METADATA_KEY: metadata}	169	return {"update_index": False, METADATA_KEY: metadata}
149	return {"update_index": False}	170	return {"update_index": False}
150	- delta_text = curr_and_prev_content[len(pre_text):]	171	+ delta_text = curr_and_prev_content[len(pre_text) :]
151	if use_reasoning_parser and use_tool_calls and is_chat_req:	172	if use_reasoning_parser and use_tool_calls and is_chat_req:
152	- return self.get_combined_stream_result(all_token_ids, prev_decode_index, curr_decode_index,	173	+ return self.get_combined_stream_result(
153	- skip_special_tokens, delta_text, metadata)	174	+ all_token_ids, prev_decode_index, curr_decode_index, skip_special_tokens, delta_text, metadata
		175	+ )
154		176
155	elif use_reasoning_parser and is_chat_req:	177	elif use_reasoning_parser and is_chat_req:
156	return self.extract_reasoning_content_streaming(all_token_ids, curr_decode_index, skip_special_tokens)	178	return self.extract_reasoning_content_streaming(all_token_ids, curr_decode_index, skip_special_tokens)
157	elif use_tool_calls and is_chat_req:	179	elif use_tool_calls and is_chat_req:
158	- return self.extract_tool_calls_streaming(all_token_ids, prev_decode_index, curr_decode_index,	180	+ return self.extract_tool_calls_streaming(
159	- skip_special_tokens, delta_text, metadata)	181	+ all_token_ids, prev_decode_index, curr_decode_index, skip_special_tokens, delta_text, metadata
		182	+ )
160	else:	183	else:
161	return {CONTENT_KEY: delta_text}	184	return {CONTENT_KEY: delta_text}
162		185
	@@ -164,40 +187,65 @@ class TokenizerWrapper:
164	reasoning_content_token_ids, content_token_ids = self.reasoning_parser.single_process_reasoning(all_token_ids)	187	reasoning_content_token_ids, content_token_ids = self.reasoning_parser.single_process_reasoning(all_token_ids)
165	return {	188	return {
166	REASON_CONTENT_KEY: self._tokenizer_decode(	189	REASON_CONTENT_KEY: self._tokenizer_decode(
167	- reasoning_content_token_ids, skip_special_tokens=skip_special_tokens),	190	+ reasoning_content_token_ids, skip_special_tokens=skip_special_tokens
168	- CONTENT_KEY: self._tokenizer_decode(content_token_ids, skip_special_tokens=skip_special_tokens)	191	+ ),
		192	+ CONTENT_KEY: self._tokenizer_decode(content_token_ids, skip_special_tokens=skip_special_tokens),
169	}	193	}
170		194
171	- def extract_reasoning_content_streaming(self, all_token_ids: list[int], curr_decode_index: int,	195	+ def extract_reasoning_content_streaming(
172	- skip_special_tokens: bool):	196	+ self, all_token_ids: list[int], curr_decode_index: int, skip_special_tokens: bool
		197	+ ):
173	reasoning_content_token_ids, content_token_ids = self.reasoning_parser.stream_process_reasoning(	198	reasoning_content_token_ids, content_token_ids = self.reasoning_parser.stream_process_reasoning(
174	- all_token_ids, curr_decode_index)	199	+ all_token_ids, curr_decode_index
		200	+ )
175	reasoning_result = {	201	reasoning_result = {
176	REASON_CONTENT_KEY: self._tokenizer_decode(	202	REASON_CONTENT_KEY: self._tokenizer_decode(
177	- reasoning_content_token_ids, skip_special_tokens=skip_special_tokens),	203	+ reasoning_content_token_ids, skip_special_tokens=skip_special_tokens
178	- CONTENT_KEY: self._tokenizer_decode(content_token_ids, skip_special_tokens=skip_special_tokens)	204	+ ),
		205	+ CONTENT_KEY: self._tokenizer_decode(content_token_ids, skip_special_tokens=skip_special_tokens),
179	}	206	}
180	return {k: v for k, v in reasoning_result.items() if v}	207	return {k: v for k, v in reasoning_result.items() if v}
181		208
182	- def extract_tool_calls_streaming(self, all_token_ids: list[int], prev_decode_index: int, curr_decode_index: int,	209	+ def extract_tool_calls_streaming(
183	- skip_special_tokens: bool, delta_text: str, metadata: Dict):	210	+ self,
		211	+ all_token_ids: list[int],
		212	+ prev_decode_index: int,
		213	+ curr_decode_index: int,
		214	+ skip_special_tokens: bool,
		215	+ delta_text: str,
		216	+ metadata: Dict,
		217	+ ):
184	# Multiple tokenizer processes, compatible with the same request toolscallprocessor with different objects	218	# Multiple tokenizer processes, compatible with the same request toolscallprocessor with different objects
185	- if hasattr(self.tool_calls_parser, 'decode_stream'):	219	+ if hasattr(self.tool_calls_parser, "decode_stream"):
186	self.tool_calls_parser.current_tool_name_sent = metadata.get("current_tool_name_sent")	220	self.tool_calls_parser.current_tool_name_sent = metadata.get("current_tool_name_sent")
187	self.tool_calls_parser.current_tool_arguments_sent = metadata.get("current_tool_arguments_sent")	221	self.tool_calls_parser.current_tool_arguments_sent = metadata.get("current_tool_arguments_sent")
188	self.tool_calls_parser.current_tool_id = metadata.get("current_tool_id")	222	self.tool_calls_parser.current_tool_id = metadata.get("current_tool_id")
189	- result = self.tool_calls_parser.decode_stream(all_token_ids, prev_decode_index, curr_decode_index,	223	+ result = self.tool_calls_parser.decode_stream(
190	- skip_special_tokens, delta_text)	224	+ all_token_ids, prev_decode_index, curr_decode_index, skip_special_tokens, delta_text
191	- result.update({"metadata": {"current_tool_name_sent": self.tool_calls_parser.current_tool_name_sent,	225	+ )
192	- "current_tool_arguments_sent": self.tool_calls_parser.current_tool_arguments_sent,	226	+ result.update(
193	- "current_tool_id": self.tool_calls_parser.current_tool_id}})	227	+ {
		228	+ "metadata": {
		229	+ "current_tool_name_sent": self.tool_calls_parser.current_tool_name_sent,
		230	+ "current_tool_arguments_sent": self.tool_calls_parser.current_tool_arguments_sent,
		231	+ "current_tool_id": self.tool_calls_parser.current_tool_id,
		232	+ }
		233	+ }
		234	+ )
194	else:	235	else:
195	logger.warning("Streaming function call parsing is not supported by the current model.")	236	logger.warning("Streaming function call parsing is not supported by the current model.")
196	result = {CONTENT_KEY: delta_text}	237	result = {CONTENT_KEY: delta_text}
197	return result	238	return result
198		239
199	- def get_combined_stream_result(self, all_token_ids: list[int], prev_decode_index: int, curr_decode_index: int,	240	+ def get_combined_stream_result(
200	- skip_special_tokens: bool, delta_text: str, metadata: Dict):	241	+ self,
		242	+ all_token_ids: list[int],
		243	+ prev_decode_index: int,
		244	+ curr_decode_index: int,
		245	+ skip_special_tokens: bool,
		246	+ delta_text: str,
		247	+ metadata: Dict,
		248	+ ):
201	"""	249	"""
202	Preferentially parse the thought chain.	250	Preferentially parse the thought chain.
203	When the end tag of the thought chain is encountered, parse the function call.	251	When the end tag of the thought chain is encountered, parse the function call.
	@@ -208,31 +256,47 @@ class TokenizerWrapper:
208	return self.extract_reasoning_content_streaming(all_token_ids, curr_decode_index, skip_special_tokens)	256	return self.extract_reasoning_content_streaming(all_token_ids, curr_decode_index, skip_special_tokens)
209	elif full_has_over_think and delta_has_over_think: # New content includes thought chain end tag	257	elif full_has_over_think and delta_has_over_think: # New content includes thought chain end tag
210	reasoning_content_token_ids, content_token_ids = self.reasoning_parser.stream_process_reasoning(	258	reasoning_content_token_ids, content_token_ids = self.reasoning_parser.stream_process_reasoning(
211	- all_token_ids, curr_decode_index)	259	+ all_token_ids, curr_decode_index
		260	+ )
212	if not content_token_ids:	261	if not content_token_ids:
213	- return {REASON_CONTENT_KEY: self._tokenizer_decode(reasoning_content_token_ids, \	262	+ return (
214	- skip_special_tokens=skip_special_tokens)} if reasoning_content_token_ids else {}	263	+ {
		264	+ REASON_CONTENT_KEY: self._tokenizer_decode(
		265	+ reasoning_content_token_ids, skip_special_tokens=skip_special_tokens
		266	+ )
		267	+ }
		268	+ if reasoning_content_token_ids
		269	+ else {}
		270	+ )
215	# case : New content = reasoning content + </think> + tool_call content	271	# case : New content = reasoning content + </think> + tool_call content
216	prev_decode_index = curr_decode_index	272	prev_decode_index = curr_decode_index
217	curr_decode_index = len(all_token_ids) - len(content_token_ids)	273	curr_decode_index = len(all_token_ids) - len(content_token_ids)
218	# update delta_text	274	# update delta_text
219	curr_and_prev_content = self._tokenizer_decode(	275	curr_and_prev_content = self._tokenizer_decode(
220	- all_token_ids[prev_decode_index:], skip_special_tokens=skip_special_tokens)	276	+ all_token_ids[prev_decode_index:], skip_special_tokens=skip_special_tokens
		277	+ )
221	pre_text = self._tokenizer_decode(	278	pre_text = self._tokenizer_decode(
222	- all_token_ids[prev_decode_index:curr_decode_index], skip_special_tokens=skip_special_tokens)	279	+ all_token_ids[prev_decode_index:curr_decode_index], skip_special_tokens=skip_special_tokens
		280	+ )
223	# No new content is added.	281	# No new content is added.
224	if len(curr_and_prev_content) <= len(pre_text):	282	if len(curr_and_prev_content) <= len(pre_text):
225	return {"update_index": False, "metadata": metadata}	283	return {"update_index": False, "metadata": metadata}
226	- delta_text = curr_and_prev_content[len(pre_text):]	284	+ delta_text = curr_and_prev_content[len(pre_text) :]
227	- tool_calls_result = self.extract_tool_calls_streaming(all_token_ids, prev_decode_index, curr_decode_index,	285	+ tool_calls_result = self.extract_tool_calls_streaming(
228	- skip_special_tokens, delta_text, metadata)	286	+ all_token_ids, prev_decode_index, curr_decode_index, skip_special_tokens, delta_text, metadata
		287	+ )
229	tool_calls_result.update(	288	tool_calls_result.update(
230	- {REASON_CONTENT_KEY:	289	+ {
231	- self._tokenizer_decode(reasoning_content_token_ids, skip_special_tokens=skip_special_tokens)})	290	+ REASON_CONTENT_KEY: self._tokenizer_decode(
		291	+ reasoning_content_token_ids, skip_special_tokens=skip_special_tokens
		292	+ )
		293	+ }
		294	+ )
232	return {k: v for k, v in tool_calls_result.items() if v}	295	return {k: v for k, v in tool_calls_result.items() if v}
233	else:	296	else:
234	- return self.extract_tool_calls_streaming(all_token_ids, prev_decode_index, curr_decode_index, \	297	+ return self.extract_tool_calls_streaming(
235	- skip_special_tokens, delta_text, metadata)	298	+ all_token_ids, prev_decode_index, curr_decode_index, skip_special_tokens, delta_text, metadata
		299	+ )
236		300
237	def _tokenizer_decode(self, outputs, **kwargs):	301	def _tokenizer_decode(self, outputs, **kwargs):
238	if self.enable_data_obfuscation:	302	if self.enable_data_obfuscation:
	@@ -256,4 +320,4 @@ class TokenizerWrapper:
256	# Manually disable the model reasoning capability. Request_Switch > Weight_Config_Switch	320	# Manually disable the model reasoning capability. Request_Switch > Weight_Config_Switch
257	if metadata and "req_enable_thinking" in metadata:	321	if metadata and "req_enable_thinking" in metadata:
258	return metadata.get("req_enable_thinking")	322	return metadata.get("req_enable_thinking")
259	- return self.enable_thinking	323	+ return self.enable_thinking

Mexamples/atb_models/tests/run_all_tests.sh+13-1

	@@ -104,7 +104,19 @@ function fn_run_pythontest()
104	! "$line" =~ "internlm2" &&	104	! "$line" =~ "internlm2" &&
105	! "$line" =~ "internlm3" &&	105	! "$line" =~ "internlm3" &&
106	! "$line" =~ "eplb_expert_data_collect" &&	106	! "$line" =~ "eplb_expert_data_collect" &&
107	- ! "$line" =~ "common_op_builder" ]]; then	107	+ ! "$line" =~ "common_op_builder" &&
		108	+ ! "$line" =~ "router_glm4v" &&
		109	+ ! "$line" =~ "router_internlmxcomposer2" &&
		110	+ ! "$line" =~ "router_internvl" &&
		111	+ ! "$line" =~ "router_janus" &&
		112	+ ! "$line" =~ "router_llava" &&
		113	+ ! "$line" =~ "router_llava_next" &&
		114	+ ! "$line" =~ "router_minicpm_qwen2_v2" &&
		115	+ ! "$line" =~ "router_mllama" &&
		116	+ ! "$line" =~ "router_qwen2_audio" &&
		117	+ ! "$line" =~ "router_qwen2_vl" &&
		118	+ ! "$line" =~ "router_vita" &&
		119	+ ! "$line" =~ "router_yivl" ]]; then
108	echo "$line" \| awk -F '"' '{print "examples/atb_models/atb_llm/" $4, $8100 "%", $10100 "%"}' >> result.txt	120	echo "$line" \| awk -F '"' '{print "examples/atb_models/atb_llm/" $4, $8100 "%", $10100 "%"}' >> result.txt
109	fi	121	fi
110	done	122	done

然后根据start的值决定是否需要拼接剩余部分