# ---
# 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
# ---

# <font size=6>
#     <b>Model_Training_with_BERT.ipynb:</b>
#     <p>Use Text Extensions for Pandas to integrate BERT tokenization with model training for named entity recognition on Pandas.</p>
# </font>

# # Introduction
#
# This notebook shows how to use the open source library [Text Extensions for Pandas](https://github.com/CODAIT/text-extensions-for-pandas) to seamlessly integrate BERT tokenization and embeddings with model training for named entity recognition using [Pandas](https://pandas.pydata.org/) DataFrames.
#
# This example will build on the analysis of the [CoNLL-2003](https://www.clips.uantwerpen.be/conll2003/ner/) corpus done in [Analyze_Model_Outputs](./Analyze_Model_Outputs.ipynb) to train a new model for named entity recognition (NER) using state-of-the-art natural language understanding with BERT tokenization and embeddings. While the model used is rather simple and will only get modest scoring results, the purpose is to demonstrate how Text Extensions for Pandas integrates BERT from [Huggingface Transformers](https://huggingface.co/transformers/index.html) with the `TensorArray` extension for model training and scoring, all within Pandas DataFrames. See [Text_Extension_for_Pandas_Overview](./Text_Extension_for_Pandas_Overview.ipynb) for `TensorArray` specification and more example usage.
#
# The notebook is divided into the following steps:
#
# 1. Retokenize the entire corpus using a "BERT-compatible" tokenizer, and map the token/entity labels from the original corpus on to the new tokenization.
# 1. Generate BERT embeddings for every token in the entire corpus in one pass, and store those embeddings in a DataFrame column (of type TensorDtype) alongside the tokens and labels.
# 1. Persist the DataFrame with computed BERT embeddings to disk as a checkpoint.
# 1. Use the embeddings to train a multinomial logistic regression model to perform named entity recognition.
# 1. Compute precision/recall for the model predictions on a test set.
#

# ## Environment Setup
#
# This notebook requires a Python 3.7 or later environment with NumPy, Pandas, scikit-learn, PyTorch and Huggingface `transformers`. 
#
# The notebook also requires the  `text_extensions_for_pandas` library. You can satisfy this dependency in two ways:
#
# * Run `pip install text_extensions_for_pandas` before running this notebook. This command adds the library to your Python environment.
# * Run this notebook out of your local copy of the Text Extensions for Pandas project's [source tree](https://github.com/CODAIT/text-extensions-for-pandas). In this case, the notebook will use the version of Text Extensions for Pandas in your local source tree **if the package is not installed in your Python environment**.

# +
import gc
import os
import sys
from typing import *
import numpy as np
import pandas as pd
import sklearn.pipeline
import sklearn.linear_model
import transformers

# And of course we need the text_extensions_for_pandas library itself.
try:
    import text_extensions_for_pandas as tp
except ModuleNotFoundError as e:
    # If we're running from within the project source tree and the parent Python
    # environment doesn't have the text_extensions_for_pandas package, use the
    # version in the local source tree.
    if not os.getcwd().endswith("notebooks"):
        raise e
    if ".." not in sys.path:
        sys.path.insert(0, "..")
    import text_extensions_for_pandas as tp
# -

