Predict file

predict.py

@@ -0,0 +1,172 @@
+import ast
+import pandas as pd
+import joblib
+import nltk
+from nltk import pos_tag
+import string
+from nltk.stem import WordNetLemmatizer
+from nltk.stem import PorterStemmer
+
+
+# Check if nltk modules are downloaded, if not download them
+nltk.download('wordnet')
+nltk.download('omw-1.4')
+nltk.download("averaged_perceptron_tagger")
+
+
+class getsentence(object):   
+    '''
+    This class is used to get the sentences from the dataset.
+    Converts from BIO format to sentences using their sentence numbers
+    '''
+    def __init__(self, data):
+        self.n_sent = 1.0
+        self.data = data
+        self.empty = False
+        self.grouped = self.data.groupby("sentence_num").apply(self._agg_func)
+        self.sentences = [s for s in self.grouped]
+   
+    def _agg_func(self, s):
+        return [(w, p) for w, p in zip(s["token"].values.tolist(),
+                                       s["pos_tag"].values.tolist())]
+
+
+def word2features(sent, i):
+    '''
+    This method is used to extract features from the words in the sentence.
+    The main features extracted are:
+    - word.lower(): The word in lowercase
+    - word.isdigit(): If the word is a digit
+    - word.punct(): If the word is a punctuation
+    - postag: The pos tag of the word
+    - word.lemma(): The lemma of the word
+    - word.stem(): The stem of the word
+    The features (not all) are also extracted for the 4 previous and 4 next words.
+    '''
+    global token_count
+    wordnet_lemmatizer = WordNetLemmatizer()
+    porter_stemmer = PorterStemmer()
+    word = sent[i][0]
+    postag = sent[i][1]
+
+    features = {
+        'bias': 1.0,
+        'word.lower()': word.lower(),
+        'word.isdigit()': word.isdigit(),
+        # Check if its punctuations
+        'word.punct()': word in string.punctuation,
+        'postag': postag,
+        # Lemma of the word
+        'word.lemma()': wordnet_lemmatizer.lemmatize(word),
+        # Stem of the word
+        'word.stem()': porter_stemmer.stem(word)
+    }
+    if i > 0:
+        word1 = sent[i-1][0]
+        postag1 = sent[i-1][1]
+        features.update({
+            '-1:word.lower()': word1.lower(),
+            '-1:word.isdigit()': word1.isdigit(),
+            '-1:word.punct()': word1 in string.punctuation,
+            '-1:postag': postag1
+        })
+        if i - 2 >= 0:
+            features.update({
+                '-2:word.lower()': sent[i-2][0].lower(),
+                '-2:word.isdigit()': sent[i-2][0].isdigit(),
+                '-2:word.punct()': sent[i-2][0] in string.punctuation,
+                '-2:postag': sent[i-2][1]
+            })
+        if i - 3 >= 0:
+            features.update({
+                '-3:word.lower()': sent[i-3][0].lower(),
+                '-3:word.isdigit()': sent[i-3][0].isdigit(),
+                '-3:word.punct()': sent[i-3][0] in string.punctuation,
+                '-3:postag': sent[i-3][1]
+            })
+        if i - 4 >= 0:
+            features.update({
+                '-4:word.lower()': sent[i-4][0].lower(),
+                '-4:word.isdigit()': sent[i-4][0].isdigit(),
+                '-4:word.punct()': sent[i-4][0] in string.punctuation,
+                '-4:postag': sent[i-4][1]
+            })
+    else:
+        features['BOS'] = True
+
+    if i < len(sent)-1:
+        word1 = sent[i+1][0]
+        postag1 = sent[i+1][1]
+        features.update({
+            '+1:word.lower()': word1.lower(),
+            '+1:word.isdigit()': word1.isdigit(),
+            '+1:word.punct()': word1 in string.punctuation,
+            '+1:postag': postag1
+        })
+        if i + 2 < len(sent):
+            features.update({
+                '+2:word.lower()': sent[i+2][0].lower(),
+                '+2:word.isdigit()': sent[i+2][0].isdigit(),
+                '+2:word.punct()': sent[i+2][0] in string.punctuation,
+                '+2:postag': sent[i+2][1]
+            })
+        if i + 3 < len(sent):
+            features.update({
+                '+3:word.lower()': sent[i+3][0].lower(),
+                '+3:word.isdigit()': sent[i+3][0].isdigit(),
+                '+3:word.punct()': sent[i+3][0] in string.punctuation,
+                '+3:postag': sent[i+3][1]
+            })
+        if i + 4 < len(sent):
+            features.update({
+                '+4:word.lower()': sent[i+4][0].lower(),
+                '+4:word.isdigit()': sent[i+4][0].isdigit(),
+                '+4:word.punct()': sent[i+4][0] in string.punctuation,
+                '+4:postag': sent[i+4][1]
+            })
+    else:
+        features['EOS'] = True
+
+    return features
+
+
+def sent2features(sent):
+    '''
+    This method is used to extract features from the sentence.
+    '''
+    return [word2features(sent, i) for i in range(len(sent))]
+
+
+print("Evaluating the model...")
+# Load file from your directory
+df_eval = pd.read_excel("testset_NER_LegalLens.xlsx")
+print("Read the evaluation dataset.")
+df_eval["tokens"] = df_eval["tokens"].apply(ast.literal_eval)
+df_eval['pos_tags'] = df_eval['tokens'].apply(lambda x: [tag[1]
+                                                         for tag in pos_tag(x)])
+data_eval = []
+for i in range(len(df_eval)):
+    for j in range(len(df_eval["tokens"][i])):
+        data_eval.append(
+            {
+                "sentence_num": i+1,
+                "id": df_eval["id"][i],
+                "token": df_eval["tokens"][i][j],
+                "pos_tag": df_eval["pos_tags"][i][j],
+            }
+        )
+data_eval = pd.DataFrame(data_eval)
+print("Dataframe created.")
+getter = getsentence(data_eval)
+sentences_eval = getter.sentences
+X_eval = [sent2features(s) for s in sentences_eval]
+print("Predicting the NER tags...")
+# Load model from your direction
+crf = joblib.load("../models/crf.pkl")
+y_pred_eval = crf.predict(X_eval)
+print("NER tags predicted.")
+df_eval["ner_tags"] = y_pred_eval
+df_eval.drop(columns=["pos_tags"], inplace=True)
+print("Saving the predictions...")
+df_eval.to_csv("predictions_NERLens.csv", index=False)
+print("Predictions saved.")