!pip install git+https://github.com/boudinfl/pke.git
!pip install datasets
!python -m spacy download en_core_web_sm

"""
Explanation: How to compute DF (document frequency) counts using pke
DF counts are required by several keyphrase extraction models (e.g. TfIdf, Kea) for weighting candidates. Below is an example on how to compute DF counts for a given text collection.
End of explanation
"""


from datasets import load_dataset

# load the inspec dataset
dataset = load_dataset('taln-ls2n/inspec')

import spacy
from tqdm.notebook import tqdm
from spacy.tokenizer import _get_regex_pattern

nlp = spacy.load("en_core_web_sm")

# Tokenization fix for in-word hyphens (e.g. 'non-linear' would be kept 
# as one token instead of default spacy behavior of 'non', '-', 'linear')
# https://spacy.io/usage/linguistic-features#native-tokenizer-additions

from spacy.lang.char_classes import ALPHA, ALPHA_LOWER, ALPHA_UPPER
from spacy.lang.char_classes import CONCAT_QUOTES, LIST_ELLIPSES, LIST_ICONS
from spacy.util import compile_infix_regex

# Modify tokenizer infix patterns
infixes = (
    LIST_ELLIPSES
    + LIST_ICONS
    + [
        r"(?<=[0-9])[+\-\*^](?=[0-9-])",
        r"(?<=[{al}{q}])\.(?=[{au}{q}])".format(
            al=ALPHA_LOWER, au=ALPHA_UPPER, q=CONCAT_QUOTES
        ),
        r"(?<=[{a}]),(?=[{a}])".format(a=ALPHA),
        # ✅ Commented out regex that splits on hyphens between letters:
        # r"(?<=[{a}])(?:{h})(?=[{a}])".format(a=ALPHA, h=HYPHENS),
        r"(?<=[{a}0-9])[:<>=/](?=[{a}])".format(a=ALPHA),
    ]
)

infix_re = compile_infix_regex(infixes)
nlp.tokenizer.infix_finditer = infix_re.finditer

# populates a docs list with spacy doc objects
docs = []
for sample in tqdm(dataset['train']):
    docs.append(nlp(sample["title"]+". "+sample["abstract"]))

import logging
from pke import compute_document_frequency
from string import punctuation

logging.basicConfig(level=logging.INFO)

# compute idf weights
compute_document_frequency(
    documents=docs,
    output_file='inspec.df.gz',
    language='en',              # language of the input files
    normalization='stemming',   # use porter stemmer
    stoplist=list(punctuation), # stoplist (punctuation marks)
    n=5                         # compute n-grams up to 5-grams
)

"""
Explanation: Preamble on keyphrase extraction datasets using 🤗 datasets
For simplicity and ease of use, we rely on the datasets module from 🤗 huggingface to load and access sample documents from keyphrase extraction datasets. Please have a look at our organization page (https://huggingface.co/taln-ls2n) for more information on which datasets are available.
End of explanation
"""


from pke import load_document_frequency_file

df = load_document_frequency_file(input_file='inspec.df.gz')

# populates a docs list with spacy doc objects
docs = []
for sample in tqdm(dataset['test']):
    docs.append(nlp(sample["title"]+". "+sample["abstract"]))

from pke.unsupervised import TfIdf
extractor = TfIdf()
outputs = []
for i, doc in enumerate(tqdm(docs)):
    extractor.load_document(input=doc, language='en')
    extractor.grammar_selection(grammar="NP: {<ADJ>*<NOUN|PROPN>+}")
    extractor.candidate_weighting(df=df)
    outputs.append([u for u,v in extractor.get_n_best(n=5, stemming=True)])

from nltk.stem.snowball import SnowballStemmer as Stemmer

# populates the references list with stemmed keyphrases
references = []
for sample in tqdm(dataset['test']):
    sample_keyphrases = []
    for keyphrase in sample["keyphrases"]:
        # tokenize keyphrase
        tokens = [token.text for token in nlp(keyphrase)]
        # normalize tokens using Porter's stemming
        stems = [Stemmer('porter').stem(tok.lower()) for tok in tokens]
        sample_keyphrases.append(" ".join(stems))
    references.append(sample_keyphrases)

import numpy as np

def evaluate(top_N_keyphrases, references):
    P = len(set(top_N_keyphrases) & set(references)) / len(top_N_keyphrases)
    R = len(set(top_N_keyphrases) & set(references)) / len(references)
    F = (2*P*R)/(P+R) if (P+R) > 0 else 0 
    return (P, R, F)


# compute the P, R, F scores for the model
scores = []
for i, output in enumerate(tqdm(outputs)):
    scores.append(evaluate(output, references[i]))
    
# compute the average scores
avg_scores = np.mean(scores, axis=0)
    
# print out the performance of the model
print("Model: TfIdf P@5: {:.3f} R@5: {:.3f} F@5: {:.3f}".format(avg_scores[0], avg_scores[1], avg_scores[2]))

"""
Explanation: Loading DF counts and benchmarking TfIdf
End of explanation
"""


