Update app.py

app.py

@@ -12,7 +12,6 @@ from tqdm import tqdm
 import gradio as gr
 import matplotlib.pyplot as plt
 from sklearn import metrics
-from sklearn.model_selection import KFold
 
 nltk.download('stopwords')
 nltk.download('punkt_tab')
@@ -22,12 +21,19 @@ PUNCT = set([".", ",", "!", "?", ":", ";", "-", "(", ")", "[", "]", "{", "}", "'
 Features_count = 6
 SEED = 42
 
+SW = set(nltk.corpus.stopwords.words("english"))
+PUNCT = set([".", ",", "!", "?", ":", ";", "-", "(", ")", "[", "]", "{", "}", "'", '"'])
+connectors = set(["of", "in", "and", "for", "to", "with", "at", "from"])
+start_words = set(["the", "a", "an", "this", "that", "these", "those", "my", "your", "his", "her", "its", "our", "their", "few", "many", "several", "all", "most", "some", "any", "every", "each", "either", "neither", "both", "another", "other", "more", "less", "fewer", "little", "much", "great", "good", "bad", "first", "second", "third", "last", "next", "previous"])
+Features_count = 6
+SEED = 42
+
 class NEI:
     def __init__(self):
         self.model = None
         self.scaler = StandardScaler()
         self.vectorizer = DictVectorizer(sparse=True)
-        self.tagset = ['Name[1]', 'No-Name[0]']
+        self.tagset = ['No-Name[0]', 'Name[1]']
     
     def load_dataset(self, file):
         sentences = []
@@ -48,51 +54,6 @@ class NEI:
                 sentences.append(sentence)
         return sentences
 
-    def sent2features(self, sentence):
-        return [self.word2features(sentence, i) for i in range(len(sentence))]
-    
-    def sent2labels(self, sentence):
-        return [label for _, _, label in sentence]
-    
-    def word2features(self, sentence, i):
-        word = sentence[i][0]
-        pos_tag = sentence[i][1]
-        features = {
-            'word': word,
-            'pos_tag': pos_tag,
-            'word.isupper': int(word.isupper()),
-            'word.islower': int(word.islower()),
-            'word.istitle': int(word.istitle()),
-            'word.isdigit': int(word.isdigit()),
-            'word.prefix2': word[:2],
-            'word.prefix3': word[:3],
-            'word.suffix2': word[-2:],
-            'word.suffix3': word[-3:],
-        }
-        # Add context features
-        if i > 0:
-            prv_word = sentence[i - 1][0]
-            prv_pos_tag = sentence[i - 1][1]
-            features.update({
-                '-1:word': prv_word,
-                '-1:pos_tag': prv_pos_tag,
-                '-1:word.isupper': int(prv_word.isupper()),
-                '-1:word.istitle': int(prv_word.istitle()),
-            })
-        else:
-            features['BOS'] = True
-        if i < len(sentence) - 1:
-            next_word = sentence[i + 1][0]
-            next_pos_tag = sentence[i + 1][1]
-            features.update({
-                '+1:word': next_word,
-                '+1:pos_tag': next_pos_tag,
-                '+1:word.isupper': int(next_word.isupper()),
-                '+1:word.istitle': int(next_word.istitle()),
-            })
-        else:
-            features['EOS'] = True
-        return features
 
     def performance(self, y_true, y_pred):
         print(classification_report(y_true, y_pred))
@@ -106,32 +67,53 @@ class NEI:
     def confusion_matrix(self,y_true,y_pred):
         matrix = metrics.confusion_matrix(y_true,y_pred)
         normalized_matrix = matrix/np.sum(matrix, axis=1, keepdims=True)
-        _, ax = plt.subplots()
-        ax.tick_params(top=True)
+
+        # disp = metrics.ConfusionMatrixDisplay(confusion_matrix=normalized_matrix, display_labels=self.tagset)
+        fig, ax = plt.subplots()
+        # disp.plot(cmap=plt.cm.GnBu, ax=ax, colorbar=True)
+        ax.xaxis.set_ticks_position('top')
+        ax.xaxis.set_label_position('top')
         plt.xticks(np.arange(len(self.tagset)), self.tagset)
         plt.yticks(np.arange(len(self.tagset)), self.tagset)
         for i in range(normalized_matrix.shape[0]):
-                for j in range(normalized_matrix.shape[1]):
-                    plt.text(j, i, format(normalized_matrix[i, j], '0.2f'), horizontalalignment="center")
+            for j in range(normalized_matrix.shape[1]):
+                text = f"{normalized_matrix[i, j]:.2f}"
+                ax.text(j, i, text, ha="center", va="center", color="black")
+        plt.title("Normalized Confusion Matrix")
+        plt.xlabel("Predicted Label")
+        plt.ylabel("True Label")
         plt.imshow(normalized_matrix,interpolation='nearest',cmap=plt.cm.GnBu)
         plt.colorbar()
         plt.savefig('Confusion_Matrix.png')
-
-    def vectorize(self, w, scaled_position):
-        title = 1 if w[0].isupper() else 0
-        allcaps = 1 if w.isupper() else 0
-        sw = 1 if w.lower() in SW else 0
-        punct = 1 if w in PUNCT else 0
-        return [title, allcaps, len(w), sw, punct, scaled_position]
+        # plt.xticks(np.arange(len(self.tagset)), self.tagset)
+        # plt.yticks(np.arange(len(self.tagset)), self.tagset)
+        # for i in range(normalized_matrix.shape[0]):
+        #         for j in range(normalized_matrix.shape[1]):
+        #             plt.text(j, i, format(normalized_matrix[i, j], '0.2f'), horizontalalignment="center")
+        # plt.imshow(normalized_matrix,interpolation='nearest',cmap=plt.cm.GnBu)
+        # plt.colorbar()
+        # plt.savefig('Confusion_Matrix.png')
+
+    def vectorize(self, w, scaled_position, prev_tag=0, next_tag=0, prev_token=None):
+        is_titlecase = 1 if w[0].isupper() else 0
+        is_allcaps = 1 if w.isupper() else 0
+        is_sw = 1 if w.lower() in SW else 0
+        is_punct = 1 if w in PUNCT else 0
+        is_surrounded_by_entities = 1 if (prev_tag > 0 and next_tag > 0) else 0
+        is_connector = 1 if (w.lower() in connectors) and (prev_tag > 0 and next_tag > 0) else 0
+        # is_start_of_sentence = 1 if (scaled_position == 0 or prev_token in [".", "!", "?"]) and w.lower() not in start_words else 0
+        # is_start_of_sentence = 1 if scaled_position == 0 else 0
+        return [is_titlecase, is_allcaps, len(w), is_sw, is_punct, is_connector, scaled_position]
 
     def create_data(self, data):
         words, features, labels = [], [], []
         for d in tqdm(data):
             tags = d["ner_tags"]
-            
             tokens = d["tokens"]
             for i, token in enumerate(tokens):
-                x = self.vectorize(token, scaled_position=(i / len(tokens)))
+                prev_tag = tags[i - 1] if i > 0 else 0
+                next_tag = tags[i + 1] if i < len(tokens) - 1 else 0
+                x = self.vectorize(token, scaled_position=(i / len(tokens)), prev_tag=prev_tag, next_tag=next_tag, prev_token=tokens[i-1] if i > 0 else None)
                 y = 1 if tags[i] > 0 else 0
                 features.append(x)
                 labels.append(y)
@@ -150,18 +132,31 @@ class NEI:
         X_val = self.scaler.transform(X_val)
         y_pred_val = self.model.predict(X_val)
         # print(classification_report(y_true=y_val, y_pred=y_pred_val))
+        print(metrics.confusion_matrix(y_val,y_pred_val))
         self.confusion_matrix(y_val,y_pred_val)
         self.performance(y_val,y_pred_val)
 
     def infer(self, sentence):
         tokens = word_tokenize(sentence)
-        features = [self.vectorize(token, i / len(tokens)) for i, token in enumerate(tokens)]
+        features = []
+
+        raw_features = [self.vectorize(token, i / len(tokens), prev_token=tokens[i-1] if i > 0 else None) for i, token in enumerate(tokens)]
+        raw_features = np.array(raw_features, dtype=np.float32)
+        scaled_features = self.scaler.transform(raw_features)
+        y_pred = self.model.predict(scaled_features)
+
+        for i, token in enumerate(tokens):
+            prev_tag = y_pred[i - 1] if i > 0 else 0
+            next_tag = y_pred[i + 1] if i < len(tokens) - 1 else 0
+
+            feature_with_context = self.vectorize(token, i / len(tokens), prev_tag, next_tag, prev_token=tokens[i-1] if i > 0 else None)
+            features.append(feature_with_context)
+
         features = np.array(features, dtype=np.float32)
         scaled_features = self.scaler.transform(features)
         y_pred = self.model.predict(scaled_features)
         return list(zip(tokens, y_pred))
 
-
 data = load_dataset("conll2003", trust_remote_code=True)
 nei_model = NEI()
 
@@ -173,7 +168,7 @@ nei_model.evaluate(data["validation"])
 
 def annotate(text):
     predictions = nei_model.infer(text)
-    annotated_output = " ".join([f"{word}_{int(label)}   " for word, label in predictions])
+    annotated_output = " ".join([f"{word}_{int(label)}  " for word, label in predictions])
     return annotated_output
 
 interface = gr.Interface(fn = annotate,
@@ -186,6 +181,6 @@ interface = gr.Interface(fn = annotate,
                              placeholder="Tagged sentence appears here...",
                          ),
                          title = "Named Entity Recognition",
-                         description = "CS626 Assignment 2 (Autumn 2024)",
+                         description = "CS626 Assignment 3 (Autumn 2024)",
                          theme=gr.themes.Soft())
 interface.launch()