cpm-bee-2b / tokenization_cpmbee.py

Gong Baitao

fix normal token with <>

31ce9ce 11 months ago

No virus

40 kB

	# coding=utf-8
	# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	"""Tokenization classes for CpmBee."""
	import json
	import os
	from typing import Any, Dict, List, Optional, Tuple, Union

	import numpy as np
	from numpy.typing import NDArray
	from typing_extensions import TypedDict

	from transformers.tokenization_utils import PaddingStrategy, PreTrainedTokenizer, TensorType
	from transformers.tokenization_utils_base import AddedToken, BatchEncoding, TextInput, TruncationStrategy
	from transformers.utils import logging


	logger = logging.get_logger(__name__)

	VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}

	PRETRAINED_VOCAB_FILES_MAP = {
	"vocab_file": {
	"openbmb/cpm-bee-10b": "https://huggingface.co/openbmb/cpm-bee-10b/blob/main/vocab.txt",
	"openbmb/cpm-bee-5b": "https://huggingface.co/openbmb/cpm-bee-5b/blob/main/vocab.txt",
	"openbmb/cpm-bee-2b": "https://huggingface.co/openbmb/cpm-bee-2b/blob/main/vocab.txt",
	"openbmb/cpm-bee-1b": "https://huggingface.co/openbmb/cpm-bee-1b/blob/main/vocab.txt",
	},
	}

	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
	"openbmb/cpm-bee-10b": 4096,
	"openbmb/cpm-bee-5b": 4096,
	"openbmb/cpm-bee-2b": 4096,
	"openbmb/cpm-bee-1b": 4096,
	}


	class _PrevExtTableStates(TypedDict):
	ext_table: Dict[int, str]
	token_id_table: Dict[str, Dict[int, int]]


	CPMBeeInputType = Union[str, Dict[str, "CPMBeeInputType"]]


	def rel_to_bucket(n_up: int, n_down: int, max_depth: int = 8):
	ret = n_up * max_depth + n_down
	if ret == 0:
	return ret
	else:
	# bucket 1 is reserved for incontext samples
	return ret + 1


	class _DictTree(TypedDict):
	value: str
	children: List["_DictTree"]
	depth: int
	segment_id: int
	need_predict: bool


	class CpmBeeTokenizer(PreTrainedTokenizer):
	"""
	Construct a CPMBee tokenizer.

	Args:
	vocab_file (`str`):
	Path to the vocabulary file.
	bos_token (`str`, optional, defaults to `"<s>"`):
	The beginning of sequence token.
	eos_token (`str`, optional, defaults to `"</s>"`):
	The end of sequence token.
	line_token (`str`, optional, defaults to `"\n"`):
	The line token.
	space_token (`str`, optional, defaults to `" "`):
	The space token.
	unk_token (`str`, optional, defaults to `"<unk>"`):
	The unknown token.
	mask_token (`str`, optional, defaults to `"<mask>"`):
	The mask token.
	pad_token (`str`, optional, defaults to `"<pad>"`):
	The token used for padding.
	padding_side (`str`, optional, defaults to `"left"`):
	The padding side. CPM-Bee will use left padding by default.
	"""

	vocab_files_names = VOCAB_FILES_NAMES
	pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
	max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	model_input_names: List[str] = [
	"input_ids",
	"attention_mask",
	"input_id_sub",
	"position",
	"context",
	"sample_ids",
	"num_segments",
	"segment",
	"segment_rel_offset",
	"segment_rel",
	]
	add_prefix_space = False

	def __init__(
	self,
	vocab_file,
	bos_token="<s>",
	eos_token="</s>",
	line_token="\n",
	space_token=" ",
	unk_token="<unk>",
	mask_token="<mask>",
	pad_token="<pad>",
	padding_side="left",
	**kwargs,
	):
	super().__init__(
	bos_token=bos_token,
	eos_token=eos_token,
	line_token=line_token,
	space_token=space_token,
	unk_token=unk_token,
	mask_token=mask_token,
	pad_token=pad_token,
	padding_side=padding_side,
	**kwargs,
	)

	self.encoder: Dict[str, int] = {}

	with open(vocab_file, "r", encoding="utf-8") as reader:
	for token in reader.readlines():
	token = token.rstrip("\n")
	if len(token) == 0:
	continue
	self.encoder[token] = len(self.encoder)

	self.encoder[" "] = self.encoder["</_>"]
	self.encoder["\n"] = self.encoder["</n>"]
	del self.encoder["</_>"]
	del self.encoder["</n>"]

	self.decoder = {v: k for k, v in self.encoder.items()}

	self._max_word_len = max([len(x) for x in self.encoder.keys()])
	self.cpmbee_special_tokens = {k: v for k, v in self.encoder.items() if k.startswith("<") and k.endswith(">")}

	self.ext_table: Dict[int, str] = {}
	self.ext_table_rev: Dict[str, int] = {}

	self.token_id_table: Dict[str, Dict[int, int]] = {}
	self.ext_special_tokens = []

	self.ext_args_for_model = [
	"input_id_subs",
	"input_pos",
	"context",
	"segment_ids",
	"segment_rel_offset",
	"segment_rel",
	"sample_ids",
	"num_segments",
	"predict_segments",
	"answer_placeholders",
	"ext_table",
	"token_id_table",
	]

	@property
	def bod_token_id(self):
	return self.encoder[self.bod_token]

	@property
	def eod_token_id(self):
	return self.encoder[self.eod_token]

	@property
	def newline_id(self):
	return self.encoder[self.line_token]

	@property
	def vocab_size(self) -> int:
	return len(self.encoder)

	def __len__(self):
	"""
	Size of the full vocabulary with the added tokens.
	"""
	return self.vocab_size + len(self.added_tokens_encoder)

	def get_vocab(self):
	return dict(self.encoder, **self.added_tokens_encoder)

	def get_piece(self, text: str) -> str:
	"""
	Match with maximum length.
	"""
	len_text = len(text)
	for i in range(len(text)):
	sub = text[: len_text - i]
	if (sub in self.encoder) or (sub in self.added_tokens_encoder):
	return sub
	return text[0]

	def tokenize(self, text: TextInput, **kwargs) -> List[str]:
	r"""
	Override the `tokenize` to meet the needs of CPMBee:
	1. Mark the special token with `<` and `>`. The `<>` will be ignored.
	2. Split sentences by the marked special tokens.
	3. Record the marked special token by `ext_table` and `ext_table_rev`.
	4. Tokenize the sentence without special tokens.
	"""
	for_cpmbee = kwargs.get("for_cpmbee", False)
	all_special_tokens_extended = {
	str(t): t for t in self.all_special_tokens_extended if isinstance(t, AddedToken)
	}

	sentence_split = [""]
	is_special_token = False
	for i, c in enumerate(text):
	if is_special_token:
	if c == "<":
	tail = sentence_split.pop(-1)
	sentence_split[-1] += tail
	sentence_split.append(c)
	is_special_token = False
	elif c == ">":
	# end of special token
	sentence_split[-1] += c
	if sentence_split[-1] == "<>":
	continue
	is_special_token = False
	sentence_split.append("")
	else:
	sentence_split[-1] += c
	else:
	if c == "<":
	is_special_token = True
	sentence_split.append(c)
	else:
	sentence_split[-1] += c
	if is_special_token:
	tail = sentence_split.pop(-1)
	sentence_split[-1] += tail

	output_tokens = []
	for i, part in enumerate(sentence_split):
	if (i & 1) == 1:
	# special token
	output_tokens.append(part)
	if for_cpmbee and (part not in self.encoder) and (part not in self.ext_table_rev):
	self.ext_table_rev[part] = len(self.ext_table_rev) + self.vocab_size
	self.ext_table[self.ext_table_rev[part]] = part
	else:
	output_tokens.extend(self._tokenize(part, for_cpmbee=for_cpmbee))

	# drop spaces
	for i, token in enumerate(output_tokens):
	if token in self.added_tokens_encoder:
	token = all_special_tokens_extended.get(token, None)
	left = output_tokens[i - 1] if i > 0 else None
	right = output_tokens[i + 1] if i < len(output_tokens) - 1 else None
	if isinstance(token, AddedToken):
	if token.rstrip and right:
	# A bit counter-intuitive but we strip the left of the string
	# since tok_extended.rstrip means the special token is eating all white spaces on its right
	output_tokens[i + 1] = right.lstrip()
	# Strip white spaces on the left
	if token.lstrip and left:
	output_tokens[i - 1] = left.rstrip() # Opposite here
	else:
	if right:
	output_tokens[i + 1] = right.lstrip()
	if left:
	output_tokens[i - 1] = left.rstrip()

	skipped_tokens = []
	for token in output_tokens:
	if not token:
	continue
	else:
	skipped_tokens.append(token)

	return skipped_tokens

	def _tokenize(self, text, **kwargs):
	"""
	Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
	vocabulary.

	Do NOT take care of added tokens. Record the unk tokens and special tokens in `ext_table` and `ext_table_rev`.
	"""
	for_cpmbee = kwargs.get("for_cpmbee", False)
	output_tokens = []

	part_st = 0
	last_unk = None
	while part_st < len(text):
	piece = self.get_piece(text[part_st:])
	if piece in self.encoder or self.added_tokens_encoder:
	if last_unk is None:
	output_tokens.append(piece)
	else:
	if for_cpmbee and (last_unk not in self.ext_table_rev):
	self.ext_table_rev[last_unk] = len(self.ext_table_rev) + self.vocab_size
	self.ext_table[self.ext_table_rev[last_unk]] = last_unk
	output_tokens.append(last_unk)
	output_tokens.append(piece)
	last_unk = None
	else:
	if last_unk is None:
	last_unk = piece
	else:
	last_unk += piece
	part_st += len(piece)
	if last_unk is not None:
	# part end with UNK
	if for_cpmbee and (last_unk not in self.ext_table_rev):
	self.ext_table_rev[last_unk] = len(self.ext_table_rev) + self.vocab_size
	self.ext_table[self.ext_table_rev[last_unk]] = last_unk
	output_tokens.append(last_unk)

	return output_tokens

	def check(self, token):
	return token in self.encoder

	def convert_tokens_to_string(self, tokens: List[str]) -> str:
	return "".join(tokens)

	def _convert_token_to_id(self, token: str):
	"""Converts a token (str) in an id using the vocab and ext_table."""
	if token in self.encoder:
	return self.encoder.get(token)
	elif token in self.ext_table_rev:
	return self.ext_table_rev[token]
	elif token in self.added_tokens_encoder:
	return self.added_tokens_encoder[token]
	else:
	return self.unk_token_id

	def _convert_id_to_token(self, index):
	"""Converts an index (integer) in a token (str) using the vocab and ext_table."""
	if index in self.ext_table:
	return self.ext_table[index]
	elif index in self.added_tokens_decoder:
	return self.added_tokens_decoder[index]
	else:
	if index >= 0:
	return self.decoder[index]

	def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
	if os.path.isdir(save_directory):
	vocab_file = os.path.join(
	save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
	)
	else:
	vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
	index = 0
	self.encoder["</n>"] = self.encoder["\n"]
	del self.encoder["\n"]
	self.encoder["</_>"] = self.encoder[" "]
	del self.encoder[" "]
	with open(vocab_file, "w", encoding="utf-8") as writer:
	for token, token_index in sorted(self.encoder.items(), key=lambda x: x[1]):
	if index != token_index:
	logger.warning(
	f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
	" Please check that the vocabulary is not corrupted!"
	)
	index = token_index
	writer.write(token + "\n")
	index += 1
	return (vocab_file,)

	def __call__(self, text, args, *kwargs):
	r"""
	CPMBee `call` method will use `_tokenize_cpmbee` when the input type is dict.
	"""
	if isinstance(text, dict):
	return self._batch_tokenize_cpmbee([text], args, *kwargs)
	elif isinstance(text, (list, tuple)):
	if isinstance(text[0], dict):
	return self._batch_tokenize_cpmbee(text, args, *kwargs)
	else:
	return super().__call__(text, args, *kwargs)
	else:
	return super().__call__(text, args, *kwargs)

	# 分词
	def _tokenize_cpmbee(self, data: TextInput, args, *kwargs) -> List[str]:
	"""
	A tokenize method to process dict data. Exclusive for CPMBee.
	"""
	if isinstance(data, str):
	data = json.loads(data)
	if not isinstance(data, Dict):
	raise TypeError(
	"CpmBeeTokenizer input data should be dict or str in dict format, but got {}".format(type(data))
	)

	# 1. prepare answer placeholder
	answer_placeholders = []

	def _put_placeholder(data: Any, path: List[str] = []):
	if isinstance(data, dict):
	ret = {}
	for k, v in data.items():
	ret[k] = _put_placeholder(v, path + [k])
	return ret
	else:
	answer_placeholders.append(path)
	return "<ans_{}>".format(len(answer_placeholders))

	data["<ans>"] = _put_placeholder(data["<ans>"])

	(
	input_ids,
	input_id_subs,
	context,
	segment_ids,
	segment_rel,
	n_segments,
	table_states,
	) = self.convert_data_to_id(data, shuffle_answer=False, max_depth=8)

	# <ans> mapping from sub to id
	sub_ans_map: Dict[int, int] = {}
	for fake_id, token_sub in table_states["token_id_table"]["<ans>"].items():
	token = table_states["ext_table"][fake_id]
	if token.startswith("<ans_") and token.endswith(">"):
	ans_id = int(token[5:-1])
	sub_ans_map[token_sub] = ans_id

	tmp_input_ids = []
	tmp_input_sub = []
	tmp_input_seg = []

	# get predict segments
	predict_segments: List[Tuple[int, int]] = []
	for i in range(input_ids.shape[0]):
	if context[i] == 0:
	if input_ids[i] == self.encoder["<ans>"]:
	# is ans
	# (segment_id, ans_id)
	predict_segments.append((segment_ids[i], sub_ans_map[input_id_subs[i]]))
	else:
	tmp_input_ids.append(input_ids[i])
	tmp_input_sub.append(input_id_subs[i])
	tmp_input_seg.append(segment_ids[i])

	if len(predict_segments) == 0:
	raise ValueError("No answer to predict")

	input_ids = np.array(tmp_input_ids, dtype=np.int32) # all context
	input_id_subs = np.array(tmp_input_sub, dtype=np.int32) # [0, 0, 0, 0, 1, 0, 0, 2, 0, ...]
	context = np.full_like(tmp_input_ids, 1, dtype=np.int8) # [1, 1, 1, ...]
	segment_ids = np.array(tmp_input_seg, dtype=np.int32) # [0, 0, 0, 1, 1, 1, 2, 2, 2, 2, ...]
	sample_ids = np.zeros(input_ids.shape, dtype=np.int32) # [0, 0, 0, 0, ...]
	segment_rel_offset = np.zeros(input_ids.shape, dtype=np.int32) # [0, 0, 0, ...]
	num_segments = np.full(input_ids.shape, n_segments, dtype=np.int32) # [n_seg, n_seg, n_seg, ...]
	input_pos = np.arange(input_ids.shape[0], dtype=np.int32) # [0, 1, 2, 3, 4, ...]

	return (
	self.prepare_for_model(
	input_ids.tolist(),
	input_id_subs=input_id_subs.tolist(),
	input_pos=input_pos.tolist(),
	context=context.tolist(),
	segment_ids=segment_ids.tolist(),
	segment_rel_offset=segment_rel_offset.tolist(),
	segment_rel=segment_rel.tolist(),
	sample_ids=sample_ids.tolist(),
	num_segments=num_segments.tolist(),
	**kwargs,
	),
	predict_segments,
	answer_placeholders,
	table_states["ext_table"],
	table_states["token_id_table"],
	)

	def _batch_tokenize_cpmbee(self, data_lst, args, *kwargs):
	"""
	Batched _token_cpmbee.
	"""
	device = kwargs.get("device", "cpu")
	return_tensors = kwargs.get("return_tensors", None)
	batch_outputs = {}
	segment_rel_pack = []
	other_info = []

	batch_ext_table_map: Dict[Tuple[int, int], int] = {}
	batch_ext_table_ids: List[int] = []
	batch_ext_table_sub: List[int] = []

	for data in data_lst:
	self.ext_table = {}
	self.ext_table_rev = {}
	self.token_id_table = {}
	(outputs, predict_segments, answer_placeholders, ext_table, token_id_table) = self._tokenize_cpmbee(
	data,
	truncation=None,
	padding=PaddingStrategy.DO_NOT_PAD.value,
	max_length=None,
	pad_to_multiple_of=None,
	return_attention_mask=False,
	return_tensors=None,
	)
	rev_ext_table = {}
	for token, mp in token_id_table.items():
	if token == "<ans>":
	continue
	token_id = self.encoder[token]
	for fake_id, token_sub in mp.items():
	if token_sub > 0:
	if (token_id, token_sub) not in batch_ext_table_map:
	batch_ext_table_map[(token_id, token_sub)] = len(batch_ext_table_ids) + self.vocab_size
	batch_ext_table_ids.append(token_id)
	batch_ext_table_sub.append(token_sub)
	rev_ext_table[batch_ext_table_map[(token_id, token_sub)]] = ext_table[fake_id]
	else:
	rev_ext_table[token_id] = ext_table[fake_id]

	segment_rel_pack.append(np.array(outputs.pop("segment_rel")))
	other_info.append(
	{
	"predict_segments": predict_segments,
	"answer_placeholders": answer_placeholders,
	"ext_table": rev_ext_table,
	}
	)

	for key, value in outputs.items():
	if key not in batch_outputs:
	batch_outputs[key] = []
	batch_outputs[key].append(value)

	max_length = max([len(item) for item in batch_outputs[self.model_input_names[0]]])
	batch_size = len(batch_outputs[self.model_input_names[0]])
	for i in range(batch_size):
	inputs = {k: v[i] for k, v in batch_outputs.items()}

	for k, v in inputs.items():
	required_input = v

	needs_to_be_padded = len(required_input) != max_length

	if needs_to_be_padded:
	difference = max_length - len(required_input)
	batch_outputs[k][i] = [self.pad_token_id] * difference + required_input

	max_num_rels = 0
	for rel in segment_rel_pack:
	max_num_rels = max(max_num_rels, rel.shape[0])
	padded_rels = np.zeros((len(segment_rel_pack), max_num_rels), dtype=np.int32)
	for i, rel in enumerate(segment_rel_pack):
	padded_rels[i, : rel.shape[0]] = rel
	batch_outputs["segment_rel"] = padded_rels
	batch_outputs["batch_ext_table_ids"] = np.array(batch_ext_table_ids, dtype=np.int32)
	batch_outputs["batch_ext_table_sub"] = np.array(batch_ext_table_sub, dtype=np.int32)
	batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors)
	if return_tensors == "pt":
	batch_outputs = batch_outputs.to(device=device)
	batch_outputs["other_info"] = other_info

	return batch_outputs

	def convert_data_to_id(
	self,
	data: Any,
	prev_ext_states: Optional[_PrevExtTableStates] = None,
	shuffle_answer: bool = True,
	max_depth: int = 8,
	):
	"""
	Parse a dict to data ids. Exclusive for CPMBee. It will
	1. parse the dict to segments and get segment_rel, which for calculating of position_bias.
	2. tokenize every segment.
	"""
	root: _DictTree = {
	"value": "<root>",
	"children": [],
	"depth": 0,
	"segment_id": 0,
	"need_predict": False,
	}

	segments = [root]

	def _build_dict_tree(data: CPMBeeInputType, depth: int, need_predict: bool) -> List[_DictTree]:
	if isinstance(data, dict):
	ret_list: List[_DictTree] = []
	curr_items = list(data.items())
	if need_predict and shuffle_answer:
	access_idx = np.arange(len(curr_items))
	np.random.shuffle(access_idx)
	curr_items = [curr_items[idx] for idx in access_idx]
	for k, v in curr_items:
	child_info: _DictTree = {
	"value": k,
	"children": [],
	"depth": depth,
	"segment_id": len(segments),
	"need_predict": False, # only leaves are contexts
	}
	segments.append(child_info)
	child_info["children"] = _build_dict_tree(
	v, depth + 1, need_predict or (depth == 1 and k == "<ans>")
	) # elements in <root>.<ans>

	ret_list.append(child_info)
	return ret_list
	else:
	assert isinstance(data, str), "Invalid data {}".format(data)
	ret: _DictTree = {
	"value": data,
	"children": [],
	"depth": depth,
	"segment_id": len(segments),
	"need_predict": need_predict,
	}
	segments.append(ret)
	return [ret]

	root["children"] = _build_dict_tree(data, 1, False)

	num_segments = len(segments)
	segment_rel = np.zeros((num_segments * num_segments,), dtype=np.int32)

	def _build_segment_rel(node: _DictTree) -> List[Tuple[int, int]]:
	ret: List[Tuple[int, int]] = [(node["segment_id"], node["depth"])]
	for child in node["children"]:
	sub = _build_segment_rel(child)
	for seg_id_1, depth_1 in sub:
	for seg_id_2, depth_2 in ret:
	n_up = min(depth_1 - node["depth"], max_depth - 1)
	n_down = min(depth_2 - node["depth"], max_depth - 1)
	segment_rel[seg_id_1 * num_segments + seg_id_2] = rel_to_bucket(
	n_up, n_down, max_depth=max_depth
	)
	segment_rel[seg_id_2 * num_segments + seg_id_1] = rel_to_bucket(
	n_down, n_up, max_depth=max_depth
	)
	ret.extend(sub)
	return ret

	_build_segment_rel(root)

	input_ids: List[int] = []
	input_id_subs: List[int] = []
	segment_bound: List[Tuple[int, int]] = []

	if prev_ext_states is not None:
	self.ext_table = prev_ext_states["ext_table"]
	self.token_id_table = prev_ext_states["token_id_table"]

	for seg in segments:
	# tokenize
	tokens = self.convert_tokens_to_ids(self.tokenize(seg["value"], for_cpmbee=True))

	token_id_subs = []
	reid_token_ids = []
	for idx in tokens:
	if idx in self.ext_table:
	# unk or special token
	token = self.ext_table[idx]
	if token.startswith("<") and token.endswith(">"):
	# special token
	if "_" in token:
	token_name = token[1:-1].split("_", maxsplit=1)[0]
	else:
	token_name = token[1:-1]
	token_name = "<{}>".format(token_name)
	else:
	token_name = "<unk>"

	if token_name not in self.token_id_table:
	self.token_id_table[token_name] = {}
	if idx not in self.token_id_table[token_name]:
	self.token_id_table[token_name][idx] = len(self.token_id_table[token_name])
	if token_name not in self.encoder:
	raise ValueError("Invalid token {}".format(token))
	reid_token_ids.append(self.encoder[token_name])
	token_id_subs.append(self.token_id_table[token_name][idx])
	else:
	reid_token_ids.append(idx)
	token_id_subs.append(0)
	tokens = [self.bos_token_id] + reid_token_ids
	token_id_subs = [0] + token_id_subs
	# eos_id 表示 no need_predict
	if not seg["need_predict"]: # eos
	tokens = tokens + [self.eos_token_id]
	token_id_subs = token_id_subs + [0]
	else:
	# no eos
	pass
	begin = len(input_ids)
	input_ids.extend(tokens)
	input_id_subs.extend(token_id_subs)
	end = len(input_ids)
	segment_bound.append((begin, end))

	ids = np.array(input_ids, dtype=np.int32)
	id_subs = np.array(input_id_subs, dtype=np.int32)
	segs = np.zeros((ids.shape[0],), dtype=np.int32) # 按segment_bound对seg编号
	context = np.zeros((ids.shape[0],), dtype=np.int8)
	for i, (begin, end) in enumerate(segment_bound):
	if not segments[i]["need_predict"]:
	context[begin:end] = 1
	segs[begin:end] = i

	curr_ext_table_states: _PrevExtTableStates = {
	"ext_table": self.ext_table,
	"token_id_table": self.token_id_table,
	}
	return ids, id_subs, context, segs, segment_rel, num_segments, curr_ext_table_states

	def prepare_for_model(
	self,
	ids: List[int],
	pair_ids: Optional[List[int]] = None,
	add_special_tokens: bool = True,
	padding: Union[bool, str, PaddingStrategy] = False,
	truncation: Union[bool, str, TruncationStrategy] = None,
	max_length: Optional[int] = None,
	stride: int = 0,
	pad_to_multiple_of: Optional[int] = None,
	return_tensors: Optional[Union[str, TensorType]] = None,
	return_token_type_ids: Optional[bool] = None,
	return_attention_mask: Optional[bool] = None,
	return_overflowing_tokens: bool = False,
	return_special_tokens_mask: bool = False,
	return_length: bool = False,
	verbose: bool = True,
	prepend_batch_axis: bool = False,
	**kwargs,
	) -> BatchEncoding:
	"""
	Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It
	adds special tokens, truncates sequences if overflowing while taking into account the special tokens and
	manages a moving window (with user defined stride) for overflowing tokens. Please Note, for pair_ids
	different than `None` and truncation_strategy = longest_first or `True`, it is not possible to return
	overflowing tokens. Such a combination of arguments will raise an error.

	Args:
	ids (`List[int]`):
	Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and
	`convert_tokens_to_ids` methods.
	pair_ids (`List[int]`, optional):
	Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize`
	and `convert_tokens_to_ids` methods.
	"""

	# Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
	padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
	padding=padding,
	truncation=truncation,
	max_length=max_length,
	pad_to_multiple_of=pad_to_multiple_of,
	verbose=verbose,
	**kwargs,
	)

	pair = bool(pair_ids is not None)
	len_ids = len(ids)
	len_pair_ids = len(pair_ids) if pair else 0

	if return_token_type_ids and not add_special_tokens:
	raise ValueError(
	"Asking to return token_type_ids while setting add_special_tokens to False "
	"results in an undefined behavior. Please set add_special_tokens to True or "
	"set return_token_type_ids to None."
	)

	if (
	return_overflowing_tokens
	and truncation_strategy == TruncationStrategy.LONGEST_FIRST
	and pair_ids is not None
	):
	raise ValueError(
	"Not possible to return overflowing tokens for pair of sequences with the "
	"`longest_first`. Please select another truncation strategy than `longest_first`, "
	"for instance `only_second` or `only_first`."
	)

	# Load from model defaults
	if return_token_type_ids is None:
	return_token_type_ids = "token_type_ids" in self.model_input_names
	if return_attention_mask is None:
	return_attention_mask = "attention_mask" in self.model_input_names

	encoded_inputs = {}

	# Compute the total size of the returned encodings
	total_len = len_ids + len_pair_ids + (self.num_special_tokens_to_add(pair=pair) if add_special_tokens else 0)

	# Truncation: Handle max sequence length
	overflowing_tokens = []
	if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length:
	ids, pair_ids, overflowing_tokens = self.truncate_sequences(
	ids,
	pair_ids=pair_ids,
	num_tokens_to_remove=total_len - max_length,
	truncation_strategy=truncation_strategy,
	stride=stride,
	)

	if return_overflowing_tokens:
	encoded_inputs["overflowing_tokens"] = overflowing_tokens
	encoded_inputs["num_truncated_tokens"] = total_len - max_length

	# Add special tokens
	if add_special_tokens:
	sequence = self.build_inputs_with_special_tokens(ids, pair_ids)
	token_type_ids = self.create_token_type_ids_from_sequences(ids, pair_ids)
	else:
	sequence = ids + pair_ids if pair else ids
	token_type_ids = [0] * len(ids) + ([0] * len(pair_ids) if pair else [])

	# Build output dictionary
	encoded_inputs["input_ids"] = sequence
	if return_token_type_ids:
	encoded_inputs["token_type_ids"] = token_type_ids
	if return_special_tokens_mask:
	if add_special_tokens:
	encoded_inputs["special_tokens_mask"] = self.get_special_tokens_mask(ids, pair_ids)
	else:
	encoded_inputs["special_tokens_mask"] = [0] * len(sequence)

	# Check lengths
	self._eventual_warn_about_too_long_sequence(encoded_inputs["input_ids"], max_length, verbose)

	# Padding
	if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask:
	encoded_inputs = self.pad(
	encoded_inputs,
	max_length=max_length,
	padding=padding_strategy.value,
	pad_to_multiple_of=pad_to_multiple_of,
	return_attention_mask=return_attention_mask,
	)

	if return_length:
	encoded_inputs["length"] = len(encoded_inputs["input_ids"])

	# for CPMBee, encode all the model arguments
	for arg in self.ext_args_for_model:
	v = kwargs.get(arg, None)
	if v is not None:
	encoded_inputs[arg] = v

	batch_outputs = BatchEncoding(
	encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis
	)

	return batch_outputs

	def prepare_for_finetune(
	self,
	data_list: List[Dict],
	max_length: int = 2048
	):
	_inputs: List[NDArray[np.int32]] = []
	_inputs_sub: List[NDArray[np.int32]] = []
	_context: List[NDArray[np.int8]] = []
	_sample_ids: List[NDArray[np.int32]] = []
	_segments: List[NDArray[np.int32]] = []
	_num_segments: List[NDArray[np.int32]] = []
	_segment_rel_offset: List[NDArray[np.int32]] = []
	_segment_rel: List[NDArray[np.int32]] = []
	_spans: List[List[int]] = []
	_raw_data: List[List[Any]] = []

	raw_data = {}
	for data in data_list:
	(
	input_ids,
	input_id_subs,
	context,
	segment_ids,
	segment_rel,
	n_segments,
	_
	) = self.convert_data_to_id(data)

	input_ids = input_ids[: max_length]
	context = context[: max_length]
	segment_ids = segment_ids[: max_length]
	raw_data["input"] = data
	raw_data["samples"] = []

	sample_ids = np.zeros(input_ids.shape, dtype=np.int32)
	segment_rel_offset = np.zeros(input_ids.shape, dtype=np.int32)
	num_segments = np.full(input_ids.shape, n_segments, dtype=np.int32)

	_inputs.append(input_ids)
	_inputs_sub.append(input_id_subs)
	_context.append(context)
	_sample_ids.append(sample_ids)
	_segments.append(segment_ids)
	_num_segments.append(num_segments)
	_segment_rel_offset.append(segment_rel_offset)
	_segment_rel.append(segment_rel)
	_spans.append([input_ids.shape[0]])
	_raw_data.append([raw_data])

	batch_size = len(_inputs)
	inputs = np.zeros((batch_size, max_length), dtype=np.int32)
	inputs_sub = np.zeros((batch_size, max_length), dtype=np.int32)
	context = np.zeros((batch_size, max_length), dtype=np.int8)
	sample_ids = np.zeros((batch_size, max_length), dtype=np.int32)
	segments = np.zeros((batch_size, max_length), dtype=np.int32)
	num_segments = np.zeros((batch_size, max_length), dtype=np.int32)
	segment_rel_offset = np.zeros((batch_size, max_length), dtype=np.int32)
	tgt = np.full((batch_size, max_length), -100, dtype=np.int32)

	max_rel = 0
	for i in range(batch_size):
	max_rel = max(max_rel, _segment_rel[i].shape[0])
	segment_rel = np.zeros((batch_size, max_rel), dtype=np.int32)
	spans = np.zeros((batch_size, max_length), dtype=np.int32)
	length = np.zeros((batch_size,), dtype=np.int32)

	batch_ext_table_map: Dict[Tuple[int, int], int] = {}
	batch_ext_table_ids: List[int] = []
	batch_ext_table_sub: List[int] = []
	raw_data_list: List[Any] = []

	for i in range(batch_size):
	instance_length = _inputs[i].shape[0]
	rel_size = _segment_rel[i].shape[0]
	inputs[i, :instance_length] = _inputs[i]
	inputs_sub[i, :instance_length] = _inputs_sub[i]
	context[i, :instance_length] = _context[i]
	sample_ids[i, :instance_length] = _sample_ids[i]
	segments[i, :instance_length] = _segments[i]
	num_segments[i, :instance_length] = _num_segments[i]
	segment_rel_offset[i, :instance_length] = _segment_rel_offset[i]
	segment_rel[i, :rel_size] = _segment_rel[i]

	span_begin = 0
	for span_id, span_end in enumerate(_spans[i]):
	spans[i, span_begin:span_end] = span_id
	span_begin = span_end
	length[i] = instance_length
	raw_data_list.extend(_raw_data[i])

	for j in range(instance_length):
	idx, idx_sub = _inputs[i][j], _inputs_sub[i][j]
	tgt_idx = idx
	if idx_sub > 0:
	# need to be in ext table
	if (idx, idx_sub) not in batch_ext_table_map:
	batch_ext_table_map[(idx, idx_sub)] = len(batch_ext_table_map)
	batch_ext_table_ids.append(idx)
	batch_ext_table_sub.append(idx_sub)
	tgt_idx = batch_ext_table_map[(idx, idx_sub)] + self.vocab_size
	if j > 1 and context[i, j - 1] == 0:
	if idx != self.bos_token_id:
	tgt[i, j - 1] = tgt_idx
	else:
	tgt[i, j - 1] = self.eos_token_id
	if context[i, instance_length - 1] == 0:
	tgt[i, instance_length - 1] = self.eos_token_id

	if len(batch_ext_table_map) == 0:
	# placeholder
	batch_ext_table_ids.append(0)
	batch_ext_table_sub.append(1)

	return BatchEncoding({
	"input_ids": inputs,
	"input_id_sub": inputs_sub,
	"length": length,
	"context": context > 0,
	"sample_ids": sample_ids,
	"num_segments": num_segments,
	"segment": segments,
	"segment_rel_offset": segment_rel_offset,
	"segment_rel": segment_rel,
	"span": spans,
	"labels": tgt,
	"ext_table_ids": np.array(batch_ext_table_ids, dtype=np.int32),
	"ext_table_sub": np.array(batch_ext_table_sub, dtype=np.int32)
	}, tensor_type="pt")