Update tokenization_cpmbee.py

tokenization_cpmbee.py

@@ -18,6 +18,7 @@ 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
@@ -866,3 +867,132 @@ class CpmBeeTokenizer(PreTrainedTokenizer):
         )
 
         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")