Initial commit

predict.py

@@ -0,0 +1,421 @@
+# -*- coding: utf-8 -*-
+
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import BertModel,AutoTokenizer
+import os
+from torch.utils import data
+import numpy as np
+import sys
+
+max_seq_length = 512
+
+
+class InputExample(object):
+
+    def __init__(self, guid, words, labels):
+        self.guid = guid
+        self.words = words
+        self.labels = labels
+
+class InputFeatures(object):
+
+    def __init__(self, input_ids, input_mask, segment_ids,  predict_mask, label_ids):
+        self.input_ids = input_ids
+        self.input_mask = input_mask
+        self.segment_ids = segment_ids
+        self.predict_mask = predict_mask
+        self.label_ids = label_ids
+class DataProcessor(object):
+
+    def get_train_examples(self, data_dir):
+        raise NotImplementedError()
+
+    def get_dev_examples(self, data_dir):
+        raise NotImplementedError()
+    def get_predict_examples(self, data_dir,predict_string):
+        raise NotImplementedError()
+    def get_labels(self):
+        raise NotImplementedError()
+
+    @classmethod
+    def _read_data(cls, input_file,isPredict = False,sentence = ''):
+        if isPredict == False:
+          with open(input_file) as f:
+              out_lists = []
+              entries = f.read().strip().split("\n\n")
+              for entry in entries:
+                  words = []
+                  ner_labels = []
+                  pos_tags = []
+                  bio_pos_tags = []
+                  for line in entry.splitlines():
+                      pieces = line.strip().split()
+                      if len(pieces) < 1:
+                          continue
+                      word = pieces[0]
+                      words.append(word)
+                      ner_labels.append(pieces[-1])
+                  
+                  out_lists.append([words,pos_tags,bio_pos_tags,ner_labels])
+        else:
+          out_lists = []
+          words = []
+          ner_labels = []
+          pos_tags = []
+          bio_pos_tags = []
+          entries = sentence.strip().split(" ")
+          for i in entries:
+            if len(i) < 1:
+                continue
+            word = i
+            words.append(word)
+            ner_labels.append('O')
+          out_lists.append([words,pos_tags,bio_pos_tags,ner_labels])
+        return out_lists
+
+class DNRTIDataProcessor(DataProcessor):
+
+    def __init__(self):
+        self._label_types =  [ 'X', '[CLS]', '[SEP]', 'O', 'B-Area', 'B-Exp', 'B-Features', 'B-HackOrg', 'B-Idus', 'B-OffAct','B-Org', 'B-Purp', 'B-SamFile','B-SecTeam','B-Time','B-Tool','B-Way','I-Area','I-Exp','I-Features','I-HackOrg','I-Idus','I-OffAct','I-Org','I-Purp','I-SamFile','I-SecTeam','I-Time','I-Tool','I-Way']
+        self._num_labels = len(self._label_types)
+        self._label_map = {label: i for i,
+                           label in enumerate(self._label_types)}
+
+    def get_train_examples(self, data_dir):
+        return self._create_examples(
+            self._read_data(os.path.join(data_dir, "train.txt")))
+
+    def get_dev_examples(self, data_dir):
+        return self._create_examples(
+            self._read_data(os.path.join(data_dir, "valid.txt")))
+    def get_predict_examples(self,data_dir, predict_string):
+        return self._create_examples(
+            self._read_data(os.path.join(data_dir, "None.txt"),True,predict_string),True)
+    def get_test_examples(self, data_dir):
+        return self._create_examples(
+            self._read_data(os.path.join(data_dir, "test.txt")))
+
+    def get_labels(self):
+        return self._label_types
+
+    def get_num_labels(self):
+        return self.get_num_labels
+
+    def get_label_map(self):
+        return self._label_map
+
+    def get_start_label_id(self):
+        return self._label_map['[CLS]']
+
+    def get_stop_label_id(self):
+        return self._label_map['[SEP]']
+
+    def _create_examples(self, all_lists,isPredict = False):
+        examples = []
+        if isPredict == False:
+          for (i, one_lists) in enumerate(all_lists):
+            
+            guid = i
+            words = one_lists[0]
+            labels = one_lists[-1]
+
+            examples.append(InputExample(
+                guid=guid, words=words, labels=labels))
+            
+        else:
+          k = 1
+          for i in all_lists:
+
+            guid = k
+            k += 1
+            words = i[0]
+            labels = i[3]
+            examples.append(InputExample(
+                guid=guid, words=words, labels=labels))
+
+        return examples
+
+    def _create_examples2(self, lines):
+        examples = []
+        for (i, line) in enumerate(lines):
+            guid = i
+            text = line[0]
+            ner_label = line[-1]
+            examples.append(InputExample(
+                guid=guid, text_a=text, labels_a=ner_label))
+        return examples
+def example2feature(example, tokenizer, label_map, max_seq_length):
+    add_label = 'X'
+    # tokenize_count = []
+    tokens = ['[CLS]']
+    predict_mask = [0]
+    label_ids = [label_map['[CLS]']]
+    for i, w in enumerate(example.words):
+        sub_words = tokenizer.tokenize(w)
+        if not sub_words:
+            sub_words = ['[UNK]']
+        tokens.extend(sub_words)
+        for j in range(len(sub_words)):
+            if j == 0:
+                predict_mask.append(1)
+                label_ids.append(label_map[example.labels[i]])
+            else:
+                predict_mask.append(0)
+                label_ids.append(label_map[add_label])
+
+    if len(tokens) > max_seq_length - 1:
+        print('Example No.{} is too long, length is {}, truncated to {}!'.format(example.guid, len(tokens), max_seq_length))
+        tokens = tokens[0:(max_seq_length - 1)]
+        predict_mask = predict_mask[0:(max_seq_length - 1)]
+        label_ids = label_ids[0:(max_seq_length - 1)]
+    tokens.append('[SEP]')
+    predict_mask.append(0)
+    label_ids.append(label_map['[SEP]'])
+
+    input_ids = tokenizer.convert_tokens_to_ids(tokens)
+    segment_ids = [0] * len(input_ids)
+    input_mask = [1] * len(input_ids)
+    feat=InputFeatures(
+                # guid=example.guid,
+                # tokens=tokens,
+                input_ids=input_ids,
+                input_mask=input_mask,
+                segment_ids=segment_ids,
+                predict_mask=predict_mask,
+                label_ids=label_ids)
+
+    return feat
+class NerDataset(data.Dataset):
+    def __init__(self, examples, tokenizer, label_map, max_seq_length):
+        self.examples=examples
+        self.tokenizer=tokenizer
+        self.label_map=label_map
+        self.max_seq_length=max_seq_length
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        feat=example2feature(self.examples[idx], self.tokenizer, self.label_map, max_seq_length)
+        return feat.input_ids, feat.input_mask, feat.segment_ids, feat.predict_mask, feat.label_ids
+
+    @classmethod
+    def pad(cls, batch):
+
+        seqlen_list = [len(sample[0]) for sample in batch]
+        maxlen = np.array(seqlen_list).max()
+
+        f = lambda x, seqlen: [sample[x] + [0] * (seqlen - len(sample[x])) for sample in batch]
+        input_ids_list = torch.LongTensor(f(0, maxlen))
+        input_mask_list = torch.LongTensor(f(1, maxlen))
+        segment_ids_list = torch.LongTensor(f(2, maxlen))
+        predict_mask_list = torch.ByteTensor(f(3, maxlen))
+        label_ids_list = torch.LongTensor(f(4, maxlen))
+
+        return input_ids_list, input_mask_list, segment_ids_list, predict_mask_list, label_ids_list
+
+def f1_score(y_true, y_pred):
+    ignore_id=3
+
+    num_proposed = len(y_pred[y_pred>ignore_id])
+    num_correct = (np.logical_and(y_true==y_pred, y_true>ignore_id)).sum()
+    num_gold = len(y_true[y_true>ignore_id])
+
+    try:
+        precision = num_correct / num_proposed
+    except ZeroDivisionError:
+        precision = 1.0
+
+    try:
+        recall = num_correct / num_gold
+    except ZeroDivisionError:
+        recall = 1.0
+
+    try:
+        f1 = 2*precision*recall / (precision + recall)
+    except ZeroDivisionError:
+        if precision*recall==0:
+            f1=1.0
+        else:
+            f1=0
+
+    return precision, recall, f1
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+BertLayerNorm = torch.nn.LayerNorm
+def log_sum_exp_1vec(vec):  
+    max_score = vec[0, np.argmax(vec)]
+    max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
+    return max_score + torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))
+
+def log_sum_exp_mat(log_M, axis=-1): 
+    return torch.max(log_M, axis)[0]+torch.log(torch.exp(log_M-torch.max(log_M, axis)[0][:, None]).sum(axis))
+
+def log_sum_exp_batch(log_Tensor, axis=-1): 
+    return torch.max(log_Tensor, axis)[0]+torch.log(torch.exp(log_Tensor-torch.max(log_Tensor, axis)[0].view(log_Tensor.shape[0],-1,1)).sum(axis))
+
+class BERT_CRF_NER(nn.Module):
+
+    def __init__(self, bert_model, start_label_id, stop_label_id, num_labels, max_seq_length, batch_size, device):
+        super(BERT_CRF_NER, self).__init__()
+        self.hidden_size = 768
+        self.start_label_id = start_label_id
+        self.stop_label_id = stop_label_id
+        self.num_labels = num_labels
+        # self.max_seq_length = max_seq_length
+        self.batch_size = batch_size
+        self.device=device
+        self.bert = bert_model
+        self.dropout = torch.nn.Dropout(0.2)
+        self.hidden2label = nn.Linear(self.hidden_size, self.num_labels)
+        self.transitions = nn.Parameter(torch.randn(self.num_labels, self.num_labels))
+        self.transitions.data[start_label_id, :] = -10000
+        self.transitions.data[:, stop_label_id] = -10000
+
+        nn.init.xavier_uniform_(self.hidden2label.weight)
+        nn.init.constant_(self.hidden2label.bias, 0.0)
+
+
+    def init_bert_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, BertLayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+    def _forward_alg(self, feats):
+        T = feats.shape[1]
+        batch_size = feats.shape[0]
+        log_alpha = torch.Tensor(batch_size, 1, self.num_labels).fill_(-10000.).to(self.device)
+        log_alpha[:, 0, self.start_label_id] = 0
+        for t in range(1, T):
+            log_alpha = (log_sum_exp_batch(self.transitions + log_alpha, axis=-1) + feats[:, t]).unsqueeze(1)
+        log_prob_all_barX = log_sum_exp_batch(log_alpha)
+        return log_prob_all_barX
+
+    def _get_bert_features(self, input_ids, segment_ids, input_mask):
+        bert_seq_out, _ = self.bert(input_ids, token_type_ids=segment_ids, attention_mask=input_mask,return_dict=False)
+        bert_seq_out = self.dropout(bert_seq_out)
+        bert_feats = self.hidden2label(bert_seq_out)
+        return bert_feats
+
+    def _score_sentence(self, feats, label_ids):
+        T = feats.shape[1]
+        batch_size = feats.shape[0]
+
+        batch_transitions = self.transitions.expand(batch_size,self.num_labels,self.num_labels)
+        batch_transitions = batch_transitions.flatten(1)
+
+        score = torch.zeros((feats.shape[0],1)).to(device)
+        for t in range(1, T):
+            score = score + \
+                batch_transitions.gather(-1, (label_ids[:, t]*self.num_labels+label_ids[:, t-1]).view(-1,1)) \
+                    + feats[:, t].gather(-1, label_ids[:, t].view(-1,1)).view(-1,1)
+        return score
+
+    def _viterbi_decode(self, feats):
+        T = feats.shape[1]
+        batch_size = feats.shape[0]
+        log_delta = torch.Tensor(batch_size, 1, self.num_labels).fill_(-10000.).to(self.device)
+        log_delta[:, 0, self.start_label_id] = 0
+        psi = torch.zeros((batch_size, T, self.num_labels), dtype=torch.long).to(self.device)  
+        for t in range(1, T):
+            log_delta, psi[:, t] = torch.max(self.transitions + log_delta, -1)
+            log_delta = (log_delta + feats[:, t]).unsqueeze(1)
+        path = torch.zeros((batch_size, T), dtype=torch.long).to(self.device)
+        max_logLL_allz_allx, path[:, -1] = torch.max(log_delta.squeeze(), -1)
+
+        for t in range(T-2, -1, -1):
+            path[:, t] = psi[:, t+1].gather(-1,path[:, t+1].view(-1,1)).squeeze()
+        return max_logLL_allz_allx, path
+
+    def neg_log_likelihood(self, input_ids, segment_ids, input_mask, label_ids):
+        bert_feats = self._get_bert_features(input_ids, segment_ids, input_mask)
+        forward_score = self._forward_alg(bert_feats)
+        gold_score = self._score_sentence(bert_feats, label_ids)
+        return torch.mean(forward_score - gold_score)
+
+    def forward(self, input_ids, segment_ids, input_mask):
+        bert_feats = self._get_bert_features(input_ids, segment_ids, input_mask)
+        score, label_seq_ids = self._viterbi_decode(bert_feats)
+        return score, label_seq_ids
+
+
+
+bert_model_scale = 'bert-base-cased'
+tokenizer = AutoTokenizer.from_pretrained(bert_model_scale, do_lower_case=True)
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+DNRTIProcessor = DNRTIDataProcessor()
+start_label_id = DNRTIProcessor.get_start_label_id()
+stop_label_id = DNRTIProcessor.get_stop_label_id()
+label_list = DNRTIProcessor.get_labels()
+label_map = DNRTIProcessor.get_label_map()
+batch_size = 512
+bert_model = BertModel.from_pretrained(bert_model_scale)
+model = BERT_CRF_NER(bert_model, start_label_id, stop_label_id, len(label_list), max_seq_length, batch_size, device)
+
+
+checkpoint = torch.load('./outputs/ner_bert_crf_checkpoint.pt', map_location='cpu', weights_only=False)
+epoch = checkpoint['epoch']
+valid_acc_prev = checkpoint['valid_acc']
+valid_f1_prev = checkpoint['valid_f1']
+pretrained_dict=checkpoint['model_state']
+net_state_dict = model.state_dict()
+pretrained_dict_selected = {k: v for k, v in pretrained_dict.items() if k in net_state_dict}
+net_state_dict.update(pretrained_dict_selected)
+model.load_state_dict(net_state_dict)
+print('Loaded the pretrain  NER_BERT_CRF  model, epoch:',checkpoint['epoch'],'valid acc:',
+      checkpoint['valid_acc'], 'valid f1:', checkpoint['valid_f1'])
+
+model.to(device)
+
+
+max_seq_length = 512
+data_dir = 'none.txt'
+
+def predict_sentence(sentence):
+    predict_string = sentence
+    predict_examples = DNRTIProcessor.get_predict_examples(data_dir,predict_string = predict_string)
+    predict_dataset = NerDataset(predict_examples,tokenizer,label_map,max_seq_length)
+
+    # model.eval()
+    with torch.no_grad():
+        demon_dataloader = data.DataLoader(dataset=predict_dataset,
+                                    batch_size=1,
+                                    shuffle=False,
+                                    num_workers=1,
+                                    collate_fn=NerDataset.pad)
+        for batch in demon_dataloader:
+            batch = tuple(t.to(device) for t in batch)
+            input_ids, input_mask, segment_ids, predict_mask, label_ids = batch
+            _, predicted_label_seq_ids = model(input_ids, segment_ids, input_mask)
+            # _, predicted = torch.max(out_scores, -1)
+            valid_predicted = torch.masked_select(predicted_label_seq_ids, predict_mask.bool())
+            # valid_label_ids = torch.masked_select(label_ids, predict_mask)
+            for i in range(1):
+                new_ids=predicted_label_seq_ids[i].cpu().numpy()[predict_mask[i].cpu().numpy()==1]
+                # print(list(map(lambda i: label_list[i], new_ids)))
+                xx = list(map(lambda i: label_list[i], new_ids))
+                result = " ".join(str(x) for x in list(map(lambda i: label_list[i], new_ids)))
+                for sss in xx:
+                    if sss != "O":
+                        pass
+                print("----------------")
+                print("source sentence: ", running_input)
+                print("result sentence: ", result)
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-I', '--input', required=True, type=str, default="", help="Please enter the OSINT or CTI senetences.")
+    arguments = parser.parse_args(sys.argv[1:])
+    running_input = arguments.input
+    if running_input:
+        predict_sentence(running_input)
+    else:
+        print("Error!")