brucewlee / lftk Goto Github PK

import spacy
import lftk
import math
import nltk
from supar import Parser


# load models
nlp = spacy.load('en_core_web_sm')
SuPar = Parser.load('crf-con-en')

def preprocess(doc, short=False, see_token=False, see_sent_token=True):
    n_token = 1
    n_sent = 1
    token_list = []
    #raw_token_list = []
    sent_token_list = []

    # sent_list is for raw string sentences
    sent_list = []

    # count tokens, sentence + make lists
    #for sent in self.NLP_doc.sents:
    for sent in doc.sents:
        n_sent += 1
        sent_list.append(sent.text)
        temp_list = []
        for token in sent:
            if token.text.isalpha():
                temp_list.append(token.text)
                if short == True:
                    n_token += 1
                    token_list.append(token.lemma_.lower())
                if short == False:
                    if len(token.text) >= 3:
                        n_token += 1
                        token_list.append(token.lemma_.lower())
        if len(temp_list) > 3:
            sent_token_list.append(temp_list)

    #self.n_token = n_token 
    #self.n_sent = n_sent
    #self.token_list = token_list
    #self.sent_token_list = sent_token_list
    
    result = {"n_token": n_token, 
                "n_sent":n_sent
                }

    if see_token == True:
        result["token"] = token_list
    if see_sent_token == True:
        result["sent_token"] = sent_token_list

    return result

def calculate_nerf(extracted_features):
    return (0.04876 * extracted_features['t_kup'] - 0.1145 * extracted_features['t_subtlex_us_zipf']) / extracted_features['t_sent'] \
        + (0.3091 * (extracted_features['n_noun'] + extracted_features['n_verb'] + extracted_features['n_num'] + extracted_features['n_adj'] + extracted_features['n_adv']) + 0.1866 * extracted_features['n_noun'] + 0.2645 * extracted_features['to_TreeH_C']) / extracted_features['t_sent'] \
        + (1.1017 * extracted_features['t_uword']) / math.sqrt(extracted_features['t_word']) - 4.125


text = 'This is simple example sentence. This is another example sentence.'
doc = nlp(text)

# initiate LFTK extractor by passing in doc
LFTK = lftk.Extractor(docs = doc)
LFTK.customize(stop_words=True, punctuations=False, round_decimal=3)

preprocessed_features = preprocess(doc, short=False, see_token=False, see_sent_token=True)
TrSF = retrieve(SuPar, preprocessed_features['sent_token'])
feature_keys = ['t_kup', 't_subtlex_us_zipf', 't_sent', 'n_noun', 'n_verb', 'n_adj', 'n_adv', 'n_num', 't_uword', 't_word']

extracted_features = LFTK.extract(features = feature_keys)
extracted_features.update(TrSF)

# convert to float
extracted_features = {k: float(v) for k, v in extracted_features.items()}

print(calculate_nerf(extracted_features))