#!/usr/bin/env python # encoding: utf-8 import sys import torch from typing import Sequence from .alphabet_atom import AlphabetAtom from .utils import gene_seq_replace class BatchConverter(object): def __init__(self, task_level_type, label_size, output_mode, seq_subword, seq_tokenizer, no_position_embeddings, no_token_type_embeddings, truncation_seq_length: int = None, truncation_matrix_length: int = None, atom_tokenizer: AlphabetAtom = None, atom_truncation_seq_length: int = None, atom_truncation_matrix_length: int = None, ignore_index: int = -100, padding_idx: int = 0, unk_idx: int = 1, cls_idx: int = 2, eos_idx: int = 3, mask_idx: int = 4, non_ignore: bool = False, mlm_probability=0.15, prepend_bos=None, append_eos=None, **kwargs): print("------BatchConverter------") print("BatchConverter, kwargs:") print(kwargs) self.task_level_type = task_level_type self.label_size = label_size self.output_mode = output_mode self.seq_tokenizer = seq_tokenizer self.seq_subword = seq_subword self.ignore_index = ignore_index self.non_ignore = non_ignore self.mlm_probability = mlm_probability self.truncation_seq_length = truncation_seq_length self.truncation_matrix_length = truncation_matrix_length # subword 则必包含两个特殊字符 if prepend_bos is None: if seq_subword is not None: self.prepend_bos = True else: self.prepend_bos = False else: self.prepend_bos = prepend_bos if append_eos is None: if seq_subword is not None: self.append_eos = True else: self.append_eos = False else: self.append_eos = append_eos self.padding_idx = padding_idx self.unk_idx = unk_idx self.cls_idx = cls_idx self.eos_idx = eos_idx self.mask_idx = mask_idx if self.seq_tokenizer is None: self.append_len = 0 else: if hasattr(seq_tokenizer, "prepend_bos"): self.prepend_bos = self.seq_tokenizer.prepend_bos if hasattr(seq_tokenizer, "append_eos"): self.append_eos = self.seq_tokenizer.append_eos if hasattr(seq_tokenizer, "padding_idx"): self.padding_idx = self.seq_tokenizer.padding_idx if hasattr(seq_tokenizer, "unk_idx"): self.unk_idx = self.seq_tokenizer.unk_idx if hasattr(seq_tokenizer, "cls_idx"): self.cls_idx = self.seq_tokenizer.cls_idx if hasattr(seq_tokenizer, "eos_idx"): self.eos_idx = self.seq_tokenizer.eos_idx if hasattr(seq_tokenizer, "mask_idx"): self.mask_idx = self.seq_tokenizer.mask_idx if hasattr(seq_tokenizer, "all_special_token_idx_list"): self.all_special_token_idx_list = self.seq_tokenizer.all_special_token_idx_list else: self.all_special_token_idx_list = [self.padding_idx, self.unk_idx, self.cls_idx, self.eos_idx, self.mask_idx] self.append_len = int(self.prepend_bos) + int(self.append_eos) # for atom self.atom_tokenizer = atom_tokenizer self.atom_truncation_seq_length = atom_truncation_seq_length self.atom_truncation_matrix_length = atom_truncation_matrix_length self.atom_prepend_bos = False self.atom_append_eos = False self.atom_padding_idx = padding_idx self.atom_unk_idx = unk_idx self.atom_cls_idx = cls_idx self.atom_eos_idx = eos_idx self.atom_mask_idx = mask_idx if self.atom_tokenizer is None: self.atom_append_len = 0 else: if hasattr(seq_tokenizer, "padding_idx"): self.padding_idx = self.seq_tokenizer.padding_idx elif hasattr(seq_tokenizer, "pad_idx"): self.padding_idx = self.seq_tokenizer.pad_idx elif hasattr(seq_tokenizer, "pad_token_id"): self.padding_idx = self.seq_tokenizer.pad_token_id if hasattr(seq_tokenizer, "unk_idx"): self.unk_idx = self.seq_tokenizer.unk_idx elif hasattr(seq_tokenizer, "unk_token_id"): self.unk_idx = self.seq_tokenizer.unk_token_id if hasattr(seq_tokenizer, "cls_idx"): self.cls_idx = self.seq_tokenizer.cls_idx elif hasattr(seq_tokenizer, "cls_token_id"): self.cls_idx = self.seq_tokenizer.cls_token_id elif hasattr(seq_tokenizer, "bos_idx"): self.cls_idx = self.seq_tokenizer.bos_idx elif hasattr(seq_tokenizer, "bos_token_id"): self.cls_idx = self.seq_tokenizer.bos_token_id if hasattr(seq_tokenizer, "eos_idx"): self.eos_idx = self.seq_tokenizer.eos_idx elif hasattr(seq_tokenizer, "eos_token_id"): self.eos_idx = self.seq_tokenizer.eos_token_id elif hasattr(seq_tokenizer, "sep_token_id"): self.eos_idx = self.seq_tokenizer.sep_token_id if hasattr(seq_tokenizer, "mask_idx"): self.mask_idx = self.seq_tokenizer.mask_idx elif hasattr(seq_tokenizer, "mask_token_id"): self.mask_idx = self.seq_tokenizer.mask_token_id if hasattr(atom_tokenizer, "all_special_token_idx_list"): self.atom_all_special_token_idx_list = self.atom_tokenizer.all_special_token_idx_list else: self.atom_all_special_token_idx_list = [self.padding_idx, self.unk_idx, self.cls_idx, self.eos_idx, self.mask_idx] self.atom_append_len = int(self.atom_prepend_bos) + int(self.atom_append_eos) print("BatchConverter: prepend_bos=%r, append_eos=%r" % (self.prepend_bos, self.append_eos)) print("BatchConverter: atom_prepend_bos=%r, atom_append_eos=%r" % (self.atom_prepend_bos, self.atom_append_eos)) self.matrix_add_special_token = False if "matrix_add_special_token" in kwargs and kwargs["matrix_add_special_token"]: self.matrix_add_special_token = kwargs["matrix_add_special_token"] if self.matrix_add_special_token: self.prepend_bos = True self.append_eos = True self.atom_prepend_bos = True self.atom_append_eos = True self.append_len = int(self.prepend_bos) + int(self.append_eos) self.atom_append_len = int(self.atom_prepend_bos) + int(self.atom_append_eos) # 减去特殊字符之后的长度 self.truncation_seq_length -= self.append_len self.truncation_matrix_length -= self.append_len # 减去特殊字符之后的长度 if self.atom_truncation_seq_length: self.atom_truncation_seq_length -= self.atom_append_len if self.atom_truncation_matrix_length: self.atom_truncation_matrix_length -= self.atom_append_len self.input_type = None if "input_type" in kwargs and kwargs["input_type"]: self.input_type = kwargs["input_type"] if "max_sentence_length" in kwargs and kwargs["max_sentence_length"]: self.max_sentence_length = kwargs["max_sentence_length"] - self.append_len print("BatchConverter: self.max_sentence_length=%d" % self.max_sentence_length) if atom_tokenizer is not None: self.atom_max_sentence_length = kwargs["max_sentence_length"] - self.atom_append_len print("BatchConverter: self.atom_max_sentence_length=%d" % self.atom_max_sentence_length) if "max_sentences" in kwargs and kwargs["max_sentences"]: self.max_sentences = kwargs["max_sentences"] print("BatchConverter: self.max_sentences=%d" % self.max_sentences) self.trunc_type = "right" if "trunc_type" in kwargs and kwargs["trunc_type"]: self.trunc_type = kwargs["trunc_type"] print("BatchConverter: self.trunc_type=%s" % self.trunc_type) self.no_position_embeddings = no_position_embeddings self.no_token_type_embeddings = no_token_type_embeddings print("BatchConverter: prepend_bos=%r, append_eos=%r" % (self.prepend_bos, self.append_eos)) print("BatchConverter: atom_prepend_bos=%r, atom_append_eos=%r" % (self.atom_prepend_bos, self.atom_append_eos)) print("-" * 50) def __parse_label__(self, max_length, task_level_type, label_size, output_mode, label): if isinstance(label, str): label = eval(label) ''' print("label:") print(label) ''' # 需要是padding长度 cur_len = max_length if task_level_type in ["token_level", "structure_level"]: if output_mode in ["multi_label", "multi-label"]: # N * seq_len * label_size new_label = [] for _ in range(cur_len): tmp = [] for _ in range(label_size): tmp.append(0 if self.non_ignore else self.ignore_index) new_label.append(tmp) else: # N * seq_len new_label = [] for _ in range(cur_len): new_label.append(0 if self.non_ignore else self.ignore_index) if label is not None and len(label) > 0: begin_idx = 0 end_idx = cur_len if self.prepend_bos: begin_idx = 1 if self.append_eos: end_idx = cur_len - 1 for idx, item in enumerate(label): idx += begin_idx if idx >= end_idx: break if output_mode in ["multi_label", "multi-label"]: for v in item: new_label[idx][v] = 1 else: new_label[idx] = item elif task_level_type == "span_level": if output_mode in ["multi_label", "multi-label"]: # N * seq_len * label_size new_label = [] for _ in range(cur_len): tmp = [] for _ in range(label_size): tmp.append(0 if self.non_ignore else self.ignore_index) new_label.append(tmp) else: # N * seq_len new_label = [] for _ in range(cur_len): new_label.append(0 if self.non_ignore else self.ignore_index) if label is not None and len(label) > 0: begin_idx = 0 end_idx = cur_len if self.prepend_bos: begin_idx = 1 if self.append_eos: end_idx = cur_len - 1 for item in label: for idx in range(item[0], item[1] + 1, 1): idx += begin_idx if idx >= end_idx: break if output_mode in ["multi_label", "multi-label"]: new_label[idx][item[2]] = 1 else: new_label[idx] = item[2] elif task_level_type in ["seq_level"]: if output_mode in ["multi_label", "multi-label"]: # N * label_size new_label = [] for _ in range(label_size): new_label.append(0 if self.non_ignore else self.ignore_index) else: # N * 1 new_label = [0 if self.non_ignore else self.ignore_index] if output_mode in ["multi_label", "multi-label"]: if label is not None and len(label) > 0: for v in label: new_label[int(v)] = 1 else: if label is not None and len(str(label)) > 0: if isinstance(label, str): new_label = [int(label)] elif isinstance(label, list): new_label = [int(label[0])] else: new_label = [label] else: raise Exception("Not support task_level_type=%s" % task_level_type) return new_label def __atom_parse_label__(self, max_length, task_level_type, label_size, output_mode, label): if isinstance(label, str): label = eval(label) ''' print("label:") print(label) ''' # 需要是padding长度 cur_len = max_length if task_level_type in ["token_level", "structure_level"]: if output_mode in ["multi_label", "multi-label"]: # N * seq_len * label_size new_label = [] for _ in range(cur_len): tmp = [] for _ in range(label_size): tmp.append(0 if self.non_ignore else self.ignore_index) new_label.append(tmp) else: # N * seq_len new_label = [] for _ in range(cur_len): new_label.append(0 if self.non_ignore else self.ignore_index) if label is not None and len(label) > 0: begin_idx = 0 end_idx = cur_len if self.atom_prepend_bos: begin_idx = 1 if self.atom_append_eos: end_idx = cur_len - 1 for idx, item in enumerate(label): idx += begin_idx if idx >= end_idx: break if output_mode in ["multi_label", "multi-label"]: for v in item: new_label[idx][v] = 1 else: new_label[idx] = item elif task_level_type == "span_level": if output_mode in ["multi_label", "multi-label"]: # N * seq_len * label_size new_label = [] for _ in range(cur_len): tmp = [] for _ in range(label_size): tmp.append(0 if self.non_ignore else self.ignore_index) new_label.append(tmp) else: # N * seq_len new_label = [] for _ in range(cur_len): new_label.append(0 if self.non_ignore else self.ignore_index) if label is not None and len(label) > 0: begin_idx = 0 end_idx = cur_len if self.atom_prepend_bos: begin_idx = 1 if self.atom_append_eos: end_idx = cur_len - 1 for item in label: for idx in range(item[0], item[1] + 1, 1): idx += begin_idx if idx >= end_idx: break if output_mode in ["multi_label", "multi-label"]: new_label[idx][item[2]] = 1 else: new_label[idx] = item[2] elif task_level_type in ["seq_level"]: if output_mode in ["multi_label", "multi-label"]: # N * label_size new_label = [] for _ in range(label_size): new_label.append(0 if self.non_ignore else self.ignore_index) else: # N * 1 new_label = [0 if self.non_ignore else self.ignore_index] if output_mode in ["multi_label", "multi-label"]: if label is not None and len(label) > 0: for v in label: new_label[int(v)] = 1 else: if label is not None and len(str(label)) > 0: if isinstance(label, str): new_label = [int(label)] elif isinstance(label, list): new_label = [int(label[0])] else: new_label = [label] else: raise Exception("Not support task_level_type=%s" % task_level_type) return new_label def __mask_tokens__(self, input_ids): labels = input_ids.clone() probability_matrix = torch.full(labels.shape, self.mlm_probability) # 特殊字符处为1 special_tokens_mask = [ 1 if v in self.all_special_token_idx_list else 0 for v in labels.tolist() ] special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool) # 将特殊字符处填充为0.0 probability_matrix.masked_fill_(special_tokens_mask, value=0.0) # 非特殊字符的位置 masked_indices = torch.bernoulli(probability_matrix).bool() # 特殊字符处为-100 labels[~masked_indices] = self.ignore_index # We only compute loss on masked tokens # 80% of the time, we replace masked input tokens with alphabet.mask_token ([MASK]) indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices input_ids[indices_replaced] = self.mask_idx # 10% of the time, we replace masked input tokens with random word indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced random_words = torch.randint(len(self.seq_tokenizer), labels.shape, dtype=torch.long) input_ids[indices_random] = random_words[indices_random] # The rest of the time (10% of the time) we keep the masked input tokens unchanged return input_ids, labels def __atom_mask_tokens__(self, input_ids): labels = input_ids.clone() probability_matrix = torch.full(labels.shape, self.mlm_probability) # 特殊字符处为1 special_tokens_mask = [ 1 if v in self.atom_all_special_token_idx_list else 0 for v in labels.tolist() ] special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool) # 将特殊字符处填充为0.0 probability_matrix.masked_fill_(special_tokens_mask, value=0.0) # 非特殊字符的位置 masked_indices = torch.bernoulli(probability_matrix).bool() # 特殊字符处为-100 labels[~masked_indices] = self.ignore_index # We only compute loss on masked tokens # 80% of the time, we replace masked input tokens with alphabet.mask_token ([MASK]) indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices input_ids[indices_replaced] = self.atom_mask_idx # 10% of the time, we replace masked input tokens with random word indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced random_words = torch.randint(len(self.atom_tokenizer), labels.shape, dtype=torch.long) input_ids[indices_random] = random_words[indices_random] # The rest of the time (10% of the time) we keep the masked input tokens unchanged return input_ids, labels def __seq_encode__(self, batch_size, seqs): ''' 该函数不加特殊字符[CLS]与[SEP] :param batch_size: :param seqs: :return: ''' if self.seq_subword: seq_encoded_list = [] for seq_str in seqs: seq_to_list = self.seq_subword.process_line(seq_str.upper()).split(" ") seq = " ".join(seq_to_list) inputs = self.seq_tokenizer.encode_plus( seq, None, add_special_tokens=False, max_length=self.truncation_seq_length, truncation=True ) seq_encoded_list.append(inputs["input_ids"]) else: seq_encoded_list = [self.seq_tokenizer.encode(seq_str.upper()) for seq_str in seqs] # 该长度已经减去了需要增加的特殊字符的个数 if self.truncation_seq_length: seq_encoded_list = [encoded[:self.truncation_seq_length] for encoded in seq_encoded_list] max_len = max(len(seq_encoded) for seq_encoded in seq_encoded_list) max_len = max_len + int(self.prepend_bos) + int(self.append_eos) # for input input_ids = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) input_ids.fill_(self.padding_idx) position_ids = None if not self.no_position_embeddings: position_ids = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) position_ids.fill_(self.padding_idx) token_type_ids = None if not self.no_position_embeddings: token_type_ids = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) token_type_ids.fill_(self.padding_idx) attention_masks = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) attention_masks.fill_(0) return seq_encoded_list, input_ids, position_ids, token_type_ids, attention_masks, max_len def __multi_seq_encode__(self, batch_size, seqs): ''' 该函数是多sentence的表征器，每个sentence都加[CLS]与[SEP] :param batch_size: :param seqs: :return: ''' assert hasattr(self, "max_sentences") and hasattr(self, "max_sentence_length") max_sentence_len = 0 max_sentence_num = 0 if self.seq_subword: seq_encoded_list = [] for cur_sample_seqs in seqs: cur_seq_encoded_list = [] if len(cur_sample_seqs) > self.max_sentences: # 每个样本最多cur_sample_seqs个 if self.trunc_type == "left": cur_sample_seqs = cur_sample_seqs[-self.max_sentences:] else: cur_sample_seqs = cur_sample_seqs[:self.max_sentences] if max_sentence_num < len(cur_sample_seqs): max_sentence_num = len(cur_sample_seqs) for seq_idx, seq_str in enumerate(cur_sample_seqs): seq_to_list = self.seq_subword.process_line(seq_str.upper()).split(" ") seq = " ".join(seq_to_list) inputs = self.seq_tokenizer.encode_plus( seq, None, add_special_tokens=False, max_length=self.max_sentence_length, truncation=True ) if self.prepend_bos: inputs["input_ids"] = [self.cls_idx] + inputs["input_ids"] if self.append_eos: inputs["input_ids"] = inputs["input_ids"] + [self.eos_idx] if max_sentence_len < len(inputs["input_ids"]): max_sentence_len = len(inputs["input_ids"]) cur_seq_encoded_list.append(inputs["input_ids"]) seq_encoded_list.append(cur_seq_encoded_list) else: seq_encoded_list = [] for cur_sample_seqs in seqs: cur_seq_encoded_list = [] if len(cur_sample_seqs) > self.max_sentences: # 每个样本最多cur_sample_seqs个 if self.trunc_type == "left": cur_sample_seqs = cur_sample_seqs[-self.max_sentences:] else: cur_sample_seqs = cur_sample_seqs[:self.max_sentences] if max_sentence_num < len(cur_sample_seqs): max_sentence_num = len(cur_sample_seqs) for seq_idx, seq_str in enumerate(cur_sample_seqs): if len(seq_str) > self.max_sentence_length: if self.trunc_type == "left": seq_str = seq_str[-self.max_sentence_length:] else: seq_str = seq_str[:self.max_sentence_length] inputs = self.seq_tokenizer.encode(seq_str.upper()) # print("len:%d, %s" % (len(seq_str), seq_str.upper())) if self.prepend_bos: inputs = [self.cls_idx] + inputs if self.append_eos: inputs = inputs + [self.eos_idx] # print("inputs:%d, " %len(inputs), inputs) cur_seq_encoded_list.append(inputs) if max_sentence_len < len(inputs): max_sentence_len = len(inputs) seq_encoded_list.append(cur_seq_encoded_list) # for input input_ids = torch.empty( ( batch_size, max_sentence_num, max_sentence_len, ), dtype=torch.int64, ) input_ids.fill_(self.padding_idx) position_ids = None if not self.no_position_embeddings: position_ids = torch.empty( ( batch_size, max_sentence_num, max_sentence_len ), dtype=torch.int64, ) position_ids.fill_(self.padding_idx) token_type_ids = None if not self.no_position_embeddings: token_type_ids = torch.empty( ( batch_size, max_sentence_num, max_sentence_len ), dtype=torch.int64, ) token_type_ids.fill_(self.padding_idx) attention_masks = torch.empty( ( batch_size, max_sentence_num, max_sentence_len ), dtype=torch.int64, ) attention_masks.fill_(0) return seq_encoded_list, input_ids, position_ids, token_type_ids, attention_masks, max_sentence_num, max_sentence_len def __atom_seq_encode__(self, batch_size, seqs): ''' 该函数不加特殊字符[CLS]与[SEP] :param batch_size: :param seqs: :return: ''' seq_encoded_list = [] for seq_idx, cur_seq in enumerate(seqs): if isinstance(cur_seq, str): # smiles cur_seq_encoded = self.atom_tokenizer.encode_smi(cur_seq, prepend_bos=False, append_eos=False) elif isinstance(cur_seq, list): # atom list cur_seq_encoded = self.atom_tokenizer.encode(cur_seq, prepend_bos=False, append_eos=False) else: raise Exception("not support molecule input type:", type(cur_seq)) # 该长度已经减去了需要增加的特殊字符的个数 if self.atom_truncation_seq_length: cur_seq_encoded = cur_seq_encoded[:self.atom_truncation_seq_length] seq_encoded_list.append(cur_seq_encoded) max_len = max(len(seq_encoded) for seq_encoded in seq_encoded_list) max_len = max_len + int(self.atom_prepend_bos) + int(self.atom_append_eos) # for input input_ids = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) input_ids.fill_(self.atom_padding_idx) position_ids = None if not self.no_position_embeddings: position_ids = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) position_ids.fill_(self.atom_padding_idx) token_type_ids = None if not self.no_position_embeddings: token_type_ids = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) token_type_ids.fill_(self.atom_padding_idx) attention_masks = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) attention_masks.fill_(0) return seq_encoded_list, input_ids, position_ids, token_type_ids, attention_masks, max_len def __vector_encode__(self, batch_size, vectors): embedding_vector_dim = vectors[0].shape[0] filled_vectors = torch.empty( ( batch_size, embedding_vector_dim ), dtype=torch.float32, ) filled_vectors.fill_(0.0) return filled_vectors, 1 def __atom_vector_encode__(self, batch_size, vectors): return self.__vector_encode__(batch_size, vectors) def __multi_vector_encode__(self, batch_size, vectors): embedding_vector_dim = vectors[0][0].shape[0] filled_vectors = torch.empty( ( batch_size, self.max_sentences, embedding_vector_dim ), dtype=torch.float32, ) filled_vectors.fill_(0.0) return filled_vectors, self.max_sentences, 1 def __matrix_encode__(self, batch_size, matrices): ''' 该函数不加特殊字符[CLS]与[SEP]的向量 :param batch_size: :param matrices: :return: ''' max_len = max(matrix.shape[0] for matrix in matrices) if self.matrix_add_special_token: max_len -= 2 if self.truncation_matrix_length: max_len = min(max_len, self.truncation_matrix_length) if self.matrix_add_special_token: max_len += 2 else: max_len = max_len + int(self.prepend_bos) + int(self.append_eos) embedding_vector_dim = matrices[0].shape[1] # for input filled_matrices = torch.empty( ( batch_size, max_len, embedding_vector_dim ), dtype=torch.float32, ) filled_matrices.fill_(0.0) attention_masks = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) attention_masks.fill_(0) return filled_matrices, attention_masks, max_len def __atom_matrix_encode__(self, batch_size, matrices): ''' 该函数不加特殊字符[CLS]与[SEP]的向量 :param batch_size: :param matrices: :return: ''' max_len = max(matrix.shape[0] for matrix in matrices) if self.matrix_add_special_token: max_len -= 2 if self.atom_truncation_matrix_length: max_len = min(max_len, self.atom_truncation_matrix_length) if self.matrix_add_special_token: max_len += 2 else: max_len = max_len + int(self.atom_prepend_bos) + int(self.atom_append_eos) embedding_vector_dim = matrices[0].shape[1] # for input filled_matrices = torch.empty( ( batch_size, max_len, embedding_vector_dim ), dtype=torch.float32, ) filled_matrices.fill_(0.0) attention_masks = torch.empty( ( batch_size, max_len, ), dtype=torch.int64, ) attention_masks.fill_(0) return filled_matrices, attention_masks, max_len def __multi_matrix_encode__(self, batch_size, matrices): ''' 该函数不加特殊字符[CLS]与[SEP]的向量 :param batch_size: :param matrices: :return: ''' max_sentence_num = max(len(cur_matrix) for cur_matrix in matrices) max_sentence_num = min(max_sentence_num, self.max_sentences) if self.trunc_type == "left": max_sentence_len = max(max(matrix.shape[0] for matrix in cur_matrix[-max_sentence_num:]) for cur_matrix in matrices) else: max_sentence_len = max(max(matrix.shape[0] for matrix in cur_matrix[:max_sentence_num]) for cur_matrix in matrices) # print("encoder max_sentence_num:%d, max_sentence_len: %d" % (max_sentence_num, max_sentence_len)) if self.matrix_add_special_token: max_sentence_len -= 2 max_sentence_len = min(max_sentence_len, self.max_sentence_length) # print("encoder max_sentence_num:%d, max_sentence_len: %d" % (max_sentence_num, max_sentence_len)) if self.matrix_add_special_token: max_sentence_len += 2 else: max_sentence_len = max_sentence_len + int(self.prepend_bos) + int(self.append_eos) # print("encoder max_sentence_num:%d, max_sentence_len: %d" % (max_sentence_num, max_sentence_len)) # print("self.max_sentence_length: %d" % self.max_sentence_length) # print("max_sentence_len: %d" % max_sentence_len) embedding_vector_dim = matrices[0][0].shape[1] # for input filled_matrices = torch.empty( ( batch_size, max_sentence_num, max_sentence_len, embedding_vector_dim ), dtype=torch.float32, ) filled_matrices.fill_(0.0) attention_masks = torch.empty( ( batch_size, max_sentence_num, max_sentence_len ), dtype=torch.int64, ) attention_masks.fill_(0) return filled_matrices, attention_masks, max_sentence_num, max_sentence_len def __call_single__(self, batch_size, seq_types, seqs, vectors, matrices, labels): max_length = sys.maxsize input_ids, position_ids, token_type_ids, seq_attention_masks = None, None, None, None seq_part_of_input = False molecule_flag = False multi_seq_flag = False if seqs: new_seqs = [] for seq_idx, seq_type in enumerate(seq_types): if seq_type == "gene": new_seqs.append(gene_seq_replace(seqs[seq_idx].upper())) elif seq_type == "molecule": if isinstance(seqs[seq_idx], str): new_seqs.append(AlphabetAtom.smiles_2_atom_seq(seqs[seq_idx])) else: new_seqs.append(seqs[seq_idx]) molecule_flag = True elif seq_type == "multi_gene": new_seqs.append([gene_seq_replace(seq).upper() for seq in seqs[seq_idx].split(",")]) multi_seq_flag = True elif seq_type == "multi_prot": new_seqs.append([seq.upper() for seq in seqs[seq_idx].split(",")]) multi_seq_flag = True else: new_seqs.append(seqs[seq_idx].upper()) if molecule_flag: # seq_encoded_list没有加特殊字符，input_ids标志位来占位， seq_max_length 根据标志位来加特殊字符长度 seq_encoded_list, input_ids, position_ids, token_type_ids, seq_attention_masks, seq_max_length = self.__atom_seq_encode__( batch_size=batch_size, seqs=new_seqs) elif multi_seq_flag: # seq_encoded_list根据标志位来加特殊字符，input_ids根据标志位来加特殊字符， seq_max_len 根据标志位来加特殊字符长度 seq_encoded_list, input_ids, position_ids, token_type_ids, seq_attention_masks, seq_max_num, seq_max_len = self.__multi_seq_encode__( batch_size=batch_size, seqs=new_seqs) ''' print("seq_max_num: %d" % seq_max_num) print("seq_max_len: %d" % seq_max_len) print(input_ids.shape) print("len(seq_encoded_list): %d" % len(seq_encoded_list)) for input_id in input_ids: print(len(input_id)) for matrix in input_id: print(matrix.shape) print("*****") ''' else: # seq_encoded_list没有加特殊字符，input_ids标志位来占位， seq_max_length 根据标志位来加特殊字符长度 seq_encoded_list, input_ids, position_ids, token_type_ids, seq_attention_masks, seq_max_length = self.__seq_encode__( batch_size=batch_size, seqs=new_seqs) if multi_seq_flag: max_length = min(max_length, seq_max_num * seq_max_len) else: max_length = min(max_length, seq_max_length) seq_part_of_input = True encoded_vectors = None vector_part_of_input = False if vectors is not None and len(vectors) > 0: if multi_seq_flag: encoded_vectors, vector_max_num, vector_max_len = self.__multi_vector_encode__(batch_size=batch_size, vectors=vectors) elif molecule_flag: encoded_vectors, vector_max_length = self.__atom_vector_encode__(batch_size=batch_size, vectors=vectors) else: encoded_vectors, vector_max_length = self.__vector_encode__(batch_size=batch_size, vectors=vectors) # max_length = min(max_length, vector_max_length) vector_part_of_input = True encoded_matrices, matrix_attention_masks = None, None matrix_part_of_input = False # print("multi_seq_flag:", multi_seq_flag) if matrices is not None and len(matrices) > 0: if multi_seq_flag: # 根据标记位填充，根据标记位填充，句子数量，根据标记位是否加上特殊字符长度 encoded_matrices, matrix_attention_masks, matrix_max_num, matrix_max_len = self.__multi_matrix_encode__( batch_size=batch_size, matrices=matrices) ''' print("matrix_max_num: %d" % matrix_max_num) print("matrix_max_len: %d" % matrix_max_len) print(encoded_matrices.shape) print("len(matrices): %d" % len(matrices)) for matrix_array in matrices: print(len(matrix_array)) for matrix in matrix_array: print(matrix.shape) print("*****") ''' elif molecule_flag: # 根据标记位填充，根据标记位填充，句子数量，根据标记位是否加上特殊字符长度 encoded_matrices, matrix_attention_masks, matrix_max_length = self.__atom_matrix_encode__(batch_size=batch_size, matrices=matrices ) else: # 根据标记位填充，根据标记位填充，句子数量，根据标记位是否加上特殊字符长度 encoded_matrices, matrix_attention_masks, matrix_max_length = self.__matrix_encode__(batch_size=batch_size, matrices=matrices) if multi_seq_flag: max_length = min(max_length, matrix_max_num * matrix_max_len) else: max_length = min(max_length, matrix_max_length) matrix_part_of_input = True has_label = False if labels: has_label = True new_labels = [] num_sentences = 1 sentence_length = 1 for sample_idx in range(batch_size): # seq if seq_part_of_input: if multi_seq_flag: # cls_idx 已经添加 pass elif not molecule_flag and self.prepend_bos: input_ids[sample_idx, 0] = self.cls_idx elif molecule_flag and self.atom_prepend_bos: input_ids[sample_idx, 0] = self.atom_cls_idx seq_encoded = seq_encoded_list[sample_idx] real_seq_len = len(seq_encoded) # seq_tensor = torch.tensor(seq_encoded, dtype=torch.int64) # print("seq_encoded：") # print(seq_encoded) if multi_seq_flag: cur_seq_num = min(len(seq_encoded), seq_max_num) if len(seq_encoded) > cur_seq_num: if self.trunc_type == "left": seq_encoded = seq_encoded[-cur_seq_num:] else: seq_encoded = seq_encoded[cur_seq_num:] if num_sentences < cur_seq_num: num_sentences = cur_seq_num # print("cur_seq_num: %d" % len(seq_encoded)) for seq_idx in range(cur_seq_num): cur_seq = seq_encoded[seq_idx] cur_seq_len = min(len(cur_seq), seq_max_len) ''' print("cur_seq:") print(cur_seq_len) print("input_ids:") print(input_ids.shape) ''' input_ids[sample_idx, seq_idx, :cur_seq_len] = torch.tensor(cur_seq[:cur_seq_len], dtype=torch.int64) seq_attention_masks[sample_idx, seq_idx, :cur_seq_len] = 1 if cur_seq_len > sentence_length: sentence_length = cur_seq_len elif molecule_flag: seq_tensor = torch.tensor(seq_encoded, dtype=torch.int64) input_ids[sample_idx, int(self.atom_prepend_bos): real_seq_len + int(self.atom_prepend_bos)] = seq_tensor cur_sentence_length = int(self.atom_prepend_bos) + real_seq_len + int(self.atom_prepend_bos) if cur_sentence_length > sentence_length: sentence_length = cur_sentence_length else: seq_tensor = torch.tensor(seq_encoded, dtype=torch.int64) input_ids[sample_idx, int(self.prepend_bos): real_seq_len + int(self.prepend_bos)] = seq_tensor cur_sentence_length = int(self.prepend_bos) + real_seq_len + int(self.prepend_bos) if cur_sentence_length > sentence_length: sentence_length = cur_sentence_length if multi_seq_flag: # eos_idx 已经添加 pass elif not molecule_flag and self.append_eos: input_ids[sample_idx, real_seq_len + int(self.prepend_bos)] = self.eos_idx elif molecule_flag and self.atom_append_eos: input_ids[sample_idx, real_seq_len + int(self.atom_prepend_bos)] = self.atom_eos_idx if multi_seq_flag: cur_len = num_sentences * sentence_length elif molecule_flag: cur_len = int(self.atom_prepend_bos) + real_seq_len + int(self.atom_append_eos) else: cur_len = int(self.prepend_bos) + real_seq_len + int(self.append_eos) if not self.no_position_embeddings: if multi_seq_flag: for pos_idx in range(0, cur_len): position_ids[sample_idx, pos_idx//sentence_length, pos_idx % sentence_length] = pos_idx % sentence_length else: for pos_idx in range(0, cur_len): position_ids[sample_idx, pos_idx] = pos_idx if not self.no_token_type_embeddings: seq_type = seq_types[sample_idx] if seq_type == "gene": type_value = 0 else: type_value = 1 if multi_seq_flag: for pos_idx in range(0, cur_len): token_type_ids[sample_idx, pos_idx//sentence_length, pos_idx % sentence_length] = type_value else: for pos_idx in range(0, cur_len): token_type_ids[sample_idx, pos_idx] = type_value if multi_seq_flag: pass else: seq_attention_masks[sample_idx, 0: cur_len] = 1 # vector if vector_part_of_input: if multi_seq_flag: cur_vector_num = min(len(vectors[sample_idx]), vector_max_num) if num_sentences < cur_vector_num: num_sentences = cur_vector_num for vector_idx in range(cur_vector_num): encoded_vectors[sample_idx, vector_idx, :] = torch.tensor(vectors[sample_idx][vector_idx], dtype=torch.float32) else: encoded_vectors[sample_idx, :] = torch.tensor(vectors[sample_idx], dtype=torch.float32) # matrix if matrix_part_of_input: ''' matrix_encoded = matrices[sample_idx] if self.matrix_add_special_token: real_seq_len = matrix_encoded.shape[0] - 2 else: real_seq_len = matrix_encoded.shape[0] if multi_seq_flag: pass elif molecule_flag: # real_seq_len = real_seq_len - int(self.atom_prepend_bos) - int(self.atom_append_eos) real_seq_len = min(real_seq_len, self.atom_truncation_matrix_length) else: # real_seq_len = real_seq_len - int(self.prepend_bos) - int(self.append_eos) real_seq_len = min(real_seq_len, self.truncation_matrix_length) # print("real_seq_len: %d" % real_seq_len) ''' if multi_seq_flag: # 多序列matrix matrix_encoded_list = matrices[sample_idx] cur_matrix_num = min(len(matrix_encoded_list), matrix_max_num) if len(matrix_encoded_list) > cur_matrix_num: if self.trunc_type == "left": matrix_encoded_list = matrix_encoded_list[:cur_matrix_num] else: matrix_encoded_list = matrix_encoded_list[-cur_matrix_num:] if num_sentences < cur_matrix_num: num_sentences = cur_matrix_num # print("matrix_encoded_list: %d" % len(matrix_encoded_list)) for matrix_idx in range(cur_matrix_num): # print("matrix_idx: %d" % matrix_idx) cur_matrix = matrix_encoded_list[matrix_idx] cur_matrix = torch.tensor(cur_matrix, dtype=torch.float32) cur_matrix_len = min(cur_matrix.shape[0], matrix_max_len) if self.matrix_add_special_token: encoded_matrices[sample_idx, matrix_idx, 0: cur_matrix_len - 1] = cur_matrix[0:cur_matrix_len - 1] encoded_matrices[sample_idx, matrix_idx, cur_matrix_len - 1] = cur_matrix[-1] matrix_attention_masks[sample_idx, matrix_idx, 0:cur_matrix_len] = 1 else: encoded_matrices[sample_idx, matrix_idx, int(self.prepend_bos): cur_matrix_len + int(self.prepend_bos)] = cur_matrix[0:cur_matrix_len] matrix_attention_masks[sample_idx, matrix_idx, 0: int(self.prepend_bos) + cur_matrix_len + int(self.append_eos)] = 1 cur_matrix_len = int(self.prepend_bos) + cur_matrix_len + int(self.append_eos) if sentence_length < cur_matrix_len: sentence_length = cur_matrix_len else: matrix_encoded = matrices[sample_idx] if self.matrix_add_special_token: real_seq_len = matrix_encoded.shape[0] - 2 else: real_seq_len = matrix_encoded.shape[0] if molecule_flag: # real_seq_len = real_seq_len - int(self.atom_prepend_bos) - int(self.atom_append_eos) real_seq_len = min(real_seq_len, self.atom_truncation_matrix_length) matrix = torch.tensor(matrix_encoded, dtype=torch.float32) if self.matrix_add_special_token: encoded_matrices[sample_idx, 0: real_seq_len + 2] \ = matrix[0: real_seq_len + 2] matrix_attention_masks[sample_idx, 0: real_seq_len + 2] = 1 cur_sentence_length = real_seq_len + 2 else: encoded_matrices[sample_idx, int(self.atom_prepend_bos): real_seq_len + int(self.atom_prepend_bos)] \ = matrix[0: real_seq_len] # matrix_attention_masks[sample_idx, int(self.atom_prepend_bos): real_seq_len + int(self.atom_prepend_bos)] = 1 matrix_attention_masks[sample_idx, 0: int(self.atom_prepend_bos) + real_seq_len + int(self.atom_append_eos)] = 1 cur_sentence_length = int(self.atom_prepend_bos) + real_seq_len + int(self.atom_prepend_bos) if cur_sentence_length > sentence_length: sentence_length = cur_sentence_length else: # real_seq_len = real_seq_len - int(self.prepend_bos) - int(self.append_eos) real_seq_len = min(real_seq_len, self.truncation_matrix_length) matrix = torch.tensor(matrix_encoded, dtype=torch.float32) if self.matrix_add_special_token: encoded_matrices[sample_idx, 0: real_seq_len + 2] = matrix[0: real_seq_len + 2] matrix_attention_masks[sample_idx, 0: real_seq_len + 2] = 1 cur_sentence_length = real_seq_len + 2 else: encoded_matrices[sample_idx, int(self.prepend_bos): real_seq_len + int(self.prepend_bos)] = matrix[0: real_seq_len] # matrix_attention_masks[sample_idx, int(self.prepend_bos): real_seq_len + int(self.prepend_bos)] = 1 matrix_attention_masks[sample_idx, 0: int(self.prepend_bos) + real_seq_len + int(self.append_eos)] = 1 cur_sentence_length = int(self.prepend_bos) + real_seq_len + int(self.prepend_bos) if cur_sentence_length > sentence_length: sentence_length = cur_sentence_length if has_label: if multi_seq_flag: # to do new_labels.append( self.__parse_label__(max_length, self.task_level_type, self.label_size, self.output_mode, labels[sample_idx])) elif molecule_flag: new_labels.append( self.__atom_parse_label__(max_length, self.task_level_type, self.label_size, self.output_mode, labels[sample_idx])) else: new_labels.append( self.__parse_label__(max_length, self.task_level_type, self.label_size, self.output_mode, labels[sample_idx])) if new_labels is not None and new_labels: if self.output_mode in ["regression"]: labels = torch.tensor(new_labels, dtype=torch.float32) else: labels = torch.tensor(new_labels, dtype=torch.int64) else: labels = None ''' print(input_ids.shape) print("encoded_matrices:") print(encoded_matrices.shape) print("num_sentences:%d" % num_sentences) print("sentence_length:%d" % sentence_length) if labels is not None: print("labels:") print(labels.shape) ''' if multi_seq_flag: if seq_part_of_input: input_ids = torch.reshape(input_ids, (input_ids.shape[0], -1)) if matrix_part_of_input: encoded_matrices = torch.reshape(encoded_matrices, (encoded_matrices.shape[0], -1, encoded_matrices.shape[-1])) if position_ids is not None: position_ids = torch.reshape(position_ids, (position_ids.shape[0], -1)) if token_type_ids is not None: token_type_ids = torch.reshape(token_type_ids, (token_type_ids.shape[0], -1)) if seq_attention_masks is not None: seq_attention_masks = torch.reshape(seq_attention_masks, (seq_attention_masks.shape[0], -1)) if matrix_attention_masks is not None: matrix_attention_masks = torch.reshape(matrix_attention_masks, (matrix_attention_masks.shape[0], -1)) ''' print(input_ids.shape) print("encoded_matrices:") print(encoded_matrices.shape) print("num_sentences:%d" % num_sentences) print("sentence_length:%d" % sentence_length) if labels is not None: print("labels:") print(labels.shape) print("-" * 50) ''' return input_ids, \ position_ids, \ token_type_ids, \ seq_attention_masks, \ encoded_vectors, \ encoded_matrices, \ matrix_attention_masks, \ num_sentences, \ sentence_length, \ labels def __call__(self, raw_batch: Sequence[dict]): batch_size = len(raw_batch) # pair if "seq_id_a" in raw_batch[0] and "seq_id_b" in raw_batch[0]: res = {} # seq_ids_a = [] seq_types_a = [] seqs_a = [] vectors_a = [] matrices_a = [] # seq_ids_b = [] seq_types_b = [] seqs_b = [] vectors_b = [] matrices_b = [] labels = [] for item in raw_batch: # seq_ids_a.append(item["seq_id_a"]) seq_types_a.append(item["seq_type_a"]) if item["seq_a"] is not None: seqs_a.append(item["seq_a"]) if item["vector_a"] is not None: vectors_a.append(item["vector_a"]) if item["matrix_a"] is not None: matrices_a.append(item["matrix_a"]) # seq_ids_b.append(item["seq_id_b"]) seq_types_b.append(item["seq_type_b"]) if item["seq_b"] is not None: seqs_b.append(item["seq_b"]) if item["vector_b"] is not None: vectors_b.append(item["vector_b"]) if item["matrix_b"] is not None: matrices_b.append(item["matrix_b"]) if "label" in item and item["label"] is not None: labels.append(item["label"]) input_ids_a, position_ids_a, token_type_ids_a, seq_attention_masks_a, encoded_vectors_a, encoded_matrices_a, matrix_attention_masks_a, num_sentences_a, sentence_length_a, labels \ = self.__call_single__(batch_size, seq_types_a, seqs_a, vectors_a, matrices_a, labels) if not hasattr(self, "max_sentences") or self.max_sentences is None: res.update({ "input_ids_a": input_ids_a, "position_ids_a": position_ids_a, "token_type_ids_a": token_type_ids_a, "seq_attention_masks_a": seq_attention_masks_a, "vectors_a": encoded_vectors_a, "matrices_a": encoded_matrices_a, "matrix_attention_masks_a": matrix_attention_masks_a, "labels": labels if labels is not None and len(labels) > 0 else None }) else: res.update({ "input_ids_a": input_ids_a, "position_ids_a": position_ids_a, "token_type_ids_a": token_type_ids_a, "seq_attention_masks_a": seq_attention_masks_a, "vectors_a": encoded_vectors_a, "matrices_a": encoded_matrices_a, "matrix_attention_masks_a": matrix_attention_masks_a, "num_sentences_a": num_sentences_a, "sentence_length_a": sentence_length_a, "labels": labels if labels is not None and len(labels) > 0 else None }) input_ids_b, position_ids_b, token_type_ids_b, seq_attention_masks_b, encoded_vectors_b, encoded_matrices_b, matrix_attention_masks_b, num_sentences_b, sentence_length_b, _ \ = self.__call_single__(batch_size, seq_types_b, seqs_b, vectors_b, matrices_b, labels=None) if not hasattr(self, "max_sentences") or self.max_sentences is None: res.update({ "input_ids_b": input_ids_b, "position_ids_b": position_ids_b, "token_type_ids_b": token_type_ids_b, "seq_attention_masks_b": seq_attention_masks_b, "vectors_b": encoded_vectors_b, "matrices_b": encoded_matrices_b, "matrix_attention_masks_b": matrix_attention_masks_b }) else: res.update({ "input_ids_b": input_ids_b, "position_ids_b": position_ids_b, "token_type_ids_b": token_type_ids_b, "seq_attention_masks_b": seq_attention_masks_b, "vectors_b": encoded_vectors_b, "matrices_b": encoded_matrices_b, "num_sentences_b": num_sentences_b, "sentence_length_b": sentence_length_b, "matrix_attention_masks_b": matrix_attention_masks_b }) return res else: res = {} # seq_ids = [] seq_types = [] seqs = [] vectors = [] matrices = [] labels = [] for item in raw_batch: # seq_ids.append(item["seq_id"]) seq_types.append(item["seq_type"]) if item["seq"] is not None: seqs.append(item["seq"]) if item["vector"] is not None: vectors.append(item["vector"]) if item["matrix"] is not None: matrices.append(item["matrix"]) if item["label"] is not None: labels.append(item["label"]) ''' print("seqs:") print(seqs) print([len(seq) for seq in seqs]) print("matrices:") print(matrices) print([matrix.shape for matrix in matrices]) print("labels:") print(labels) print([len(eval(label)) for label in labels]) ''' input_ids, position_ids, token_type_ids, seq_attention_masks, encoded_vectors, encoded_matrices, matrix_attention_masks, num_sentences, sentence_length, labels = self.__call_single__( batch_size, seq_types, seqs, vectors, matrices, labels=labels) if not hasattr(self, "max_sentences") or self.max_sentences is None: res.update({ "input_ids": input_ids, "position_ids": position_ids, "token_type_ids": token_type_ids, "seq_attention_masks": seq_attention_masks, "vectors": encoded_vectors, "matrices": encoded_matrices, "matrix_attention_masks": matrix_attention_masks, "labels": labels if labels is not None and len(labels) > 0 else None }) else: res.update({ "input_ids": input_ids, "position_ids": position_ids, "token_type_ids": token_type_ids, "seq_attention_masks": seq_attention_masks, "vectors": encoded_vectors, "matrices": encoded_matrices, "matrix_attention_masks": matrix_attention_masks, "num_sentences": num_sentences, "sentence_length": sentence_length, "labels": labels if labels is not None and len(labels) > 0 else None }) ''' for item in res.items(): key_name = item[0] print(key_name, ":") if item[1] is not None: print(item[1]) print(item[1].shape) else: print("None") ''' return res