Class TriConstraintDecoder

extraction/extract_constraint.py:372–473

# ET + Src -> ((Tri)(Tri)), ETText2Tri 使用
class TriConstraintDecoder(ConstraintDecoder):
    def __init__(self, tokenizer, type_schema, *args, **kwargs):
        super().__init__(tokenizer, *args, **kwargs)
        self.tree_end = '<tree-end>'
        self.type_schema = type_schema
        self.type_tree = get_label_name_tree(type_schema.role_list,
                                             tokenizer=self.tokenizer,
                                             end_symbol=self.tree_end)
        self.type_start = self.tokenizer.convert_tokens_to_ids([type_start])[0]
        self.type_end = self.tokenizer.convert_tokens_to_ids([type_end])[0]

    def check_state(self, tgt_generated):
        if tgt_generated[-1] == self.tokenizer.pad_token_id: # t5-base
            return 'start', -1

        special_token_set = {self.type_start, self.type_end}
        special_index_token = list(
            filter(lambda x: x[1] in special_token_set, list(enumerate(tgt_generated))))
        # print(special_index_token)
        last_special_index, last_special_token = special_index_token[-1]

        if len(special_index_token) == 1:
            if last_special_token != self.type_start:
                return 'error', 0

        bracket_position = find_bracket_position(
            tgt_generated, _type_start=self.type_start, _type_end=self.type_end)
        start_number, end_number = len(bracket_position[self.type_start]), len(
            bracket_position[self.type_end]) # 计算左右括号的数量

        if start_number == end_number:
            return 'end_generate', -1
        if start_number == end_number + 1:
            state = 'start_first_generation'
        elif start_number == end_number + 2:
            state = 'generate_span'
        else:
            state = 'error'
        return state, last_special_index


    def get_state_valid_tokens(self, src_sentence, tgt_generated):
        """

        :param src_sentence: ET </s> src </s>
        :param tgt_generated:
        :return:
            List[str], valid token list
        """
        old_src = src_sentence
        if self.tokenizer.eos_token_id in src_sentence:
            if src_sentence.count(self.tokenizer.eos_token_id) > 1: # 有新增的</s>
                first_index = src_sentence.index(self.tokenizer.eos_token_id) # index函数会定位第一个出现的位置
                second_index = first_index + 1 + src_sentence[first_index + 1:].index(self.tokenizer.eos_token_id) # 注意要加上 first_index + 1的偏移
                src_sentence = src_sentence[first_index + 1: second_index] # 输入端 原句 src
                
            else:
                src_sentence = src_sentence[:src_sentence.index(self.tokenizer.eos_token_id)]