# # Named Entity Recognition with BERT on CoNLL-2003
#
# [CoNLL](https://www.conll.org/), the SIGNLL Conference on Computational Natural Language Learning, is an annual academic conference for natural language processing researchers. Each year's conference features a competition involving a challenging NLP task. The task for the 2003 competition involved identifying mentions of [named entities](https://en.wikipedia.org/wiki/Named-entity_recognition) in English and German news articles from the late 1990's. The corpus for this 2003 competition is one of the most widely-used benchmarks for the performance of named entity recognition models. Current [state-of-the-art results](https://paperswithcode.com/sota/named-entity-recognition-ner-on-conll-2003) on this corpus produce an F1 score (harmonic mean of precision and recall) of 0.93. The best F1 score in the original competition was 0.89.
#
# For more information about this data set, we recommend reading the conference paper about the competition results, ["Introduction to the CoNLL-2003 Shared Task: Language-Independent Named Entity Recognition,"](https://www.aclweb.org/anthology/W03-0419/).
#
# **Note that the data set is licensed for research use only. Be sure to adhere to the terms of the license when using this data set!**
#
# The developers of the CoNLL-2003 corpus defined a file format for the corpus, based on the file format used in the earlier [Message Understanding Conference](https://en.wikipedia.org/wiki/Message_Understanding_Conference) competition. This format is generally known as "CoNLL format" or "CoNLL-2003 format".
#
# In the following cell, we use the facilities of Text Extensions for Pandas to download a copy of the CoNLL-2003 data set. Then we read the CoNLL-2003-format file containing the `test` fold of the corpus and translate the data into a collection of Pandas [DataFrame](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html) objects, one Dataframe per document. Finally, we display the Dataframe for the first document of the `test` fold of the corpus.

# Download and cache the data set.
# NOTE: This data set is licensed for research use only. Be sure to adhere
#  to the terms of the license when using this data set!
data_set_info = tp.io.conll.maybe_download_conll_data("outputs")
data_set_info

# ## Show how to retokenize with a BERT tokenizer.
#
# The BERT model is originally from the paper [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805) by <NAME>, <NAME>, <NAME>, <NAME>. The model is pre-trained with masked language modeling and next sentence prediction objectives, which make it effective for masked token prediction and NLU. 
#
# With the CoNLL-2003 corpus loaded, it will need to be retokenized using a "BERT-compatible" tokenizer. Then we can map the token/entity labels from the original corpus on to the new tokenization.
#
# We will start by showing the retokenizing process for a single document before doing the same on the entire corpus.

# +
# Read in the corpus in its original tokenization.
corpus_raw = {}
for fold_name, file_name in data_set_info.items():
    df_list = tp.io.conll.conll_2003_to_dataframes(file_name, 
                                                   ["pos", "phrase", "ent"],
                                                   [False, True, True])
    corpus_raw[fold_name] = [
        df.drop(columns=["pos", "phrase_iob", "phrase_type"])
        for df in df_list
    ]

test_raw = corpus_raw["test"]

# Pick out the dataframe for a single example document.
example_df = test_raw[5]
example_df
# -

# The `example_df` contains columns `span` and `sentence` of dtypes `SpanDtype` and `TokenSpanDtype`. These represent spans from the target text, and here they contain tokens of the text and the sentence containing that token. See the notebook [Text_Extension_for_Pandas_Overview](./Text_Extension_for_Pandas_Overview.ipynb) for more on `SpanArray` and `TokenSpanArray`.

example_df.dtypes

# ### Convert IOB-Tagged Data to Lists of Entity Mentions
#
# The data we've looked at so far has been in [IOB2 format](https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging)). 
# Each row of our DataFrame represents a token, and each token is tagged with an entity type (`ent_type`) and an IOB tag (`ent_iob`). The first token of each named entity mention is tagged `B`, while subsequent tokens are tagged `I`. Tokens that aren't part of any named entity are tagged `O`.
#
# IOB2 format is a convenient way to represent a corpus, but it is a less useful representation for analyzing the result quality of named entity recognition models. Most tokens in a typical NER corpus will be tagged `O`, any measure of error rate in terms of tokens will over-emphasizing the tokens that are part of entities. Token-level error rate implicitly assigns higher weight to named entity mentions that consist of multiple tokens, further unbalancing error metrics. And most crucially, a naive comparison of IOB tags can result in marking an incorrect answer as correct. Consider a case where the correct sequence of labels is `B, B, I` but the model has output `B, I, I`; in this case, last two tokens of model output are both incorrect (the model has assigned them to the same entity as the first token), but a naive token-level comparison will consider the last token to be correct.
#
# The CoNLL 2003 competition used the number of errors in extracting *entire* entity mentions to measure the result quality of the entries. We will use the same metric in this notebook. To compute entity-level errors, we convert the IOB-tagged tokens into pairs of `<entity span,  entity type>`. 
# Text Extensions for Pandas includes a function `iob_to_spans()` that will handle this conversion for you.

