|
- # Copyright (c) Alibaba, Inc. and its affiliates.
-
- from typing import Any, Dict
-
- from modelscope.metainfo import Preprocessors
- from modelscope.preprocessors import Preprocessor
- from modelscope.preprocessors.builder import PREPROCESSORS
- from modelscope.utils.constant import Fields, ModeKeys
- from modelscope.utils.logger import get_logger
-
- logger = get_logger()
-
-
- @PREPROCESSORS.register_module(
- Fields.nlp, module_name=Preprocessors.document_segmentation)
- class DocumentSegmentationTransformersPreprocessor(Preprocessor):
-
- def __init__(self,
- model_dir: str,
- model_max_length: int,
- mode: str = ModeKeys.INFERENCE,
- question_column_name='labels',
- context_column_name='sentences',
- example_id_column_name='example_id',
- label_list=['B-EOP', 'O']):
- """The preprocessor for document segmentation task, based on transformers' tokenizer.
-
- Args:
- model_dir: The model dir containing the essential files to build the tokenizer.
- model_max_length: The max length the model supported.
- mode: The mode for this preprocessor.
- question_column_name: The key for the question column, default `labels`.
- context_column_name: The key for the context column, default `sentences`.
- example_id_column_name: The key for the example id column, default `example_id`.
- label_list: The label list, default `['B-EOP', 'O']`
- """
-
- super().__init__(mode)
- from transformers import BertTokenizerFast
- self.tokenizer = BertTokenizerFast.from_pretrained(model_dir, )
- self.question_column_name = question_column_name
- self.context_column_name = context_column_name
- self.example_id_column_name = example_id_column_name
- self.label_list = label_list
- self.label_to_id = {
- label: id
- for id, label in enumerate(self.label_list)
- }
- self.target_specical_ids = set()
- self.target_specical_ids.add(self.tokenizer.eos_token_id)
- self.max_seq_length = model_max_length
-
- def __call__(self, examples, model_cfg=None) -> Dict[str, Any]:
- questions = examples[self.question_column_name]
- contexts = examples[self.context_column_name]
- example_ids = examples[self.example_id_column_name]
- num_examples = len(questions)
-
- sentences = []
- for sentence_list in contexts:
- sentence_list = [_ + '[EOS]' for _ in sentence_list]
- sentences.append(sentence_list)
-
- try:
- tokenized_examples = self.tokenizer(
- sentences,
- is_split_into_words=True,
- add_special_tokens=False,
- return_token_type_ids=True,
- return_attention_mask=True,
- )
- except Exception as e:
- logger.error(e)
- return {}
-
- segment_ids = []
- token_seq_labels = []
- for example_index in range(num_examples):
- example_input_ids = tokenized_examples['input_ids'][example_index]
- example_labels = questions[example_index]
- example_labels = [
- self.label_to_id[_] if _ in self.label_to_id else -100
- for _ in example_labels
- ]
- example_token_labels = []
- segment_id = []
- cur_seg_id = 1
- para_segment_id = []
- cut_para_seg_id = 1
- for token_index in range(len(example_input_ids)):
- if example_input_ids[token_index] in self.target_specical_ids:
- example_token_labels.append(example_labels[cur_seg_id - 1])
- segment_id.append(cur_seg_id)
- cur_seg_id += 1
- else:
- example_token_labels.append(-100)
- segment_id.append(cur_seg_id)
-
- if example_token_labels[token_index] != -100:
- para_segment_id.append(cut_para_seg_id)
- cut_para_seg_id += 1
- else:
- para_segment_id.append(cut_para_seg_id)
-
- if model_cfg is not None and model_cfg[
- 'type'] == 'ponet' and model_cfg['level'] == 'topic':
- segment_ids.append(para_segment_id)
- else:
- segment_ids.append(segment_id)
- token_seq_labels.append(example_token_labels)
-
- tokenized_examples['segment_ids'] = segment_ids
- tokenized_examples['token_seq_labels'] = token_seq_labels
-
- new_segment_ids = []
- new_token_seq_labels = []
- new_input_ids = []
- new_token_type_ids = []
- new_attention_mask = []
- new_example_ids = []
- new_sentences = []
-
- for example_index in range(num_examples):
- example_input_ids = tokenized_examples['input_ids'][example_index]
- example_token_type_ids = tokenized_examples['token_type_ids'][
- example_index]
- example_attention_mask = tokenized_examples['attention_mask'][
- example_index]
- example_segment_ids = tokenized_examples['segment_ids'][
- example_index]
- example_token_seq_labels = tokenized_examples['token_seq_labels'][
- example_index]
- example_sentences = contexts[example_index]
- example_id = example_ids[example_index]
- example_total_num_sentences = len(questions[example_index])
- example_total_num_tokens = len(
- tokenized_examples['input_ids'][example_index])
- accumulate_length = [
- i for i, x in enumerate(tokenized_examples['input_ids']
- [example_index])
- if x == self.tokenizer.eos_token_id
- ]
- samples_boundary = []
- left_index = 0
- sent_left_index = 0
- sent_i = 0
-
- # for sent_i, length in enumerate(accumulate_length):
- while sent_i < len(accumulate_length):
- length = accumulate_length[sent_i]
- right_index = length + 1
- sent_right_index = sent_i + 1
- if right_index - left_index >= self.max_seq_length - 1 or right_index == example_total_num_tokens:
- samples_boundary.append([left_index, right_index])
-
- sample_input_ids = [
- self.tokenizer.cls_token_id
- ] + example_input_ids[left_index:right_index]
- sample_input_ids = sample_input_ids[:self.max_seq_length]
-
- sample_token_type_ids = [
- 0
- ] + example_token_type_ids[left_index:right_index]
- sample_token_type_ids = sample_token_type_ids[:self.
- max_seq_length]
-
- sample_attention_mask = [
- 1
- ] + example_attention_mask[left_index:right_index]
- sample_attention_mask = sample_attention_mask[:self.
- max_seq_length]
-
- sample_segment_ids = [
- 0
- ] + example_segment_ids[left_index:right_index]
- sample_segment_ids = sample_segment_ids[:self.
- max_seq_length]
-
- sample_token_seq_labels = [
- -100
- ] + example_token_seq_labels[left_index:right_index]
- sample_token_seq_labels = sample_token_seq_labels[:self.
- max_seq_length]
-
- if sent_right_index - 1 == sent_left_index:
- left_index = right_index
- sample_input_ids[-1] = self.tokenizer.eos_token_id
- sample_token_seq_labels[-1] = -100
- else:
- left_index = accumulate_length[sent_i - 1] + 1
- if sample_token_seq_labels[-1] != -100:
- sample_token_seq_labels[-1] = -100
-
- if sent_right_index - 1 == sent_left_index or right_index == example_total_num_tokens:
- sample_sentences = example_sentences[
- sent_left_index:sent_right_index]
- sent_left_index = sent_right_index
- sent_i += 1
- else:
- sample_sentences = example_sentences[
- sent_left_index:sent_right_index - 1]
- sent_left_index = sent_right_index - 1
-
- if (len([_ for _ in sample_token_seq_labels if _ != -100
- ])) != len(sample_sentences) - 1 and (len([
- _
- for _ in sample_token_seq_labels if _ != -100
- ])) != len(sample_sentences):
- tmp = []
- for w_i, w, l in zip(
- sample_input_ids,
- self.tokenizer.decode(sample_input_ids).split(
- ' '), sample_token_seq_labels):
- tmp.append((w_i, w, l))
- while len(sample_input_ids) < self.max_seq_length:
- sample_input_ids.append(self.tokenizer.pad_token_id)
- sample_token_type_ids.append(0)
- sample_attention_mask.append(0)
- sample_segment_ids.append(example_total_num_sentences
- + 1)
- sample_token_seq_labels.append(-100)
-
- new_input_ids.append(sample_input_ids)
- new_token_type_ids.append(sample_token_type_ids)
- new_attention_mask.append(sample_attention_mask)
- new_segment_ids.append(sample_segment_ids)
- new_token_seq_labels.append(sample_token_seq_labels)
- new_example_ids.append(example_id)
- new_sentences.append(sample_sentences)
- else:
- sent_i += 1
- continue
-
- output_samples = {}
-
- output_samples['input_ids'] = new_input_ids
- output_samples['token_type_ids'] = new_token_type_ids
- output_samples['attention_mask'] = new_attention_mask
-
- output_samples['segment_ids'] = new_segment_ids
- output_samples['example_id'] = new_example_ids
- output_samples['labels'] = new_token_seq_labels
- output_samples['sentences'] = new_sentences
-
- return output_samples
|
