tim1900/bert-chunker

1from transformers.modeling_utils import PreTrainedModel2from torch import nn3from transformers.models.bert.configuration_bert import BertConfig4from transformers.models.bert.modeling_bert import BertModel5import torch6import torch.nn.functional as F7class BertChunker(PreTrainedModel):8 9    config_class = BertConfig10 11    def __init__(self, config, ):12        super().__init__(config)13 14        self.model = BertModel(config)15        self.chunklayer = nn.Linear(384, 2)16 17    def forward(self, input_ids=None, attention_mask=None,labels=None, **kwargs):18        model_output = self.model(19            input_ids=input_ids, attention_mask=attention_mask, **kwargs20        )21        token_embeddings = model_output[0]22        logits = self.chunklayer(token_embeddings)23        model_output["logits"]=logits24        loss = None25        logits = logits.contiguous()26        if labels:27            labels = labels.contiguous()28            # Flatten the tokens29            loss_fct = nn.CrossEntropyLoss()#用-10030            # loss_fct = nn.CrossEntropyLoss(ignore_index=50257)31            logits = logits.view(-1, logits.shape[-1])32            labels = labels.view(-1)33            # Enable model parallelism34            labels = labels.to(labels.device)35            loss = loss_fct(logits, labels)36            model_output["loss"]=loss37 38        return model_output39 40    def chunk_text(self, text:str, tokenizer, prob_threshold=0.5)->list[str]:41    # slide context window42        MAX_TOKENS=25543        tokens=tokenizer(text, return_tensors="pt",truncation=False)44        input_ids=tokens['input_ids'].to(self.device)45        attention_mask=tokens['attention_mask'][:,0:MAX_TOKENS]46        attention_mask=attention_mask.to(self.device)47        CLS=input_ids[:,0].unsqueeze(0)48        SEP=input_ids[:,-1].unsqueeze(0)49        input_ids=input_ids[:,1:-1]50        self.eval()51        split_str_poses=[]52 53        windows_start =054        windows_end= 055 56        while windows_end <= input_ids.shape[1]:57            windows_end= windows_start + MAX_TOKENS-258 59            ids=torch.cat((CLS, input_ids[:,windows_start:windows_end],SEP),1)60 61            ids=ids.to(self.device)62            63            output=self(input_ids=ids,attention_mask=torch.ones(1, ids.shape[1],device=self.device))64            logits = output['logits'][:, 1:-1,:]65 66            chunk_probabilities = F.softmax(logits, dim=-1)[:,:,1]67            chunk_decision = (chunk_probabilities>prob_threshold)68            greater_rows_indices = torch.where(chunk_decision)[1].tolist()69 70            # null or not71            if len(greater_rows_indices)>0 and (not (greater_rows_indices[0] == 0 and len(greater_rows_indices)==1)):72 73                split_str_pos=[tokens.token_to_chars(sp + windows_start + 1).start for sp in greater_rows_indices]74 75                split_str_poses += split_str_pos76 77                windows_start = greater_rows_indices[-1] + windows_start78 79            else:80 81                windows_start = windows_end82 83        substrings = [text[i:j] for i, j in zip([0] + split_str_poses, split_str_poses+[len(text)])]84        return substrings85 86    def chunk_text_fast(87        self, text: str, tokenizer, batchsize=20, prob_threshold=0.588    ) -> list[str]:89    # chunk the text faster with a fixed context window, batchsize is the number of windows run per batch.90        self.eval()91 92        split_str_poses=[]93        MAX_TOKENS = 25594        USEFUL_TOKENS = MAX_TOKENS - 2 # delete cls and sep95        tokens = tokenizer(text, return_tensors="pt", truncation=False)96        input_ids = tokens["input_ids"]97 98 99        CLS = tokenizer.cls_token_id100 101        SEP = tokenizer.sep_token_id102 103        input_ids = input_ids[:, 1:-1].squeeze().contiguous()# delete cls and sep104 105        token_num = input_ids.shape[0]106        seq_num = input_ids.shape[0] // (USEFUL_TOKENS)107        left_token_num = input_ids.shape[0] % (USEFUL_TOKENS)108 109        if seq_num > 0:110 111            reshaped_input_ids = input_ids[: seq_num * USEFUL_TOKENS].view( seq_num, USEFUL_TOKENS )112 113            i = torch.arange(seq_num).unsqueeze(1)114            j = torch.arange(USEFUL_TOKENS).repeat(seq_num, 1)115 116            bias = 1 # 1 bias by cls token117            position_id = i * (USEFUL_TOKENS) + j + bias 118            position_id = position_id.to(self.device)119            reshaped_input_ids = torch.cat(120                (121                    torch.full((reshaped_input_ids.shape[0], 1), CLS),122                    reshaped_input_ids,123                    torch.full((reshaped_input_ids.shape[0], 1), SEP),124                ),125                1,126            )127 128            batch_num = seq_num // batchsize129            left_seq_num = seq_num % batchsize130            for i in range(batch_num):131                batch_input = reshaped_input_ids[i : i + batchsize, :].to(self.device)132                attention_mask = torch.ones(batch_input.shape[0], batch_input.shape[1]).to(self.device)133                output = self(input_ids=batch_input, attention_mask=attention_mask)134                logits = output['logits'][:, 1:-1,:]#delete cls and sep135                # is_left_greater = ((logits[:,:, 0] + 0) < logits[:,:, 1])136 137                chunk_probabilities = F.softmax(logits, dim=-1)[:,:,1]138                chunk_decision = (chunk_probabilities>prob_threshold)139                140                pos = chunk_decision * position_id[i : i + batchsize, :]141                pos = pos[pos>0].tolist()142                split_str_poses += [tokens.token_to_chars(p).start for p in pos]143            if left_seq_num > 0:144                batch_input = reshaped_input_ids[-left_seq_num:, :].to(self.device)145                attention_mask = torch.ones(batch_input.shape[0], batch_input.shape[1]).to(self.device)146                output = self(input_ids=batch_input, attention_mask=attention_mask)147                logits = output['logits'][:, 1:-1,:]#delete cls and sep148                chunk_probabilities = F.softmax(logits, dim=-1)[:,:,1]149                chunk_decision = (chunk_probabilities>prob_threshold)150                pos = chunk_decision * position_id[-left_seq_num:, :]151                pos = pos[pos>0].tolist()152                split_str_poses += [tokens.token_to_chars(p).start for p in pos]153 154        if left_token_num > 0:155            left_input_ids = torch.cat([torch.tensor([CLS]), input_ids[-left_token_num:], torch.tensor([SEP])])156            left_input_ids = left_input_ids.unsqueeze(0).to(self.device)157            attention_mask = torch.ones(left_input_ids.shape[0], left_input_ids.shape[1]).to(self.device)158            output = self(input_ids=left_input_ids, attention_mask=attention_mask)159            logits = output['logits'][:, 1:-1,:]#delete cls and sep160            chunk_probabilities = F.softmax(logits, dim=-1)[:,:,1]161            chunk_decision = (chunk_probabilities>prob_threshold)162            bias = token_num - (left_input_ids.shape[1] - 2) + 1163            pos = (torch.where(chunk_decision)[1] + bias).tolist()164            split_str_poses += [tokens.token_to_chars(p).start for p in pos]165     166        substrings = [text[i:j] for i, j in zip([0] + split_str_poses, split_str_poses+[len(text)])]167        return substrings168