# Convert the corpus IOB2 tagged DataFrame to one with entity span and type columns.
spans_df = tp.io.conll.iob_to_spans(example_df)
spans_df

# ### Initialize our BERT Tokenizer and Model
#
# Here we configure and initialize the [Huggingface transformers BERT tokenizer and model](https://huggingface.co/transformers/model_doc/bert.html). Text Extensions for Pandas provides a `make_bert_tokens()` function that will use the tokenizer to create BERT tokens as a span column in a DataFrame, suitable to compute BERT embeddings with.

# +
# Huggingface transformers BERT Configuration.
bert_model_name = "dslim/bert-base-NER"

tokenizer = transformers.BertTokenizerFast.from_pretrained(bert_model_name, 
                                                           add_special_tokens=True)

# Retokenize the document's text with the BERT tokenizer as a DataFrame 
# with a span column.
bert_toks_df = tp.io.bert.make_bert_tokens(example_df["span"].values[0].target_text, tokenizer)
bert_toks_df
# -

# BERT tokenization includes special zero-length tokens.
bert_toks_df[bert_toks_df["special_tokens_mask"]]

# Align the BERT tokens with the original tokenization.
bert_spans = tp.TokenSpanArray.align_to_tokens(bert_toks_df["span"],
                                               spans_df["span"])
pd.DataFrame({
    "original_span": spans_df["span"],
    "bert_spans": bert_spans,
    "ent_type": spans_df["ent_type"]
})

# Generate IOB2 tags and entity labels that align with the BERT tokens.
# See https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging)
bert_toks_df[["ent_iob", "ent_type"]] = tp.io.conll.spans_to_iob(bert_spans, 
                                                        spans_df["ent_type"])
bert_toks_df[10:20]

# Create a Pandas categorical type for consistent encoding of categories
# across all documents.
ENTITY_TYPES = ["LOC", "MISC", "ORG", "PER"]
token_class_dtype, int_to_label, label_to_int = tp.io.conll.make_iob_tag_categories(ENTITY_TYPES)
token_class_dtype

# The traditional way to transform NER to token classification is to 
# treat each combination of {I,O,B} X {entity type} as a different
# class. Generate class labels in that format.
classes_df = tp.io.conll.add_token_classes(bert_toks_df, token_class_dtype)
classes_df

# ## Show how to compute BERT embeddings
#
# We are going to use the BERT embeddings as the feature vector to train our model. First, we will show how they are computed 

# +
# Initialize the BERT model that will be used to generate embeddings.
bert = transformers.BertModel.from_pretrained(bert_model_name)

# Force garbage collection in case this notebook is running on a low-RAM environment.
gc.collect()

# Compute BERT embeddings with the BERT model and add result to our example DataFrame.
embeddings_df = tp.io.bert.add_embeddings(classes_df, bert)
embeddings_df[["token_id", "span", "input_id", "ent_iob", "ent_type", "token_class", "embedding"]].iloc[10:20]
# -

embeddings_df[["span", "ent_iob", "ent_type", "embedding"]].iloc[70:75]

# The `embedding` column is an extension type `TensorDtype` that holds a 
#`TensorArray` provided by Text Extensions for Pandas.
embeddings_df["embedding"].dtype

# A `TensorArray` can be constructed with a NumPy array of arbitrary dimensions, added to a DataFrame, then used with standard Pandas functionality. See the notebook [Text_Extension_for_Pandas_Overview](./Text_Extensions_for_Pandas.ipynb) for more on `TensorArray`.

# Zero-copy conversion to NumPy can be done by first unwrapping the
# `TensorArray` with `.array` and calling `to_numpy()`.
embeddings_arr = embeddings_df["embedding"].array.to_numpy()
embeddings_arr.dtype, embeddings_arr.shape

# ## Generate BERT tokens and BERT embeddings for the entire corpus
#
# Text Extensions for Pandas has a convenience function that will combine the above cells to create BERT tokens and embeddings. We will use this to add embeddings to the entire corpus.

