EntityExtractor.py

#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('cpu')
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