Upload EntityExtractor.py

EntityExtractor.py

@@ -0,0 +1,280 @@
+#import nlpaug
+#import nlpaug.augmenter.word as naw
+import warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+import nltk
+nltk.download('punkt')
+import pandas as pd
+from nltk import pos_tag
+from nltk.corpus import stopwords
+import string
+
+from gensim.models.phrases import Phrases, Phraser
+import numpy as np
+import re
+from gensim.models import Word2Vec
+import pickle
+import os
+from pathos.multiprocessing import ProcessingPool as Pool
+import itertools
+from time import time
+nltk.download('stopwords')
+#import parmap
+nltk.download('averaged_perceptron_tagger')
+import torch
+device = torch.device('cuda')
+from torch.utils.data import Dataset
+from transformers import BertTokenizer
+from ast import literal_eval
+import os.path
+import os
+from torch.nn.utils import clip_grad_norm_
+from torch.utils.data import DataLoader
+from torch.nn.functional import softmax
+from torch.nn import CrossEntropyLoss
+from torch.optim import Adam
+import time
+from sklearn import metrics
+import statistics
+from transformers import get_linear_schedule_with_warmup
+device = torch.device('cuda')
+import torch
+from torch.utils.data import Dataset
+from transformers import BertTokenizer
+import pandas as pd
+
+from ast import literal_eval
+import os.path
+
+
+
+nltk.download('punkt')
+import pandas as pd
+import string
+
+from gensim.models.phrases import Phrases, Phraser
+#from anytree import Node, RenderTree, PreOrderIter
+
+from pathos.multiprocessing import ProcessingPool as Pool
+import itertools
+from time import time
+import os
+nltk.download('stopwords')
+#import parmap
+from torch.nn.utils import clip_grad_norm_
+from torch.utils.data import DataLoader
+from transformers import get_linear_schedule_with_warmup
+import torch.nn as nn
+
+
+from transformers import *
+
+nltk.download('punkt')
+nltk.download('wordnet')
+nltk.download('omw-1.4')
+
+
+
+device = torch.device('cuda')
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+
+MAX_SEQ_LEN = 256
+
+
+MASK_TOKEN = '[MASK]'
+BATCH_SIZE=32
+
+def generate_production_batch(batch):
+    tok=[(instance.tokens for instance in batch)]
+
+    tok=list( itertools.chain.from_iterable(tok))
+    tok=list( itertools.chain.from_iterable([[' '.join(i)] for i in tok]))
+    encoded = tokenizer.__call__(tok, add_special_tokens=True,
+                                          max_length=MAX_SEQ_LEN, pad_to_max_length=True, 
+                                          return_tensors='pt')
+    input_ids = encoded['input_ids']
+    attn_mask = encoded['attention_mask']
+
+    entity_indices = indices_for_entity_ranges([instance.entity_range for instance in batch])
+
+    return input_ids, attn_mask, entity_indices, batch
+
+
+def indices_for_entity_ranges(ranges):
+    max_e_len = max(end - start for start, end in ranges)
+    indices = torch.tensor([[[min(t, end)] * HIDDEN_OUTPUT_FEATURES
+                             for t in range(start, start + max_e_len + 1)]
+                            for start, end in ranges])
+    return indices
+
+#print(os.getcwd())
+open_file = open("./labels.pkl", "rb")
+LABELS = pickle.load(open_file)
+NUM_CLASSES = len(LABELS) 
+open_file.close()     
+with open('./labels_map.pkl', 'rb') as f:
+    LABEL_MAP = pickle.load(f)
+
+open_file = open("./labels.pkl", "rb")
+LABELS = pickle.load(open_file)
+open_file.close()     
+with open('./labels_map.pkl', 'rb') as f:
+    LABEL_MAP = pickle.load(f)
+
+
+class EntityDataset(Dataset):
+
+    def __init__(self, df, size=None):
+        # filter inapplicable rows
+        self.df = df[df.apply(lambda x: EntityDataset.instance_from_row(x) is not None, axis=1)]
+        print(len(self.df))
+
+        # sample data if a size is specified
+        if size is not None and size < len(self):
+            self.df = self.df.sample(size, replace=False)
+
+    @staticmethod
+    def from_df(df, size=None):
+        dataset = EntityDataset(df, size=size)
+        print('Obtained dataset of size', len(dataset))
+        return dataset
+
+
+    @staticmethod
+    def instance_from_row(row):
+        unpacked_arr = literal_eval(row['entityMentions']) if type(row['entityMentions']) is str else row['entityMentions']
+        entity= unpacked_arr[0]['text']
+        text = row['sentText']
+        return EntityDataset.get_instance(text, entity)
+    @staticmethod
+    def get_instance(text, entity, label=None):
+        tokens = tokenizer.tokenize(text)
+
+        i = 0
+        found_entity = True
+        entity_range = (0,100)
+
+        if found_entity:
+            return PairRelInstance(tokens, entity, entity_range, None, text)
+    def __len__(self):
+        return len(self.df.index)
+
+    def __getitem__(self, idx):
+        return EntityDataset.instance_from_row(self.df.iloc[idx])
+
+
+
+class PairRelInstance:
+
+    def __init__(self, tokens, entity, entity_range, label, text):
+        self.tokens = tokens
+        self.entity = entity
+        self.entity_range = entity_range
+        self.label = label
+        self.text = text
+TRAINED_WEIGHTS = 'bert-base-uncased'
+HIDDEN_OUTPUT_FEATURES = 768
+
+
+
+class PairRelInstance:
+
+    def __init__(self, tokens, entity, entity_range, label, text):
+        self.tokens = tokens
+        self.entity = entity
+        self.entity_range = entity_range
+        self.label = label
+        self.text = text
+
+def input_text_format(text ):
+
+    if text is not None:
+        return text, [{'text': text}]
+
+    return None
+def prep(s):
+      return s.replace('_', ' ').lower()
+class BertEntityExtractor:
+
+    def __init__(self):
+        self.net = EntityBertNet()
+
+    @staticmethod
+    def load_saved(path):
+        extr = BertEntityExtractor()
+        extr.net = EntityBertNet()
+        extr.net.load_state_dict(torch.load(path,map_location=torch.device('cpu'))) #,map_location=torch.device('cpu')
+        extr.net.eval()
+        return extr
+    def load_trained_model():
+        entity_extractor_path = './entity_model2.pt'
+        entity_extractor = BertEntityExtractor.load_saved(entity_extractor_path) 
+        return entity_extractor
+    def input_text(self,texts):
+        mapping1=[input_text_format(texts)]
+        entity_texts = [t for t in mapping1
+                        if t is not None]
+                    
+        df = pd.DataFrame(entity_texts, columns=['sentText', 'entityMentions'])
+        df['sentText']=str(df['sentText'][0])
+        data = EntityDataset.from_df(df)
+        return data,df
+    def extract_entity_probabilities(self, file_path=None, dataset=None, size=None):
+        # load data
+        if file_path is not None:
+            data, _ = EntityDataset.from_file(file_path, size=size)
+        else:
+            if dataset is None:
+                raise AttributeError('file_path and data cannot both be None')
+            data = dataset
+
+        loader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=False, num_workers=4,
+                            collate_fn=generate_production_batch)
+        #print("loader"+str(loader))
+        self.net.to(device)
+        self.net.eval()
+
+        probs =[]
+
+        with torch.no_grad():
+            for input_ids, attn_mask, entity_indices, instances in loader:
+                # send batch to gpu
+                input_ids, attn_mask, entity_indices = tuple(i.to(device) for i in [input_ids, attn_mask,
+                                                                                    entity_indices])
+
+                # forward pass
+                output_scores = softmax(self.net(input_ids, attn_mask, entity_indices), dim=1)
+                for i,(ins, score) in enumerate(zip(instances, output_scores.tolist())):
+                    probs.append(score)
+        return probs
+        #print(probs)
+
+        return {t: statistics.mean(t_probs) if len(t_probs) > 0 else None for t, t_probs in probs.items()}
+
+
+class EntityBertNet(nn.Module):
+
+    def __init__(self):
+        super(EntityBertNet, self).__init__()
+        config = BertConfig.from_pretrained(TRAINED_WEIGHTS)
+        self.bert_base = BertModel.from_pretrained(TRAINED_WEIGHTS, config=config)
+        self.fc = nn.Linear(HIDDEN_OUTPUT_FEATURES, NUM_CLASSES)
+
+    def forward(self, input_ids, attn_mask, entity_indices):
+        # BERT
+        bert_output, _ = self.bert_base(input_ids=input_ids, attention_mask=attn_mask,return_dict=False)
+        # max pooling at entity locations
+        entity_pooled_output = EntityBertNet.pooled_output(bert_output, entity_indices)
+
+        # fc layer (softmax activation done in loss function)
+        x = self.fc(entity_pooled_output)
+        return x
+
+    @staticmethod
+    def pooled_output(bert_output, indices):
+        #print(bert_output)
+        outputs = torch.gather(input=bert_output, dim=1, index=indices)
+        pooled_output, _ = torch.max(outputs, dim=1)
+        return pooled_output
+
+