# Example usage of the convenience function to create BERT tokens and embeddings.
tp.io.bert.conll_to_bert(example_df, tokenizer, bert, token_class_dtype)

# When this notebook is running on a resource-constrained environment like [Binder](https://mybinder.org/),
# there may not be enough RAM available to hold all the embeddings in memory.
# So we use [Gaussian random projection](https://scikit-learn.org/stable/modules/random_projection.html#gaussian-random-projection) to reduce the size of the embeddings.
# The projection shrinks the embeddings by a factor of 3 at the expense of a small
# decrease in model accuracy.
#
# Change the constant `SHRINK_EMBEDDINGS` in the following cell to `False` if you want to disable this behavior.

# +
SHRINK_EMBEDDINGS = True
PROJECTION_DIMS = 256
RANDOM_SEED=42

import sklearn.random_projection
projection = sklearn.random_projection.GaussianRandomProjection(
    n_components=PROJECTION_DIMS, random_state=RANDOM_SEED)

def maybe_shrink_embeddings(df):
    if SHRINK_EMBEDDINGS:
        df["embedding"] = tp.TensorArray(projection.fit_transform(df["embedding"]))
    return df


# +
# Run the entire corpus through our processing pipeline.
bert_toks_by_fold = {}
for fold_name in corpus_raw.keys():
    print(f"Processing fold '{fold_name}'...")
    raw = corpus_raw[fold_name]
    bert_toks_by_fold[fold_name] = tp.jupyter.run_with_progress_bar(
        len(raw), lambda i: maybe_shrink_embeddings(tp.io.bert.conll_to_bert(
            raw[i], tokenizer, bert, token_class_dtype)))
    
bert_toks_by_fold["dev"][20]
# -

# ## Collate the data structures we've generated so far

# Create a single DataFrame with the entire corpus's embeddings.
corpus_df = tp.io.conll.combine_folds(bert_toks_by_fold)
corpus_df

# ## Checkpoint
#
# With the `TensorArray` from Text Extensions for Pandas, the computed embeddings can be persisted as a tensor along with the rest of the DataFrame using standard Pandas input/output methods. Since this is a costly operation and the embeddings are deterministic, it can save lots of time to checkpoint the data here and save the results to disk. This will allow us to continue working with model training without needing to re-compute the BERT embeddings again.
#  
# ### Save DataFrame with Embeddings Tensor

# Write the tokenized corpus with embeddings to a Feather file.
# We can't currently serialize span columns that cover multiple documents (see issue #73 https://github.com/CODAIT/text-extensions-for-pandas/issues/73),
# so drop span columns from the contents we write to the Feather file.
cols_to_drop = [c for c in corpus_df.columns if "span" in c]
corpus_df.drop(columns=cols_to_drop).to_feather("outputs/corpus.feather")

# ### Load DataFrame with Previously Computed Embeddings

# Read the serialized embeddings back in so that you can rerun the model 
# training parts of this notebook (the cells from here onward) without 
# regenerating the embeddings.
corpus_df = pd.read_feather("outputs/corpus.feather")
corpus_df

# ## Training a model on the BERT embeddings
#
# Now we will use the loaded BERT embeddings to train a multinomial model to predict the token class from the embeddings tensor.

# Extract the training set DataFrame.
train_df = corpus_df[corpus_df["fold"] == "train"]
train_df

# +
# Train a multinomial logistic regression model on the training set.
MULTI_CLASS = "multinomial"
    
# How many iterations to run the BGFS optimizer when fitting logistic
# regression models. 100 ==> Fast; 10000 ==> Full convergence
LBGFS_ITERATIONS = 10000

base_pipeline = sklearn.pipeline.Pipeline([
    # Standard scaler. This only makes a difference for certain classes
    # of embeddings.
    #("scaler", sklearn.preprocessing.StandardScaler()),
    ("mlogreg", sklearn.linear_model.LogisticRegression(
        multi_class=MULTI_CLASS,
        verbose=10,
        max_iter=LBGFS_ITERATIONS
    ))
])

