DataReader.py

from keras.preprocessing import sequence
import numpy as np
import nltk
from nltk.tokenize import sent_tokenize

import _pickle as cPickle

class DataReader:

    def __init__ (self, input_resources_pickle_file =None):
        # Some constants
        self.num_classes = 0
        self.num_embedding_features = 0
        self.max_sentence_len_train =0
        
        # Other stuff
        self.wordvecs = None
        self.word_to_ix_map = {}                
        self.n_sentences_all = 0
        self.tag_to_vector_map = {}
        self.vector_to_tag_map = {}                
        

        if not (input_resources_pickle_file is None):
            self.load_resources_pickle_file (input_resources_pickle_file)                 

    ##################################################
    # decode_prediction_sequence
    ##################################################
    def decode_prediction_sequence (self, pred_seq):
        
        pred_tags = []
        for class_prs in pred_seq:
            class_vec = np.zeros(self.num_classes, dtype=np.int32)
            class_vec[np.argmax(class_prs)] = 1
            if tuple(class_vec.tolist()) in self.vector_to_tag_map:
                pred_tags.append(self.vector_to_tag_map[tuple(class_vec.tolist())])
            else:
                print(tuple(class_vec.tolist()))
        return pred_tags
    
    ##################################################
    # load_resources_pickle_file
    ##################################################
    def load_resources_pickle_file (self, input_resources_pickle_file):
        
        print("Loading the resources pickle file {}".format(input_resources_pickle_file))
        
        with open(input_resources_pickle_file, 'rb') as f:           
            pickle_content = cPickle.load(f, encoding='bytes')   

        self.word_to_ix_map = pickle_content["word_to_ix_map"]
        self.wordvecs = pickle_content["wordvecs"]
        self.num_embedding_features = pickle_content["num_embedding_features"] 
        self.num_classes = pickle_content["num_classes"] 
        self.max_sentence_len_train = pickle_content["max_sentence_len_train"] 
        self.tag_to_vector_map = pickle_content["tag_to_vector_map"]
        self.vector_to_tag_map = pickle_content["vector_to_tag_map"]
        self.zero_vec_pos = pickle_content["zero_vec_pos"] 

    ##################################################
    # load_embedding_lookup_table
    ##################################################
    def load_embedding_lookup_table (self, embeddings_file):
        
        ###Load the Word2Vec Model###
        print("Loading the W2V model from file {}".format(embeddings_file))
        #W2V_model = cPickle.load(open(embeddings_file, "rb"))
        with open(embeddings_file, 'rb') as f:           
            W2V_model = cPickle.load(f, encoding='bytes')                     
            
        vocab = list(W2V_model.keys())       
        
        self.word_to_ix_map = {}
        self.wordvecs = []
        
        ###Create LookUp Table for words and their word vectors###
        
        print("Creating the lookup table")
        for index, word in enumerate(vocab):
            self.word_to_ix_map[word] = index            
            self.wordvecs.append(W2V_model[vocab[index]])            
            
           
        self.wordvecs = np.array(self.wordvecs)
        print("Number of entries in the lookup table = {}".format(len(self.wordvecs)))
        self.num_embedding_features = len(self.wordvecs[0])
        print("embedding size = {}".format(self.num_embedding_features))
        
        # Add a zero vector for the Paddings
        self.wordvecs = np.vstack((self.wordvecs, np.zeros(self.num_embedding_features)))
        self.zero_vec_pos = self.wordvecs.shape[0] - 1
        
        print("Done")
        return (self.wordvecs)    
    
    
    ##################################################
    ##  read_and_parse_training_data  
    ##################################################
    def read_and_parse_training_data (self, train_file, output_resources_pickle_file):
        
        print("Loading the training data from file {}".format(train_file))
        with open(train_file, 'r') as f_train:            
            
            self.tag_to_vector_map = {}    # For storing one hot vector notation for each Tag
            self.vector_to_tag_map = {} 
            self.num_classes = 0            # Used to put 1 in the one hot vector notation
            raw_data_train = []
            raw_words_train = []
            raw_tags_train = []        

            # Process all lines in the file
            for line in f_train:
                line = line.strip()
                if not line:
                    raw_data_train.append( (tuple(raw_words_train), tuple(raw_tags_train)))
                    raw_words_train = []
                    raw_tags_train = []
                    continue
                
                word, tag = line.split('\t')
                
                raw_words_train.append(word)
                raw_tags_train.append(tag)
                
                if tag not in self.tag_to_vector_map:                    
                    self.tag_to_vector_map[tag] = None                    
                    self.num_classes += 1
                    
        print("number of training examples = " + str(len(raw_data_train)))               
        
        all_tags = sorted(list(self.tag_to_vector_map.keys()))

        for tag_class_id, tag in enumerate(all_tags):
            one_hot_vec = np.zeros(self.num_classes +1, dtype=np.int32)
            one_hot_vec[tag_class_id] = 1
            self.tag_to_vector_map[tag] = tuple(one_hot_vec)
            self.vector_to_tag_map[tuple(one_hot_vec)] = tag

         #Adding a None Tag        
        one_hot_vec = np.zeros(self.num_classes +1, dtype=np.int32)
        one_hot_vec[self.num_classes] = 1
        self.tag_to_vector_map['NONE'] = tuple(one_hot_vec)
        self.vector_to_tag_map[tuple(one_hot_vec)] = 'NONE'
        self.num_classes += 1

        self.n_sentences_all = len(raw_data_train)

        # Find the maximum sequence length for Training data
        self.max_sentence_len_train = 0
        for seq in raw_data_train:
            if len(seq[0]) > self.max_sentence_len_train:
                self.max_sentence_len_train = len(seq[0])                
        
         ############## Create Train Vectors################
        all_X_train, all_Y_train = [], []        
        unk_words = []
        count = 0
        for word_seq, tag_seq in raw_data_train:  
            
            elem_wordvecs, elem_tags = [], []            
            for ix in range(len(word_seq)):
                w = word_seq[ix]
                t = tag_seq[ix]
                w = w.lower()
                if w in self.word_to_ix_map :
                    count += 1
                    elem_wordvecs.append(self.word_to_ix_map[w])
                    elem_tags.append(self.tag_to_vector_map[t])

                elif "UNK" in self.word_to_ix_map :
                    unk_words.append(w)
                    elem_wordvecs.append(self.word_to_ix_map["UNK"])
                    elem_tags.append(self.tag_to_vector_map[t])
                
                else:
                    unk_words.append(w)
                    w = "UNK"       
                    new_wv = 2 * np.random.randn(self.num_embedding_features) - 1 # sample from normal distribution
                    norm_const = np.linalg.norm(new_wv)
                    new_wv /= norm_const
                    self.wordvecs = np.vstack((self.wordvecs, new_wv))
                    self.word_to_ix_map[w] = self.wordvecs.shape[0] - 1
                    elem_wordvecs.append(self.word_to_ix_map[w])
                    elem_tags.append(list(self.tag_to_vector_map[t]))

            
            # Pad the sequences for missing entries to make them all the same length
            nil_X = self.zero_vec_pos
            nil_Y = np.array(self.tag_to_vector_map['NONE'])
            pad_length = self.max_sentence_len_train - len(elem_wordvecs)
            all_X_train.append( ((pad_length)*[nil_X]) + elem_wordvecs)
            all_Y_train.append( ((pad_length)*[nil_Y]) + elem_tags)

        all_X_train = np.array(all_X_train)
        all_Y_train = np.array(all_Y_train)
        
        print("UNK WORD COUNT = " + str(len(unk_words)))
        print("Found WORDS COUNT = " + str(count))
        print("TOTAL WORDS COUNT= " + str(count+len(unk_words)))    

        self.save_resources(output_resources_pickle_file)
       
        print("Done")
        
        return (all_X_train, all_Y_train)

    ##################################################
    # save_resources
    ##################################################
    def save_resources(self, output_resources_pickle_file):
        print("saving the resources into the file {}".format(output_resources_pickle_file))
        pickle_content = {}       
        pickle_content["word_to_ix_map"] = self.word_to_ix_map
        pickle_content["wordvecs"] = self.wordvecs
        pickle_content["num_embedding_features"] = self.num_embedding_features
        pickle_content["num_classes"] = self.num_classes
        pickle_content["max_sentence_len_train"] = self.max_sentence_len_train
        pickle_content["tag_to_vector_map"] = self.tag_to_vector_map
        pickle_content["vector_to_tag_map"] = self.vector_to_tag_map
        pickle_content["zero_vec_pos"] = self.zero_vec_pos
        
        cPickle.dump(pickle_content, open(output_resources_pickle_file, "wb"))
        print("Done")

    ################################################## 
    #  read_and_parse_test_data 
    ################################################## 
    def read_and_parse_test_data (self, test_file):       
        
        print("Loading test data from file {}".format(test_file))
        with open(test_file, 'r') as f_test:                        
            
            data_set = []            
            sentence_words = []
            sentence_tags = []        

            # Process all lines in the file
            for line in f_test:
                line = line.strip()
                if not line:
                    data_set.append( (tuple(sentence_words), tuple(sentence_tags)))
                    sentence_words = []
                    sentence_tags = []
                    continue
                
                word, tag = line.split('\t')                 
                
                sentence_words.append(word)
                sentence_tags.append(tag)                
                                    
        print("number of test examples = " + str(len(data_set)))   
        self.n_sentences_all = len(data_set)
    
        #Create TEST feature vectors
        all_X_test, all_Y_test = [], []
        num_tokens_list = []
        unk_words = []
        count = 0

        for word_seq, tag_seq in data_set:              
            if len(word_seq) > self.max_sentence_len_train:
               print("skip the extra words in the long sentence")
               word_seq = word_seq[:self.max_sentence_len_train]
               tag_seq = tag_seq[:self.max_sentence_len_train]

            elem_wordvecs, elem_tags = [], []            
            for ix in range(len(word_seq)):
                w = word_seq[ix]
                w = w.lower()

                t = tag_seq[ix]

                #ignore the word if it has uncovered ground truth entity type
                if not (t in self.tag_to_vector_map):
                    continue
                
                if w in self.word_to_ix_map:
                    count += 1
                    elem_wordvecs.append(self.word_to_ix_map[w])
                    elem_tags.append(self.tag_to_vector_map[t])
                    
                elif "UNK" in self.word_to_ix_map :
                    unk_words.append(w)
                    elem_wordvecs.append(self.word_to_ix_map["UNK"])
                    elem_tags.append(self.tag_to_vector_map[t])
                    
                else:
                    unk_words.append(w)
                    w = "UNK"
                    self.word_to_ix_map[w] = self.wordvecs.shape[0] - 1
                    elem_wordvecs.append(self.word_to_ix_map[w])
                    elem_tags.append(self.tag_to_vector_map[t])
                
            # Pad the sequences for missing entries to make all the sentences the same length
            nil_X = self.zero_vec_pos
            nil_Y = np.array(self.tag_to_vector_map['NONE'])
            num_tokens_list.append(len(elem_wordvecs))
            pad_length = self.max_sentence_len_train - len(elem_wordvecs)
            all_X_test.append( ((pad_length)*[nil_X]) + elem_wordvecs)
            all_Y_test.append( ((pad_length)*[nil_Y]) + elem_tags)

        all_X_test = np.array(all_X_test)
        all_Y_test = np.array(all_Y_test)
        
        print("UNK WORD COUNT = " + str(len(unk_words)))
        print("Found WORDS COUNT = " + str(count))
        print("TOTAL WORDS COUNT = " + str(count+len(unk_words)))         
        
        print("Done")
        
        return (all_X_test, all_Y_test, data_set, num_tokens_list)                                         
        
        
     ##################################################
     #  get_feature_vectors_2  
     ################################################## 
    def get_feature_vectors_2 (self, data_file):        

        print("Loading unlabeled data from file {}".format(data_file))
        with open(data_file, 'r') as f_data:                                    
            all_sentences_words = []
                 

            # Process all lines in the file
            for line in f_data:
                text = line.strip()                                

                #break the input text into sentences before tokenization
                sentences = sent_tokenize(text)
                
                for sent in sentences:
                    sentence_words = nltk.word_tokenize(sent)                             
                    all_sentences_words.append( tuple(sentence_words) )                                                           
        
        self.n_sentences_all = len(all_sentences_words)   
        print("number of unlabeled examples = {}".format(self.n_sentences_all))
        return self.create_feature_vectors(all_sentences_words)

    ##################################################
    #  get_feature_vectors_1  
    ################################################## 
    def get_feature_vectors_1 (self, data_list):        

        print("Reading unlabeled data from dataframe")   
        # list of list of tokens
        all_sentences_words = []           

        # Process all lines in the file
        for line in data_list:
            text = line.strip()                                

            #break the input text into sentences before tokenization
            sentences = sent_tokenize(text)
            
            for sent in sentences:
                sentence_words = nltk.word_tokenize(sent)                             
                all_sentences_words.append( tuple(sentence_words) )                                                                                             
        
        self.n_sentences_all = len(all_sentences_words)        
        print("number of unlabeled examples = {}".format(self.n_sentences_all))
        return self.create_feature_vectors(all_sentences_words)

    ################################################## 
    #   create_feature_vectors
    ################################################## 
    def create_feature_vectors(self, all_sentences_words):
        all_X_data = []
        word_seq_list = []
        num_tokens_list = []
        unk_words = []
        count = 0

        for word_seq in all_sentences_words:  
            if len(word_seq) > self.max_sentence_len_train:
                print("skip the extra words in the long sentence")
                word_seq = word_seq[:self.max_sentence_len_train]

            word_seq_list.append(word_seq)

            elem_wordvecs = []            
            for ix in range(len(word_seq)):
                w = word_seq[ix]
                
                w = w.lower()
                if w in self.word_to_ix_map:
                    count += 1
                    elem_wordvecs.append(self.word_to_ix_map[w])                
                    
                elif "UNK" in self.word_to_ix_map :
                    unk_words.append(w)
                    elem_wordvecs.append(self.word_to_ix_map["UNK"])                    
                    
                else:
                    unk_words.append(w)
                    w = "UNK"
                    self.word_to_ix_map[w] = self.wordvecs.shape[0] - 1
                    elem_wordvecs.append(self.word_to_ix_map[w])                    
                
            # Pad the sequences for missing entries to make them all the same length
            nil_X = self.zero_vec_pos
            num_tokens_list.append(len(elem_wordvecs))
            pad_length = self.max_sentence_len_train - len(elem_wordvecs)
            all_X_data.append( ((pad_length)*[nil_X]) + elem_wordvecs)            

        all_X_data = np.array(all_X_data)        
        
        print("UNK WORD COUNT = " + str(len(unk_words)))
        print("Found WORDS COUNT = " + str(count))
        print("TOTAL WORDS = " + str(count+len(unk_words)))         
        
        print("Done")
        
        return (all_X_data, word_seq_list, num_tokens_list)