# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ### NER Model Training
# This notebook contains all the training runs and hyperparameter tuning used to train the NER model.

import pandas as pd
import numpy as np
from sklearn.feature_extraction import DictVectorizer
from sklearn.feature_extraction.text import HashingVectorizer
#from sklearn.linear_model import Perceptron
from sklearn.model_selection import train_test_split
#from sklearn.linear_model import SGDClassifier
#from sklearn.linear_model import PassiveAggressiveClassifier
#from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import classification_report
import pickle

PATH = '../GLG-Old Models, datasets/'
df = pd.read_csv(PATH + 'ner_dataset.csv', encoding = "ISO-8859-1")
#df = df[:500000]
df.head()

df.shape

df.fillna(method='ffill',inplace=True)

df.head()

# +
#X = df.drop('Tag',axis=1)
#v = DictVectorizer(sparse=False)
#X = v.fit_transform(X.to_dict('records'))
#y = df['Tag']

classes = list(np.unique(df['Tag']))

#X_train, X_text, y_train, y_test = train_test_split(X, y, test_size=0.15,random_state=4)
# -

import sklearn_crfsuite
from sklearn_crfsuite import scorers
from sklearn_crfsuite import metrics
from collections import Counter


class SentenceGetter(object):
    
    def __init__(self, data):
        self.n_sent = 1
        self.data = data
        self.empty = False
        agg_func = lambda s: [(w, p, t) for w, p, t in zip(s['Word'].values.tolist(), 
                                                           s['POS'].values.tolist(), 
                                                           s['Tag'].values.tolist())]
        self.grouped = self.data.groupby('Sentence #').apply(agg_func)
        self.sentences = [s for s in self.grouped]
        
    def get_next(self):
        try: 
            s = self.grouped['Sentence: {}'.format(self.n_sent)]
            self.n_sent += 1
            return s 
        except:
            return None
getter = SentenceGetter(df)
sentences = getter.sentences


def word2features(sent, i):
    word = sent[i][0]
    postag = sent[i][1]
    
    features = {
        'bias': 1.0, 
        'word.lower()': word.lower(), 
        'word[-3:]': word[-3:],
        'word[-2:]': word[-2:],
        'word.isupper()': word.isupper(),
        'word.istitle()': word.istitle(),
        'word.isdigit()': word.isdigit(),
        'postag': postag,
        'postag[:2]': postag[:2],
    }
    if i > 0:
        word1 = sent[i-1][0]
        postag1 = sent[i-1][1]
        features.update({
            '-1:word.lower()': word1.lower(),
            '-1:word.istitle()': word1.istitle(),
            '-1:word.isupper()': word1.isupper(),
            '-1:postag': postag1,
            '-1:postag[:2]': postag1[:2],
        })
    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.istitle()': word1.istitle(),
            '+1:word.isupper()': word1.isupper(),
            '+1:postag': postag1,
            '+1:postag[:2]': postag1[:2],
        })
    else:
        features['EOS'] = True
    return features
def sent2features(sent):
    return [word2features(sent, i) for i in range(len(sent))]
def sent2labels(sent):
    return [label for token, postag, label in sent]
def sent2tokens(sent):
    return [token for token, postag, label in sent]


X = [sent2features(s) for s in sentences]
y = [sent2labels(s) for s in sentences]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=4)

# +
## X_p = list(np.array([sent2features(s) for s in sentences]).flatten())
## y_p = list(np.array([sent2labels(s) for s in sentences]).flatten())
## X_train_p, X_test_p, y_train_p, y_test_p = train_test_split(X_p, y_p, test_size=0.1, random_state=4,stratify=y)
# -

# 08:54
crf = sklearn_crfsuite.CRF(
#    algorithm='lbfgs',
    algorithm='l2sgd',
#    c1=0.1,
#    c2=0.1,
    max_iterations=100,
    all_possible_transitions=True
)
crf.fit(X_train, y_train)

crf = sklearn_crfsuite.CRF(
    algorithm='lbfgs',
    c1=0.1,
    c2=0.1,
    max_iterations=100,
    all_possible_transitions=True
)
crf.fit(X_train, y_train)

pd.DataFrame(y_train)

# Test size = 0.33
y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))

# Test size = 0.33
y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred, labels=classes))

#Default, test_size = 0.5
y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))

filename = 'crf_ner_model.sav'
pickle.dump(crf, open(filename, 'wb'))

print(metrics.flat_classification_report(y_train, y_train))

#SGD, test_size = 0.5
y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))

#SGD
y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))

# SGD
y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred, labels=classes))

y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred, labels=classes))

y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))

y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))

y_pred = crf.predict(X_test)
print(metrics.flat_classification_report(y_test, y_pred))#, labels = classes))

# ### Predict on Some Sample Queries

# +
import re

def prep_query(phrase):
    split_query = re.findall(r"[\w']+|[.,!?;]", phrase)
    
    pos_tags = pos_tag(split_query)
    
    df_query = pd.DataFrame({'Sentence #':['Sentence: 1'] * len(pos_tags),
                            'Word':[pair[0] for pair in pos_tags],
                            'POS':[pair[1] for pair in pos_tags],
                            'Tag':[None] * len(pos_tags)})
       
    return df_query


# +
s = "<NAME> is a former host on The Apprentice. He is an American businessman and former President."
s = 'hello how are you'
s = 'The Second World War started in 1914 and ended in 1918'
s = 'The Korean War started in 1939 and ended in 1945'
s = 'Iraq and Iran were once at war. Saddam Hussein was involved'
s = 'The World Cup is a quadrennial sporting event. FIFA is the governing body involved.'


x = prep_query(s)

# +
getter_query = SentenceGetter(x)
sentences_query = getter_query.sentences

X_query = [sent2features(s) for s in sentences_query]
X_words = [s[0] for s in sentences_query[0]]

pred = crf.predict(X_query)

list(zip(pred[0],X_words))
# -



StarcoderdataJupyterScriptsDedupFiltered