X_train = train_df["embedding"].values
Y_train = train_df["token_class_id"]
base_model = base_pipeline.fit(X_train, Y_train)
base_model


# -

# ## Make Predictions on Token Class from BERT Embeddings
#
# Using our model, we can now predict the token class from the test set using the computed embeddings.

# Define a function that will let us make predictions on a fold of the corpus.
def predict_on_df(df: pd.DataFrame, id_to_class: Dict[int, str], predictor):
    """
    Run a trained model on a DataFrame of tokens with embeddings.

    :param df: DataFrame of tokens for a document, containing a TokenSpan column
     called "embedding" for each token.
    :param id_to_class: Mapping from class ID to class name, as returned by
     :func:`text_extensions_for_pandas.make_iob_tag_categories`
    :param predictor: Python object with a `predict_proba` method that accepts
     a numpy array of embeddings.
    :returns: A copy of `df`, with the following additional columns:
     `predicted_id`, `predicted_class`, `predicted_iob`, `predicted_type`
     and `predicted_class_pr`.
    """
    result_df = df.copy()
    class_pr = tp.TensorArray(predictor.predict_proba(result_df["embedding"]))
    result_df["predicted_id"] = np.argmax(class_pr, axis=1)
    result_df["predicted_class"] = [id_to_class[i]
                                    for i in result_df["predicted_id"].values]
    iobs, types = tp.io.conll.decode_class_labels(result_df["predicted_class"].values)
    result_df["predicted_iob"] = iobs
    result_df["predicted_type"] = types
    result_df["predicted_class_pr"] = class_pr
    return result_df


# Make predictions on the test set.
test_results_df = predict_on_df(corpus_df[corpus_df["fold"] == "test"], int_to_label, base_model)
test_results_df.head()

# Take a slice to show a region with more entities.
test_results_df.iloc[40:50]

# ## Compute Precision and Recall
#
# With our model predictions on the test set, we can now compute precision and recall. To do this, we will use the following steps:
#
# 1. Split up test set predictions by document, so we can work on the document level.
# 1. Join the test predictions with token information into one DataFrame per document.
# 1. Convert each DataFrame from IOB2 format to span, entity type pairs as done before.
# 1. Compute accuracy for each document as a DataFrame.
# 1. Aggregate per-document accuracy to get overal precision/recall.

# +
# Split model outputs for an entire fold back into documents and add
# token information.

# Get unique documents per fold.
fold_and_doc = test_results_df[["fold", "doc_num"]] \
        .drop_duplicates() \
        .to_records(index=False)

# Index by fold, doc and token id, then make sure sorted.
indexed_df = test_results_df \
        .set_index(["fold", "doc_num", "token_id"], verify_integrity=True) \
        .sort_index()

# Join predictions with token information, for each document.
test_results_by_doc = {}
for collection, doc_num in fold_and_doc:
    doc_slice = indexed_df.loc[collection, doc_num].reset_index()
    doc_toks = bert_toks_by_fold[collection][doc_num][
        ["token_id", "span", "ent_iob", "ent_type"]
    ].rename(columns={"id": "token_id"})
    joined_df = doc_toks.copy().merge(
        doc_slice[["token_id", "predicted_iob", "predicted_type"]])
    test_results_by_doc[(collection, doc_num)] = joined_df
    
# Test results are now in one DataFrame per document.
test_results_by_doc[("test", 0)].iloc[40:60]

# +
# Convert IOB2 format to spans, entity type with `tp.io.conll.iob_to_spans()`.
test_actual_spans = {k: tp.io.conll.iob_to_spans(v) for k, v in test_results_by_doc.items()}
test_model_spans = {k:
        tp.io.conll.iob_to_spans(v, iob_col_name = "predicted_iob",
                                 entity_type_col_name = "predicted_type")
            .rename(columns={"predicted_type": "ent_type"})
        for k, v in test_results_by_doc.items()}

test_model_spans[("test", 0)].head()
# -

