import importlib import os import random import re from collections import defaultdict from dataclasses import dataclass from copy import deepcopy from functools import partial from typing import Any, Dict, Optional, Tuple, Union import torch import torch.nn.functional as F from diffusers.utils import BaseOutput from transformers.tokenization_utils_fast import PreTrainedTokenizerFast from rosetta.utils import CondImage, ImageInfo, JointImageInfo, default @dataclass class Conversation(object): name: str roles: Tuple[str, str] = ("User", "Assistant") sep: str = "\n" sep2: str = None sep_sp: str = None stop_token_ids: list[int] = None pretrain_roles: Tuple[str, str] = ("", "") pretrain_sep: str = "" pretrain_sep2: str = "" pretrain_sep_sp: str = "" add_pad: bool = False add_bos: bool = True add_eos: bool = False def get_role_prefix(self, role): if role == "": return "" return f"<|im_start|>{role}\n" def empty(self, name=None): return Conversation( name=name or self.name, roles=self.roles, sep=self.sep, sep2=self.sep2, sep_sp=self.sep_sp, stop_token_ids=self.stop_token_ids, pretrain_roles=self.pretrain_roles, pretrain_sep=self.pretrain_sep, pretrain_sep2=self.pretrain_sep2, pretrain_sep_sp=self.pretrain_sep_sp, add_pad=self.add_pad, add_bos=self.add_bos, add_eos=self.add_eos, ) conv_templates: Dict[str, Conversation] = {} def register_conv_template(template: Conversation): assert template.name not in conv_templates, f"{template.name} has been registered." conv_templates[template.name] = template register_conv_template(Conversation( name="qwen-vl-30b-a3b-instruct", roles=("user", "assistant"), sep="<|im_end|>\n", sep2="<|im_end|>", sep_sp="\n\n", stop_token_ids=[151643, 151645], add_bos=False, )) for _template_name in ( "qwen3-06b-base-upcycling-moe-lm-deepseek", "qwen3-06b-base-upcycling-ours-lm", "qwen3-06b-upcycling-moe-mm-deepseek", "qwen3-06b-upcycling-ours-mm", "qwen3-06b-base-mot-lm", "qwen3-06b-mot", ): register_conv_template(conv_templates["qwen-vl-30b-a3b-instruct"].empty(name=_template_name)) def get_conversation_template(name: str) -> Conversation: return deepcopy(conv_templates[name]) ASSETS_BASE = os.getenv("ASSETS_BASE", "./public_assets").rstrip("/") TOKENIZER_BASE = os.getenv("TOKENIZER_BASE", f"{ASSETS_BASE}/pretrained_llm").rstrip("/") TOKENIZER_PATH = { "qwen3-0.6b-base": f"{TOKENIZER_BASE}/Qwen3-0.6B-Base", } class TokenizerEncodeOutput(BaseOutput): tokens: torch.Tensor = None text_slices: Optional[list[slice]] = None gen_image_slices: Optional[list[slice]] = None vae_image_slices: Optional[list[slice]] = None vit_image_slices: Optional[list[slice]] = None joint_image_slices: Optional[list[slice]] = None all_image_slices: Optional[list[slice]] = None text_mask: Optional[torch.Tensor] = None gen_image_mask: Optional[torch.Tensor] = None vae_image_mask: Optional[torch.Tensor] = None vit_image_mask: Optional[torch.Tensor] = None real_pos: Optional[torch.Tensor] = None cond_timestep_scatter_index: Optional[torch.Tensor] = None gen_timestep_scatter_index: Optional[torch.Tensor] = None und_token_indices: Optional[torch.Tensor] = None gen_token_indices: Optional[torch.Tensor] = None class BaseMultimodalTokenizerFast(PreTrainedTokenizerFast): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) special_tokens = self.special_tokens_map.get('additional_special_tokens', []) if len(special_tokens) > 0: special_token_ids = self.convert_tokens_to_ids(special_tokens) self._sp_dict = dict(zip(special_tokens, special_token_ids)) else: self._sp_dict = dict() self.setup_special_tokens() def setup_special_tokens(self): predefined_name_mapping = { "answer": "", "end_of_answer": "", "boi": "<｜boi｜>", "eoi": "<｜eoi｜>", "img": "<｜img｜>", } for name, mapping in predefined_name_mapping.items(): setattr(self, f"{name}_token", mapping) setattr(self, f"{name}_token_id", self.convert_tokens_to_ids(mapping)) if len(self._sp_dict) > 0: name_mapping = dict( cfg_token="<｜cfg｜>", timestep_token="<｜timestep｜>", joint_img_sep_token="<｜joint_img_sep｜>", ) for name, token in name_mapping.items(): if token in self._sp_dict: setattr(self, name, token) setattr(self, f"{name}_id", self._sp_dict[token]) def size_token(self, size: int): assert len(self._sp_dict) > 0, "Size tokens are not defined in the tokenizer." return f"<｜img_size_{size}｜>" def size_token_id(self, size: int): return self._sp_dict[self.size_token(size)] def ratio_token(self, ratio_idx: int): assert len(self._sp_dict) > 0, "Ratio tokens are not defined in the tokenizer." return f"<｜img_ratio_{ratio_idx}｜>" def ratio_token_id(self, ratio_idx: int): return self._sp_dict[self.ratio_token(ratio_idx)] def encode_text( self, *texts, uncond_enabled: Optional[bool | list[bool]] = None, uncond_p: Optional[float] = None, max_length: Optional[int] = None, pad: Optional[str] = None, ): if pad is not None: assert max_length is not None, "max_length should be provided when pad is not None." if uncond_enabled is None: uncond_enabled = [True] * len(texts) elif isinstance(uncond_enabled, bool): uncond_enabled = [uncond_enabled] * len(texts) assert len(uncond_enabled) == len(texts), ( f"Length of uncond_flags should be equal to the number of texts, " f"but got {len(uncond_enabled)} and {len(texts)}." ) do_uncond_drop = (uncond_p is not None) and (random.random() < uncond_p) text_tokens = [] cum_length = 0 for text, uncond_flag in zip(texts, uncond_enabled): if max_length is not None and cum_length >= max_length: break if isinstance(text, str): text_token = self.encode(text, add_special_tokens=False) else: text_token = text if uncond_flag and do_uncond_drop: text_token = [self.cfg_token_id] * len(text_token) if max_length is not None and (cum_length + len(text_token)) > max_length: text_token = text_token[:max_length - cum_length] text_tokens.extend(text_token) cum_length += len(text_token) if pad is not None and (pad_length := max_length - len(text_tokens)) > 0: if pad == 'left': text_tokens = [self.pad_token_id] * pad_length + text_tokens elif pad == 'right': text_tokens = text_tokens + [self.pad_token_id] * pad_length else: raise ValueError(f"Unsupported padding method: {pad}.") return text_tokens @staticmethod def _check_key_number_matched(keys, data): assert set(keys) == set(data.keys()), ( f"Keys in the template and token source should be matched, but got {keys} and {list(data.keys())}." ) key_counts = {k: 0 for k in keys} for key in keys: key_counts[key] += 1 for key, count in key_counts.items(): assert len(data[key]) == count, ( f"Number of `{key}` in the token source should be matched with the template, but got " f"{data[key]}({len(data[key])}) and {count}." ) def _add_meta_info_token( self, token_seq, token_count, extra_token_pos, add_timestep_token: bool = False, add_image_shape_token: bool = False, base_size=None, ratio_idx=None, token_height=None, token_width=None, image_type=None, media_type=None, und_token_type: list[str] = [], gen_token_type: list[str] = [], und_token_indices: list[int] = [], gen_token_indices: list[int] = [], token_count_start: int = 0, ): add_mot_indices = partial(self.process_mot_indices, und_token_indices=und_token_indices, gen_token_indices=gen_token_indices, und_token_type=und_token_type, gen_token_type=gen_token_type) if add_image_shape_token: token_seq.extend([self.size_token_id(base_size), self.ratio_token_id(ratio_idx)]) token_count += 2 add_mot_indices(token_type="vae_info", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count if add_timestep_token: token_seq.extend([self.timestep_token_id]) extra_token_pos['timestep'].append(token_count) if media_type is not None: if media_type == "gen_image": extra_token_pos['gen_timestep'].append(token_count) elif media_type in ["cond_joint_image", "cond_vae_image"]: extra_token_pos['cond_timestep'].append(token_count) else: raise ValueError(f"Unsupported image type: {media_type}.") token_count += 1 add_mot_indices(token_type="vae", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count return token_count, token_count_start def _shorten_text(self, text): text = re.sub(f"({self.img_token})+", lambda m: f"[{self.img_token}]{{{len(m.group(0)) // len(self.img_token)}}}", text) text = re.sub(f"({self.pad_token})+", lambda m: f"[{self.pad_token}]{{{len(m.group(0)) // len(self.pad_token)}}}", text) return text @staticmethod def process_mot_indices(token_type: str, token_indices: list[int], und_token_indices: list[int], gen_token_indices: list[int], und_token_type: list[str] = [], gen_token_type: list[str] = []): if token_type in und_token_type: und_token_indices.extend(token_indices) elif token_type in gen_token_type: gen_token_indices.extend(token_indices) def encode_sequence( self, template: str, token_source: dict[str, list[list[int] | dict[str, Any]]], total_length=None, add_eos=True, add_pad=True, add_bos=True, drop_last: str | bool = 'auto', add_image_shape_token=False, und_token_type: list[str] = [], gen_token_type: list[str] = [], ): if drop_last is True and total_length is None: raise ValueError("total_length should be provided when drop_last is True.") keys = template.split('-') index_indicator = {k: 0 for k in token_source} for k, v in token_source.items(): assert isinstance(v, (list, tuple)), ( f"Value of `{k}` in the token source should be a list or tuple, but got {type(v)}." ) self._check_key_number_matched(keys, token_source) token_seq = [] token_count = 0 extra_token_pos = defaultdict(list) und_token_indices = [] gen_token_indices = [] add_mot_indices = partial(self.process_mot_indices, und_token_indices=und_token_indices, gen_token_indices=gen_token_indices, und_token_type=und_token_type, gen_token_type=gen_token_type) if add_bos and self.bos_token_id is not None: token_seq.append(self.bos_token_id) add_mot_indices(token_type="bos", token_indices=[token_count]) token_count += 1 drop_last_break = False for i, key in enumerate(keys): source = token_source[key][index_indicator[key]] token_count_start = token_count if key == "text": token_seq.extend(source) extra_token_pos["_start"].append(token_count) token_count += len(source) extra_token_pos["_end"].append(token_count - 1) add_mot_indices(token_type="text", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count elif key == "gen_image": extra_count = \ 2 \ + (1 if source['add_timestep_token'] else 0) \ + (2 if source['add_image_shape_token'] else 0) if drop_last is True and token_count + extra_count + source['length'] > total_length: drop_last_break = True break token_seq.append(self.boi_token_id) extra_token_pos["boi"].append(token_count) add_mot_indices(token_type="boi", token_indices=[token_count]) token_count += 1 token_count_start = token_count token_count, token_count_start = self._add_meta_info_token( token_seq=token_seq, token_count=token_count, extra_token_pos=extra_token_pos, add_timestep_token=source['add_timestep_token'], add_image_shape_token=source['add_image_shape_token'], base_size=source.get('base_size'), ratio_idx=source.get('ratio_idx'), token_height=source.get('token_height'), token_width=source.get('token_width'), image_type=key, media_type=key, und_token_type=und_token_type, gen_token_type=gen_token_type, und_token_indices=und_token_indices, gen_token_indices=gen_token_indices, token_count_start=token_count_start, ) token_seq.extend( [self.img_token_id] * source['length'] ) extra_token_pos["_start"].append(token_count) extra_token_pos["_start"].append(token_count) token_count += source['length'] extra_token_pos["_end"].append(token_count - 1) extra_token_pos["_end"].append(token_count - 1) add_mot_indices(token_type="vae", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count token_seq.extend([self.eoi_token_id]) extra_token_pos["eoi"].append(token_count) add_mot_indices(token_type="eoi", token_indices=[token_count]) token_count += 1 token_count_start = token_count elif key == "cond_joint_image": assert isinstance(source['length'], list) and len( source['length']) == 2, "cond_joint_image length should be a list of two integers" extra_count = \ 2 + 1 \ + (1 if source['add_timestep_token'] else 0) \ + (2 if source['add_image_shape_token'] else 0) if drop_last is True and token_count + extra_count + sum(source['length']) > total_length: drop_last_break = True break token_seq.append(self.boi_token_id) extra_token_pos["boi"].append(token_count) token_count += 1 add_mot_indices(token_type="boi", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count token_count, token_count_start = self._add_meta_info_token( token_seq=token_seq, token_count=token_count, extra_token_pos=extra_token_pos, add_timestep_token=source['add_timestep_token'], add_image_shape_token=source['add_image_shape_token'], base_size=source.get('base_size'), ratio_idx=source.get('ratio_idx'), token_height=source.get('token_height'), token_width=source.get('token_width'), image_type=key, media_type=key, und_token_type=und_token_type, gen_token_type=gen_token_type, und_token_indices=und_token_indices, gen_token_indices=gen_token_indices, token_count_start=token_count_start, ) token_seq.extend( [self.img_token_id] * source['length'][0] ) extra_token_pos["_start"].append(token_count) extra_token_pos["_start"].append(token_count) extra_token_pos["_start"].append(token_count) token_count += source['length'][0] extra_token_pos["_end"].append(token_count - 1) extra_token_pos["_end"].append(token_count - 1) add_mot_indices(token_type="vae", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count token_seq.extend([self.joint_img_sep_token_id]) extra_token_pos["joint_img_sep"].append(token_count) add_mot_indices(token_type="joint_image_sep", token_indices=[token_count]) token_count += 1 token_count_start = token_count token_seq.extend( [self.img_token_id] * source['length'][1] ) extra_token_pos["_start"].append(token_count) extra_token_pos["_start"].append(token_count) token_count += source['length'][1] extra_token_pos["_end"].append(token_count - 1) extra_token_pos["_end"].append(token_count - 1) extra_token_pos["_end"].append(token_count - 1) add_mot_indices(token_type="vit", token_indices=list(range(token_count_start, token_count))) token_seq.extend( [self.eoi_token_id] ) extra_token_pos["eoi"].append(token_count) add_mot_indices(token_type="eoi", token_indices=[token_count]) token_count += 1 token_count_start = token_count elif key == "cond_vae_image": extra_count = \ 2 \ + (1 if source['add_timestep_token'] else 0) \ + (2 if source['add_image_shape_token'] else 0) if drop_last is True and token_count + extra_count + source['length'] > total_length: drop_last_break = True break token_seq.append(self.boi_token_id) extra_token_pos["boi"].append(token_count) add_mot_indices(token_type="boi", token_indices=[token_count]) token_count += 1 token_count_start = token_count token_count, token_count_start = self._add_meta_info_token( token_seq=token_seq, token_count=token_count, extra_token_pos=extra_token_pos, add_timestep_token=source['add_timestep_token'], add_image_shape_token=source['add_image_shape_token'], base_size=source.get('base_size'), ratio_idx=source.get('ratio_idx'), token_height=source.get('token_height'), token_width=source.get('token_width'), image_type=key, media_type=key, und_token_type=und_token_type, gen_token_type=gen_token_type, und_token_indices=und_token_indices, gen_token_indices=gen_token_indices, token_count_start=token_count_start, ) token_seq.extend( [self.img_token_id] * source['length'] ) extra_token_pos["_start"].append(token_count) extra_token_pos["_start"].append(token_count) token_count += source['length'] extra_token_pos["_end"].append(token_count - 1) extra_token_pos["_end"].append(token_count - 1) add_mot_indices(token_type="vae", token_indices=list(range(token_count_start, token_count))) token_seq.extend( [self.eoi_token_id] ) extra_token_pos["eoi"].append(token_count) add_mot_indices(token_type="eoi", token_indices=[token_count]) token_count += 1 token_count_start = token_count elif key == "cond_vit_image": extra_count = 2 if drop_last is True and token_count + extra_count + source['length'] > total_length: drop_last_break = True break token_seq.append(self.boi_token_id) add_mot_indices(token_type="boi", token_indices=[token_count]) token_count += 1 token_count_start = token_count token_seq.extend([self.img_token_id] * source['length']) extra_token_pos["_start"].append(token_count) extra_token_pos["_start"].append(token_count) token_count += source['length'] extra_token_pos["_end"].append(token_count - 1) extra_token_pos["_end"].append(token_count - 1) add_mot_indices(token_type="vit", token_indices=list(range(token_count_start, token_count))) token_count_start = token_count token_seq.append(self.eoi_token_id) extra_token_pos["eoi"].append(token_count) add_mot_indices(token_type="eoi", token_indices=[token_count]) token_count += 1 token_count_start = token_count else: raise ValueError(f"Not supported key: {key}") index_indicator[key] += 1 if add_eos is True and not drop_last_break: token_seq.append(self.eos_token_id) extra_token_pos["eos"].append(token_count) add_mot_indices(token_type="eos", token_indices=[token_count]) token_count += 1 token_count_start = token_count elif add_eos == 'auto' and not drop_last_break: if token_seq[-1] != self.eos_token_id and (total_length is None or token_count < total_length): token_seq.append(self.eos_token_id) extra_token_pos["eos"].append(token_count) add_mot_indices(token_type="eos", token_indices=[token_count]) token_count += 1 token_count_start = token_count if total_length: if token_count > total_length and drop_last: for start_key, end_key in [ ("_start", "_end"), ("_start", "_end"), ("_start", "_end"), ("_start", "_end"), ]: if start_key in extra_token_pos and end_key in extra_token_pos: assert all( (start > total_length or end + 1 < total_length) for start, end in zip(extra_token_pos[start_key], extra_token_pos[end_key]) ), ("Clip position should not be in the middle of the media tokens.\n" f"Below is the text:\n{self._shorten_text(self.decode(token_seq))}") token_seq = token_seq[:total_length] und_token_indices = [idx for idx in und_token_indices if idx < total_length] gen_token_indices = [idx for idx in gen_token_indices if idx < total_length] pad_num = max(0, total_length - len(token_seq)) if add_pad and pad_num: token_seq.extend([self.pad_token_id] * pad_num) extra_token_pos["first_pad"].append(token_count) add_mot_indices(token_type="pad", token_indices=list(range(token_count, token_count + pad_num))) if len(und_token_indices) > 0 and len(gen_token_indices) > 0: assert und_token_indices[-1] < len(token_seq) and gen_token_indices[-1] < len(token_seq), f"{und_token_indices[-1]=}, {gen_token_indices[-1]=}, {len(token_seq)=}" return token_seq, extra_token_pos, und_token_indices, gen_token_indices @staticmethod def parse_extra_token_pos(extra_token_pos, prefix, tokens, rng=None): if rng is None: rng = slice(None) image_slices = [ slice(start, end + 1) for start, end in zip(extra_token_pos[f'<{prefix}>_start'][rng], extra_token_pos[f'<{prefix}>_end'][rng]) ] if f'<{prefix}>_start' in extra_token_pos and f'<{prefix}>_end' in extra_token_pos else [] if image_slices: image_mask = torch.zeros_like(tokens, dtype=torch.bool) for image_slice in image_slices: image_mask[image_slice] = True else: image_mask = None return image_slices, image_mask def encode_general( self, sections: Optional[list[dict[str, Any]]] = None, max_token_length: Optional[int] = None, add_eos: bool | str = 'auto', use_text_mask: bool = True, add_pad: bool | str = 'auto', add_bos: bool = True, drop_last: bool | str = 'auto', und_token_type: list[str] = [], gen_token_type: list[str] = [], disable_ignore: bool = False, ): if sections is None: raise ValueError("sections must be provided.") template = '-'.join([section['type'] for section in sections]) sections = deepcopy(sections) token_source = defaultdict(list) text_mask_specs = [] for section in sections: if section['type'] == 'text': text = self.encode_text( section['text'] if 'text' in section else section['tokens'], uncond_enabled=section.get('uncond_enabled'), uncond_p=section.get('uncond_p'), max_length=section.get('max_length'), ) token_source['text'].append(text) text_mask_specs.append(dict( ignore=section.get('ignore', False), start_offset=section.get('start_offset', 0), end_offset=section.get('end_offset', 0), )) elif section['type'] == 'gen_image': token_source['gen_image'].append(dict( length=section['token_length'], add_timestep_token=section.get('add_timestep_token', False), add_image_shape_token=section.get('add_image_shape_token', False), base_size=section.get('base_size'), ratio_idx=section.get('ratio_idx'), token_height=section.get('token_height'), token_width=section.get('token_width'), )) elif section['type'] in ['cond_joint_image', 'cond_vae_image', 'cond_vit_image']: token_source[section['type']].append(dict( length=section['token_length'], add_timestep_token=section.get('add_timestep_token', False), add_image_shape_token=section.get('add_image_shape_token', False), base_size=section.get('base_size'), ratio_idx=section.get('ratio_idx'), token_height=section.get('token_height'), token_width=section.get('token_width'), )) else: raise ValueError(f"Invalid section type: {section['type']}") full_token_seq, extra_token_pos, und_token_indices, gen_token_indices = self.encode_sequence( template=template, token_source=dict(token_source), total_length=max_token_length, add_eos=add_eos, add_pad=add_pad, add_bos=add_bos, drop_last=drop_last, und_token_type=und_token_type, gen_token_type=gen_token_type, ) full_seq_token_tensor = torch.tensor(full_token_seq, dtype=torch.long) und_token_indices = torch.tensor(und_token_indices, dtype=torch.long) gen_token_indices = torch.tensor(gen_token_indices, dtype=torch.long) cond_timestep_scatter_index = torch.tensor(extra_token_pos['cond_timestep'], dtype=torch.long) \ if 'cond_timestep' in extra_token_pos else None gen_timestep_scatter_index = torch.tensor(extra_token_pos['gen_timestep'], dtype=torch.long) \ if 'gen_timestep' in extra_token_pos else None gen_image_slices, gen_image_mask = self.parse_extra_token_pos( extra_token_pos, 'img', full_seq_token_tensor) vae_image_slices, vae_image_mask = self.parse_extra_token_pos( extra_token_pos, 'vae_img', full_seq_token_tensor) vit_image_slices, vit_image_mask = self.parse_extra_token_pos( extra_token_pos, 'vit_img', full_seq_token_tensor) joint_image_slices, _ = self.parse_extra_token_pos( extra_token_pos, 'joint_img', full_seq_token_tensor) all_image_slices = [ slice(start, end + 1) for start, end in zip(extra_token_pos['_start'], extra_token_pos['_end']) ] if '_start' in extra_token_pos and '_end' in extra_token_pos else [] text_slices = [ slice(start, end + 1) for start, end in zip(extra_token_pos['_start'], extra_token_pos['_end']) ] if '_start' in extra_token_pos and '_end' in extra_token_pos else [] assert len(text_slices) <= len(text_mask_specs), \ (f"Number of text slices ({len(text_slices)}) should be less than or equal to " f"number of text mask specs ({len(text_mask_specs)})") if use_text_mask: text_mask = torch.zeros_like(full_seq_token_tensor, dtype=torch.float32) for text_slice, mask_spec in zip(text_slices, text_mask_specs): if not mask_spec['ignore'] or disable_ignore: real_slice = slice( text_slice.start + mask_spec['start_offset'], text_slice.stop + mask_spec['end_offset'] ) text_mask[real_slice] = 1.0 else: text_mask = None real_pos = torch.tensor(extra_token_pos.get('first_pad', [full_seq_token_tensor.shape[0]]), dtype=torch.long) if len(und_token_type) == 0 and len(gen_token_type) == 0: und_token_indices = None gen_token_indices = None return TokenizerEncodeOutput( tokens=full_seq_token_tensor, text_slices=text_slices, gen_image_slices=gen_image_slices, vae_image_slices=vae_image_slices, vit_image_slices=vit_image_slices, joint_image_slices=joint_image_slices, all_image_slices=all_image_slices, text_mask=text_mask, gen_image_mask=gen_image_mask, vae_image_mask=vae_image_mask, vit_image_mask=vit_image_mask, real_pos=real_pos, cond_timestep_scatter_index=cond_timestep_scatter_index, gen_timestep_scatter_index=gen_timestep_scatter_index, und_token_indices=und_token_indices, gen_token_indices=gen_token_indices, ) def apply_general_template( self, message_list, conv_template, max_length=None, add_assistant_prefix=False, answer="auto", bot_task="auto", sequence_template="instruct", uncond_p=0.0, cfg_factor=1, batchify=False, image_base_size=None, und_token_type=None, gen_token_type=None, use_text_mask=False, ): if bot_task == "img_ratio": assert image_base_size is not None, "image_base_size should be provided for img_ratio task." if batchify: assert isinstance(message_list[0], list), \ f"When batchify is True, message_list should be a list of list, but got [{type(message_list[0])}, ...]." return self.batch_gen_infer( infer_fn=self.apply_general_template, infer_fn_kwargs_list=[dict( message_list=message_list_i, conv_template=conv_template, max_length=max_length, add_assistant_prefix=add_assistant_prefix, answer=answer, bot_task=bot_task, sequence_template=sequence_template, image_base_size=image_base_size, und_token_type=und_token_type, gen_token_type=gen_token_type, use_text_mask=use_text_mask, ) for message_list_i in message_list], do_classifier_free_guidance=cfg_factor > 1, uncondition_repeat_times=cfg_factor - 1, und_token_type=und_token_type, gen_token_type=gen_token_type, ) uncond_kwargs = dict( uncond_enabled=uncond_p == 1.0, uncond_p=uncond_p, ) def process_successive_message(_message_list, _cur_message_idx, role, prefix, suffix, answer_prefix="", answer_suffix=""): _sub_sections = [] while _cur_message_idx < len(message_list) and _message_list[_cur_message_idx]['role'] == role: message = _message_list[_cur_message_idx] if message['type'] == 'text': text = message['content'] if role == "system": _sub_sections.append(dict(type="text", text=text)) elif role == "assistant": _sub_sections.append(dict( type="text", text=f"{answer_prefix}{text}{answer_suffix}", **uncond_kwargs)) else: _sub_sections.append(dict(type="text", text=text, **uncond_kwargs)) elif message['type'] == 'gen_image': info = message['content'] assert isinstance(info, ImageInfo), f"Expected ImageInfo, but got {type(info)}" if role == "assistant": _sub_sections.append(dict(type="text", text=answer_prefix)) _sub_sections.append(dict(type=message['type'], **info.meta_info)) if role == "assistant": _sub_sections.append(dict(type="text", text=answer_suffix)) elif message['type'] in ['cond_joint_image', 'cond_vae_image', 'cond_vit_image']: info = message['content'] assert isinstance(info, (ImageInfo, JointImageInfo)), \ f"Expected ImageInfo or JointImageInfo, but got {type(info)}" _sub_sections.append(dict(type=message['type'], **info.meta_info)) else: raise ValueError(f"Unknown message type: {message['type']}") _cur_message_idx += 1 if len(_sub_sections) > 0: _sub_sections.insert(0, dict(type='text', text=prefix)) _sub_sections.append(dict(type='text', text=suffix)) return _sub_sections, _cur_message_idx if (answer == "auto" and sequence_template == "instruct") or answer is True: answer_prefix, answer_suffix = self.answer_token, self.end_of_answer_token else: answer_prefix, answer_suffix = "", "" if sequence_template == "pretrain": system_suffix = conv_template.pretrain_sep_sp user_prefix = conv_template.get_role_prefix(conv_template.pretrain_roles[0]) user_suffix = conv_template.pretrain_sep bot_prefix = conv_template.get_role_prefix(conv_template.pretrain_roles[1]) bot_suffix = conv_template.pretrain_sep2 else: system_suffix = conv_template.sep_sp user_prefix = conv_template.get_role_prefix(conv_template.roles[0]) user_suffix = f"{conv_template.sep}" bot_prefix = conv_template.get_role_prefix(conv_template.roles[1]) bot_suffix = f"{conv_template.sep2}" sections = [] cur_message_idx = 0 final_role = None while cur_message_idx < len(message_list): sub_sections, cur_message_idx = process_successive_message( message_list, cur_message_idx, role="system", prefix="", suffix=system_suffix) sections.extend(sub_sections) if len(sub_sections) > 0: final_role = "system" sub_sections, cur_message_idx = process_successive_message( message_list, cur_message_idx, role="user", prefix=user_prefix, suffix=user_suffix) sections.extend(sub_sections) if len(sub_sections) > 0: final_role = "user" sub_sections, cur_message_idx = process_successive_message( message_list, cur_message_idx, role="assistant", prefix=bot_prefix, suffix=bot_suffix, answer_prefix=answer_prefix, answer_suffix=answer_suffix, ) sections.extend(sub_sections) if len(sub_sections) > 0: final_role = "assistant" if add_assistant_prefix: if final_role == "assistant": _bot_prefix = "" if len(sections) > 0 and sections[-1]['type'] == 'text' and sections[-1]['text'] == bot_suffix: sections = sections[:-1] else: _bot_prefix = bot_prefix bot_response_prefix = dict( auto=lambda: f"{_bot_prefix}{answer_prefix}", image=lambda: "", img_ratio=lambda: f"{_bot_prefix}{answer_prefix}{self.boi_token}{self.size_token(image_base_size)}", )[bot_task]() sections.append(dict(type='text', text=bot_response_prefix)) if und_token_type is None: und_token_type = [] if gen_token_type is None: gen_token_type = [] output = self.encode_general( sections=sections, use_text_mask=use_text_mask, add_eos=conv_template.add_eos, add_pad=conv_template.add_pad, add_bos=conv_template.add_bos, und_token_type=und_token_type, gen_token_type=gen_token_type, ) if max_length is not None: if output.tokens.shape[-1] > max_length: raise ValueError( f"Encoded token length {output.tokens.shape[-1]} exceeds max_length {max_length}.\n" f"Please set a larger max_length or check the input messages:\n{message_list}" ) return output, sections def apply_chat_template( self, batch_prompt: Optional[list[str]] = None, batch_message_list: Optional[list[list[dict[str, Any]]]] = None, mode: str = "gen_text", batch_gen_image_info: Optional[list[ImageInfo]] = None, batch_cond_images: Optional[Union[list[CondImage], list[list[CondImage]]]] = None, max_length: Optional[int] = None, bot_task: str = "auto", image_base_size: Optional[int] = None, sequence_template: str = "pretrain", cfg_factor: int = 1, add_assistant_prefix: Optional[bool] = None, conv_template: Optional[Conversation] = None, und_token_type: list[str] = None, gen_token_type: list[str] = None, use_text_mask: bool = False, **kwargs, ) -> dict[str, Any]: allowed_tasks = ["image", "auto", "img_ratio"] assert bot_task in allowed_tasks, f"bot_task should be one of {allowed_tasks}, but got {bot_task}." if batch_message_list is None: batch_size = len(batch_prompt) if not isinstance(batch_gen_image_info, list): batch_gen_image_info = [batch_gen_image_info] * batch_size if batch_cond_images is not None: assert len(batch_cond_images) == batch_size, \ (f"batch_cond_image_info should have the same length as batch_size ({batch_size}), " f"but got {len(batch_cond_images)}.") batch_cond_images = [ cond_images if isinstance(cond_images, list) else [cond_images] for cond_images in batch_cond_images ] else: batch_cond_images = [[] for _ in range(batch_size)] batch_message_list = [] for prompt, gen_image_info, cond_images in zip( batch_prompt, batch_gen_image_info, batch_cond_images, ): message_list = [] if len(cond_images) > 0: message_list.extend([ dict(role="user", type=cond_image.section_type, content=cond_image.i) for cond_image in cond_images ]) message_list.append(dict(role="user", type="text", content=prompt)) if mode == "gen_image": message_list.append(dict( role="assistant", type="gen_image", content=gen_image_info)) batch_message_list.append(message_list) output, sections = self.apply_general_template( message_list=batch_message_list, conv_template=conv_template, max_length=max_length, add_assistant_prefix=default(add_assistant_prefix, mode == "gen_text"), bot_task=bot_task, sequence_template=sequence_template, cfg_factor=cfg_factor, batchify=True, image_base_size=image_base_size, und_token_type=und_token_type, gen_token_type=gen_token_type, use_text_mask=use_text_mask, **kwargs, ) return dict(output=output, sections=sections) def pad(self, tensor_list, dim=0, pad_val=None, key=None): if pad_val is None: pad_val = self.pad_token_id max_len = max([t.shape[dim] for t in tensor_list]) padded_tensor_list = [] for t in tensor_list: if t.shape[dim] < max_len: assert pad_val is not False, f"Not allowed/implemented pad for key: {key}" t = F.pad(t, (0, max_len - t.shape[dim]), value=pad_val) padded_tensor_list.append(t) return padded_tensor_list def batch_gen_infer( self, infer_fn, infer_fn_kwargs_list: list[dict[str, int]] = None, do_classifier_free_guidance=False, uncondition_repeat_times: int = 1, und_token_type: Optional[list] = None, gen_token_type: Optional[list] = None, ): cond_results_list = None uncond_results_list = None for infer_fn_kwargs in infer_fn_kwargs_list: cond_kwargs = {"uncond_p": 0.0} if do_classifier_free_guidance else {} results = infer_fn( **infer_fn_kwargs, **cond_kwargs, ) assert isinstance(results, tuple), f"Expected tuple output from tokenizer template, got {type(results)}." if cond_results_list is None: cond_results_list = [[] for _ in results] uncond_results_list = [[] for _ in results] for i, result in enumerate(results): cond_results_list[i].append(result) if do_classifier_free_guidance: uncond_results = infer_fn( **infer_fn_kwargs, uncond_p=1.0, ) if isinstance(uncond_results, TokenizerEncodeOutput): uncond_results_list.append(uncond_results) else: for i, result in enumerate(uncond_results): uncond_results_list[i].append(result) def make_batch(batch_cond_item, batch_uncond_item): first = batch_cond_item[0] if isinstance(first, (list, tuple)): stacked_item = batch_cond_item + batch_uncond_item * uncondition_repeat_times elif isinstance(first, TokenizerEncodeOutput): stacked_item = {} for key in list(first.keys()): merged_list = [cond_item[key] for cond_item in batch_cond_item] + \ [uncond_item[key] for uncond_item in batch_uncond_item] * uncondition_repeat_times if isinstance(first[key], torch.Tensor): if 'mask' in key: pad_val = 0.0 elif key == 'tokens': pad_val = self.pad_token_id elif key in ['und_token_indices', 'gen_token_indices']: continue else: pad_val = False if key not in ('und_token_indices', 'gen_token_indices'): stacked_item[key] = torch.stack(self.pad(merged_list, pad_val=pad_val, key=key), dim=0) elif isinstance(first[key], list): stacked_item[key] = merged_list elif first[key] is None: pass else: raise ValueError(f"Unsupported type of {key}: {type(first[key])}.") stacked_item = TokenizerEncodeOutput(stacked_item) if 'und_token_indices' in first.keys() and first['und_token_indices'] is not None and 'gen_token_indices' in first.keys() and first['gen_token_indices'] is not None: und_token_indices_merged_list = [cond_item['und_token_indices'] for cond_item in batch_cond_item] + [uncond_item['und_token_indices'] for uncond_item in batch_uncond_item] * uncondition_repeat_times gen_token_indices_merged_list = [cond_item['gen_token_indices'] for cond_item in batch_cond_item] + [uncond_item['gen_token_indices'] for uncond_item in batch_uncond_item] * uncondition_repeat_times sequence_length = stacked_item["tokens"].shape[1] def _safe_max_index(indices): return indices.max().item() if indices.numel() > 0 else -1 max_index = [ max( _safe_max_index(und_token_indices_merged_list[i]), _safe_max_index(gen_token_indices_merged_list[i]), ) for i in range(len(und_token_indices_merged_list)) ] for i, (und_token_indices_item, max_index_item) in enumerate(zip(und_token_indices_merged_list, max_index)): if max_index_item == sequence_length - 1: continue und_token_indices_merged_list[i] = torch.cat([und_token_indices_item, torch.arange(max_index_item + 1, sequence_length)]) max_gen_count = max(g.shape[0] for g in gen_token_indices_merged_list) max_extra_needed = 0 for i in range(len(gen_token_indices_merged_list)): pad_needed = max_gen_count - gen_token_indices_merged_list[i].shape[0] pad_available = max(0, sequence_length - 1 - max_index[i]) max_extra_needed = max(max_extra_needed, pad_needed - pad_available) if max_extra_needed > 0: for key in list(stacked_item.keys()): if key == 'tokens': stacked_item[key] = F.pad(stacked_item[key], (0, max_extra_needed), value=self.pad_token_id) elif 'mask' in key and isinstance(stacked_item[key], torch.Tensor): stacked_item[key] = F.pad(stacked_item[key], (0, max_extra_needed), value=0.0) new_positions = torch.arange(sequence_length, sequence_length + max_extra_needed) for i in range(len(und_token_indices_merged_list)): und_token_indices_merged_list[i] = torch.cat([und_token_indices_merged_list[i], new_positions]) sequence_length += max_extra_needed for i in range(len(gen_token_indices_merged_list)): pad_needed = max_gen_count - gen_token_indices_merged_list[i].shape[0] if pad_needed > 0: moved = und_token_indices_merged_list[i][-pad_needed:] und_token_indices_merged_list[i] = und_token_indices_merged_list[i][:-pad_needed] gen_token_indices_merged_list[i] = torch.cat([gen_token_indices_merged_list[i], moved]) stacked_item['und_token_indices'] = torch.stack(und_token_indices_merged_list, dim=0) stacked_item['gen_token_indices'] = torch.stack(gen_token_indices_merged_list, dim=0) elif ('und_token_indices' in first.keys() and first['und_token_indices'] is not None) or ('gen_token_indices' in first.keys() and first['gen_token_indices'] is not None): raise ValueError(f"Only one of 'und_token_indices' and 'gen_token_indices' exists.") stacked_item = TokenizerEncodeOutput(stacked_item) else: raise TypeError(f"Making batch on type {type(first)} is not supported.") return stacked_item stacked_outputs = [] for cond_results, uncond_results in zip(cond_results_list, uncond_results_list): stacked_outputs.append(make_batch(cond_results, uncond_results)) return tuple(stacked_outputs) class Qwen3BaseTokenizerFast(BaseMultimodalTokenizerFast): def setup_special_tokens(self): predefined_name_mapping = { "bos": "<|im_start|>", "eos": "<|im_end|>", "answer": "", "end_of_answer": "", "boi": "<|vision_start|>", "eoi": "<|vision_end|>", "img": "<|image_pad|>", } for name, mapping in predefined_name_mapping.items(): setattr(self, f"{name}_token", mapping) setattr(self, f"{name}_token_id", self.convert_tokens_to_ids(mapping)) if len(self._sp_dict) > 0: name_mapping = dict( cfg_token="<｜cfg｜>", timestep_token="<｜timestep｜>", joint_img_sep_token="<｜joint_img_sep｜>", ) for name, token in name_mapping.items(): setattr(self, name, token) setattr(self, f"{name}_id", self._sp_dict.get(token)) def load_tokenizer( tokenizer_name: str, tokenizer_class: str, ) -> "BaseMultimodalTokenizerFast": assert '.' in tokenizer_class, ( f"Invalid tokenizer class: {tokenizer_class}. A valid tokenizer name should be in the form of " f".." ) if tokenizer_class.startswith("transformers."): module_name, tokenizer_cls = tokenizer_class.rsplit('.', 1) module_spec = importlib.import_module(module_name) TokenizerSpec = getattr(module_spec, tokenizer_cls) # noqa else: module_name, tokenizer_cls = tokenizer_class.rsplit('.', 1) if module_name == "tokenizer": module_spec = importlib.import_module("rosetta.tokenizer") else: module_spec = importlib.import_module(module_name) TokenizerSpec = getattr(module_spec, tokenizer_cls) # noqa if tokenizer_name in TOKENIZER_PATH: tokenizer_name = TOKENIZER_PATH[tokenizer_name] return TokenizerSpec.from_pretrained(tokenizer_name)