# Compute per-document statistics into a single DataFrame.
test_stats_by_doc = tp.io.conll.compute_accuracy_by_document(test_actual_spans, test_model_spans)
test_stats_by_doc

# Collection-wide precision and recall can be computed by aggregating
# our DataFrame.
tp.io.conll.compute_global_accuracy(test_stats_by_doc)

# ### Adjusting the BERT Model Output
#
# The above results aren't bad for a first shot, but taking a look a some of the predictions will show that sometimes the tokens have been split up into multiple entities. This is because the BERT tokenizer uses WordPiece to make subword tokens, see https://huggingface.co/transformers/tokenizer_summary.html and https://static.googleusercontent.com/media/research.google.com/ja//pubs/archive/37842.pdf for more information.
#
# This is going to cause a problem when computing precision/recall because we are comparing exact spans, and if the entity is split, it will be counted as a false negative _and_ possibly one or more false positives. Luckily we can fix up with Text Extension for Pandas.
#
# Let's drill down to see an example of the issue and how to correct it.

# Every once in a while, the BERT model will split a token in the original data
# set into multiple entities. For example, look at document 202 of the test set:
test_model_spans[("test", 202)].head(10)

# Notice `[150, 151): 'W'` and `[151, 156): 'idnes'`. These outputs are part
# of the same original token, but have been split by the model.

# +
# We can use spanner algebra in `tp.spanner.overlap_join()`
# to fix up these outputs.
spans_df = test_model_spans[("test", 202)]
toks_df = test_raw[202]

# First, find which tokens the spans overlap with:
overlaps_df = (
    tp.spanner.overlap_join(spans_df["span"], toks_df["span"],
                            "span", "corpus_token")
        .merge(spans_df)
)
overlaps_df.head(10)
# -

# Next, compute the minimum span that covers all the corpus tokens
# that overlap with each entity span.
agg_df = (
    overlaps_df
    .groupby("span")
    .aggregate({"corpus_token": "sum", "ent_type": "first"})
    .reset_index()
)
agg_df.head(10)

# Finally, take unique values and covert character-based spans to token
# spans in the corpus tokenization (since the new offsets might not match a
# BERT tokenizer token boundary).
cons_df = (
    tp.spanner.consolidate(agg_df, "corpus_token")[["corpus_token", "ent_type"]]
        .rename(columns={"corpus_token": "span"})
)
cons_df["span"] = tp.TokenSpanArray.align_to_tokens(toks_df["span"],
                                                    cons_df["span"])
cons_df.head(10)

# Text Extensions for Pandas contains a single function that repeats the actions of the 
# previous 3 cells.
tp.io.bert.align_bert_tokens_to_corpus_tokens(test_model_spans[("test", 202)], test_raw[202]).head(10)

# Run all of our DataFrames through `align_bert_tokens_to_corpus_tokens()`.
keys = list(test_model_spans.keys())
new_values = tp.jupyter.run_with_progress_bar(
    len(keys), 
    lambda i: tp.io.bert.align_bert_tokens_to_corpus_tokens(test_model_spans[keys[i]], test_raw[keys[i][1]]))
test_model_spans = {k: v for k, v in zip(keys, new_values)}
test_model_spans[("test", 202)].head(10)

# Compute per-document statistics into a single DataFrame.
test_stats_by_doc = tp.io.conll.compute_accuracy_by_document(test_actual_spans, test_model_spans)
test_stats_by_doc

# Collection-wide precision and recall can be computed by aggregating
# our DataFrame.
tp.io.conll.compute_global_accuracy(test_stats_by_doc)

# These results are a bit better than before, and while the F1 score is not high compared to todays standards, it is decent enough for a simplistic model. More importantly, we did show it was fairly easy to create a model for named entity recognition and analyze the output by leveraging the functionalitiy of Pandas DataFrames along with [Text Extensions for Pandas](https://github.com/CODAIT/text-extensions-for-pandas) `SpanArray`, `TensorArray` and integration with BERT from [Huggingface Transformers](https://huggingface.co/transformers/index.html).

StarcoderdataJupyterScriptsDedupFiltered