<jupyter_start><jupyter_text>CoNLL_4.ipynbThis notebook contains the fourth part of the model training and analysis code from our CoNLL-2020 paper, ["Identifying Incorrect Labels in the CoNLL-2003 Corpus"](https://www.aclweb.org/anthology/2020.conll-1.16/).If you're new to the Text Extensions for Pandas library, we recommend that you startby reading through the notebook [`Analyze_Model_Outputs.ipynb`](https://github.com/CODAIT/text-extensions-for-pandas/blob/master/notebooks/Analyze_Model_Outputs.ipynb), which explains the portions of the library that we use in the notebooks in this directory. SummaryThis notebook repeats the model training process from `CoNLL_3.ipynb`, but performs a 10-fold cross-validation. This process involves training a total of 170 models -- 10 groups of 17. Next, this notebook evaluates each group of models over the holdout set from the associated fold of the cross-validation. Then it aggregates together these outputs and uses the same techniques used in `CoNLL_2.ipynb` to flag potentially-incorrect labels. Finally, the notebook writes out CSV files containing ranked lists of potentially-incorrect labels.  Libraries and constants<jupyter_code># Libraries
import numpy as np
import pandas as pd
import os
import sys
import time
import torch
import transformers
from typing import *
import sklearn.model_selection
import sklearn.pipeline
import matplotlib.pyplot as plt

# And of course we need the text_extensions_for_pandas library itself.
PROJECT_ROOT = "../.."
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("corpus"):
        raise e
    if PROJECT_ROOT not in sys.path:
        sys.path.insert(0, PROJECT_ROOT)
    import text_extensions_for_pandas as tp

# Code shared among notebooks is kept in util.py, in this directory.
import util

# BERT Configuration
# Keep this in sync with `CoNLL_3.ipynb`.
#bert_model_name = "bert-base-uncased"
#bert_model_name = "bert-large-uncased"
bert_model_name = "dslim/bert-base-NER"
tokenizer = transformers.BertTokenizerFast.from_pretrained(bert_model_name, 
                                                           add_special_tokens=True)
bert = transformers.BertModel.from_pretrained(bert_model_name)

# If False, use cached values, provided those values are present on disk
_REGENERATE_EMBEDDINGS = False
_REGENERATE_MODELS = False

# Number of dimensions that we reduce the BERT embeddings down to when
# training reduced-quality models.
#_REDUCED_DIMS = [8, 16, 32, 64, 128, 256]
_REDUCED_DIMS = [32, 64, 128, 256]

# How many models we train at each level of dimensionality reduction
_MODELS_AT_DIM = [4] * len(_REDUCED_DIMS)

# Consistent set of random seeds to use when generating dimension-reduced
# models. Index is [index into _REDUCED_DIMS, model number], and there are
# lots of extra entries so we don't need to resize this matrix.
from numpy.random import default_rng
_MASTER_SEED = 42
rng = default_rng(_MASTER_SEED)
_MODEL_RANDOM_SEEDS = rng.integers(0, 1e6, size=(8, 8))

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

# Parameters for splitting the corpus into folds
_KFOLD_RANDOM_SEED = _MASTER_SEED
_KFOLD_NUM_FOLDS = 10<jupyter_output><empty_output><jupyter_text>Read inputsRead in the corpus, retokenize it with the BERT tokenizer, add BERT embeddings, and convertto a single dataframe.<jupyter_code># 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
# The raw dataset 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
    ]
# Retokenize with the BERT tokenizer
bert_data = {
    key: tp.jupyter.run_with_progress_bar(
        len(val), 
        lambda i: tp.io.bert.conll_to_bert(val[i], tokenizer, bert, token_class_dtype,
                                           compute_embeddings=_REGENERATE_EMBEDDINGS))
    for key, val in corpus_raw.items()
}
# Create a single dataframe of annotated tokens for the entire corpus
if _REGENERATE_EMBEDDINGS:
    corpus_df = tp.combine_folds(bert_data)
    # We can't currently serialize span columns that cover multiple documents (see issue 73),
    # so the Feather file won't contain them. Drop these columns for consistency when
    # we regenerate the embeddings here.
    cols_to_drop = [c for c in corpus_df.columns if "span" in c]
    corpus_df.drop(columns=cols_to_drop, inplace=True)
else:
    # Use embeddings computed in CoNLL_3.ipynb
    _EMBEDDINGS_FILE = "outputs/corpus.feather"
    if not os.path.exists(_EMBEDDINGS_FILE):
        raise ValueError(f"Precomputed embeddings not found at {_EMBEDDINGS_FILE}. "
                         f"Please rerun CoNLL_3.ipynb to regenerate this file, or "
                         f"set _REGENERATE_EMBEDDINGS to True in the previous cell.")
    corpus_df = pd.read_feather("outputs/corpus.feather")

corpus_df.head()<jupyter_output><empty_output><jupyter_text>Prepare folds for a 10-fold cross-validationWe divide the documents of the corpus into 10 random samples.<jupyter_code># IDs for each of the keys
doc_keys = corpus_df[["fold", "doc_num"]].drop_duplicates().reset_index(drop=True)
doc_keys
# We want to split the documents randomly into _NUM_FOLDS sets, then
# for each stage of cross-validation train a model on the union of
# (_NUM_FOLDS - 1) of them while testing on the remaining fold.
# sklearn.model_selection doesn't implement this approach directly,
# but we can piece it together with some help from Numpy.
#from numpy.random import default_rng
rng = np.random.default_rng(seed=_KFOLD_RANDOM_SEED)
iloc_order = rng.permutation(len(doc_keys.index))
kf = sklearn.model_selection.KFold(n_splits=_KFOLD_NUM_FOLDS)

train_keys = []
test_keys = []
for train_ix, test_ix in kf.split(iloc_order):
    # sklearn.model_selection.KFold gives us a partitioning of the
    # numbers from 0 to len(iloc_order). Use that partitioning to 
    # choose elements from iloc_order, then use those elements to 
    # index into doc_keys.
    train_iloc = iloc_order[train_ix]
    test_iloc = iloc_order[test_ix]
    train_keys.append(doc_keys.iloc[train_iloc])
    test_keys.append(doc_keys.iloc[test_iloc])

train_keys[1].head(10)<jupyter_output><empty_output><jupyter_text>Dry run: Train and evaluate models on the first foldTrain models on the first of our 10 folds and manually examine some of the model outputs.<jupyter_code># Gather the training set together by joining our list of documents
# with the entire corpus on the composite key <fold, doc_num>
train_inputs_df = corpus_df.merge(train_keys[0])
train_inputs_df
# Repeat the same process for the test set
test_inputs_df = corpus_df.merge(test_keys[0])
test_inputs_df<jupyter_output><empty_output><jupyter_text>Train an ensemble of models<jupyter_code>import importlib
util = importlib.reload(util)
import sklearn.linear_model

def train_models(train_df: pd.DataFrame) \
        -> Dict[str, sklearn.base.BaseEstimator]:
    """
    Train an ensemble of models with different levels of noise.
    
    :param train_df: DataFrame of labeled training documents, with one
     row per token. Must contain the columns "embedding" (precomputed 
     BERT embeddings) and "token_class_id" (integer ID of token type)
     
    :returns: A mapping from mnemonic model name to trained model
    """
    X = train_df["embedding"].values
    Y = train_df["token_class_id"]
    models = {}  # Type: Dict[str, sklearn.base.BaseEstimator]

    
    for i in range(len(_REDUCED_DIMS)):
        num_dims = _REDUCED_DIMS[i]
        num_models = _MODELS_AT_DIM[i]
        for j in range(num_models):
            model_name = f"{num_dims}_{j + 1}"
            seed = _MODEL_RANDOM_SEEDS[i, j]
            print(f"Training model '{model_name}' (#{j + 1} "
                  f"at {num_dims} dimensions) with seed {seed}")
            models[model_name] = util.train_reduced_model(X, Y, num_dims, seed)
    
    print(f"Training model using all of "
          f"{X._tensor.shape[1]}-dimension embeddings.")
    models[f"{X._tensor.shape[1]}_1"] = (
        sklearn.linear_model.LogisticRegression(
            multi_class="multinomial",
            verbose=10,
            max_iter=10000
        )
        .fit(X, Y)
    )
    return models

def maybe_train_models(train_df: pd.DataFrame, fold_num: int):
    import pickle
    _CACHED_MODELS_FILE = f"outputs/fold_{fold_num}_models.pickle"
    if _REGENERATE_MODELS or not os.path.exists(_CACHED_MODELS_FILE):
        m = train_models(train_df)
        print(f"Trained {len(m)} models.")
        with open(_CACHED_MODELS_FILE, "wb") as f:
            pickle.dump(m, f)
    else:
        # Use a cached model when using cached embeddings
        with open(_CACHED_MODELS_FILE, "rb") as f:
            m = pickle.load(f)
            print(f"Loaded {len(m)} models from {_CACHED_MODELS_FILE}.")
    return m

models = maybe_train_models(train_inputs_df, 0)
print(f"Model names after loading or training: {', '.join(models.keys())}")
# Uncomment this code if you need to have the cells that follow ignore
# some of the models saved to disk.
# _MODEL_SIZES_TO_KEEP = [32, 64, 128, 256]
# _RUNS_TO_KEEP = [4] * len(_MODEL_SIZES_TO_KEEP)
# _OTHER_MODELS_TO_KEEP = ["768_1"]

# to_keep = _OTHER_MODELS_TO_KEEP.copy()
# for size in _MODEL_SIZES_TO_KEEP:
#     for num_runs in _RUNS_TO_KEEP:
#         for i in range(num_runs):
#             to_keep.append(f"{size}_{i+1}")

# models = {k: v for k, v in models.items() if k in to_keep}

# print(f"Model names after filtering: {', '.join(models.keys())}")<jupyter_output><empty_output><jupyter_text>Evaluate the models on this fold's test set<jupyter_code>def eval_models(models: Dict[str, sklearn.base.BaseEstimator],
                test_df: pd.DataFrame):
    """
    Bulk-evaluate an ensemble of models generated by :func:`train_models`.
    
    :param models: Output of :func:`train_models`
    :param test_df:  DataFrame of labeled test documents, with one
     row per token. Must contain the columns "embedding" (precomputed 
     BERT embeddings) and "token_class_id" (integer ID of token type)
    
    :returns: A dictionary from model name to results of 
     :func:`util.analyze_model`
    """
    todo = [(name, model) for name, model in models.items()]
    results = tp.jupyter.run_with_progress_bar(
        len(todo),
        lambda i: util.analyze_model(test_df, int_to_label, todo[i][1], 
                                     bert_data, corpus_raw, expand_matches=True),
        "model"
    )
    return {t[0]: result for t, result in zip(todo, results)}

evals = eval_models(models, test_inputs_df)
# Summarize how each of the models does on the test set.
def make_summary_df(evals_df: pd.DataFrame) -> pd.DataFrame:
    global_scores = [r["global_scores"] for r in evals_df.values()]
    return pd.DataFrame({
        "name": list(evals_df.keys()),
        "dims": [n.split("_")[0] for n in evals_df.keys()],
        "num_true_positives": [r["num_true_positives"] for r in global_scores],
        "num_entities": [r["num_entities"] for r in global_scores],
        "num_extracted": [r["num_extracted"] for r in global_scores],
        "precision": [r["precision"] for r in global_scores],
        "recall": [r["recall"] for r in global_scores],
        "F1": [r["F1"] for r in global_scores]
    })

summary_df = make_summary_df(evals)
# Plot the tradeoff between dimensionality and F1 score
x = summary_df["dims"]
y = summary_df["F1"]

plt.figure(figsize=(4,4))
plt.scatter(x, y)
#plt.yscale("log")
#plt.xscale("log")
plt.ylim([0.75, 1.0])
plt.xlabel("Number of Dimensions")
plt.ylabel("F1 Score")

# Also dump the raw data to a local file.
pd.DataFrame({"num_dims": x, "f1_score": y}).to_csv("outputs/dims_vs_f1_score_xval.csv",
                                                    index=False)

plt.show()<jupyter_output><empty_output><jupyter_text>Aggregate the model results and compare with the gold standard<jupyter_code>full_results = util.merge_model_results(evals)
full_results
# Drop Boolean columns for now
results = full_results[["fold", "doc_offset", "span", "ent_type", "gold", "num_models"]]
results
(results[results["gold"] == True][["num_models", "span"]]
 .groupby("num_models").count()
 .rename(columns={"span": "count"}))
(results[results["gold"] == False][["num_models", "span"]]
 .groupby("num_models").count()
 .rename(columns={"span": "count"}))
# Pull out some hard-to-find examples, sorting by document to make labeling easier
hard_to_get = results[results["gold"]].sort_values(["num_models", "fold", "doc_offset"]).head(20)
hard_to_get<jupyter_output><empty_output><jupyter_text>TODO: Relabel the above 20 examples with a Markdown table (copy from CSV)<jupyter_code># Hardest results not in the gold standard for models to avoid
hard_to_avoid = results[~results["gold"]].sort_values(["num_models", "fold", "doc_offset"], ascending=[False, True, True]).head(20)
hard_to_avoid<jupyter_output><empty_output><jupyter_text>TODO: Relabel the above 20 examples (copy from CSV)  Remainder of ExperimentFor each of the 10 folds, train a model on the fold's training set and runanalysis on the fold's test set.<jupyter_code>def handle_fold(fold_ix: int) -> Dict[str, Any]:
    """
    The per-fold processing of the previous section's cells, collapsed into 
    a single function.
    
    :param fold_ix: 0-based index of fold
    
    :returns: a dictionary that maps data structure name to data structure
    """
    # To avoid accidentally picking up leftover data from a previous cell,
    # local to this function are named with a leading underscore
    _train_inputs_df = corpus_df.merge(train_keys[fold_ix])
    _test_inputs_df = corpus_df.merge(test_keys[fold_ix])
    _models = maybe_train_models(_train_inputs_df, fold_ix)
    _evals = eval_models(_models, _test_inputs_df)
    _summary_df = make_summary_df(_evals)
    _full_results = util.merge_model_results(_evals)
    _results = _full_results[["fold", "doc_offset", "span", 
                              "ent_type", "gold", "num_models"]]
    return {
        "models": _models,
        "summary_df": _summary_df,
        "full_results": _full_results,
        "results": _results
    }

# Start with the (already computed) results for fold 0
results_by_fold = [
    {
        "models": models,
        "summary_df": summary_df,
        "full_results": full_results,
        "results": results
    }
]

for fold in range(1, _KFOLD_NUM_FOLDS):
    print(f"Starting fold {fold}.")
    results_by_fold.append(handle_fold(fold))
    print(f"Done with fold {fold}.")
    
# Combine all the results into a single dataframe for the entire corpus
all_results = pd.concat([r["results"] for r in results_by_fold])
all_results<jupyter_output><empty_output><jupyter_text>Generate CSV files for manual labeling<jupyter_code># Reformat for output
dev_and_test_results = all_results[all_results["fold"].isin(["dev", "test"])]
in_gold_to_write, not_in_gold_to_write = util.csv_prep(dev_and_test_results, "num_models")
in_gold_to_write
not_in_gold_to_write
in_gold_to_write.to_csv("outputs/CoNLL_4_in_gold.csv", index=False)
not_in_gold_to_write.to_csv("outputs/CoNLL_4_not_in_gold.csv", index=False)
# Repeat for the contents of the original training set
train_results = all_results[all_results["fold"] == "train"]
in_gold_to_write, not_in_gold_to_write = util.csv_prep(train_results, "num_models")
in_gold_to_write
not_in_gold_to_write
in_gold_to_write.to_csv("outputs/CoNLL_4_train_in_gold.csv", index=False)
not_in_gold_to_write.to_csv("outputs/CoNLL_4_train_not_in_gold.csv", index=False)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09757579439820159, 'fraction_numerical': 0.008870526763472872, 'mean_word_length': 4.044437634079069, 'pattern_counts': {'":': 20, '<': 42, '<?xml version=': 0, '>': 45, 'https://': 2, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '215348', 'license': 'Apache-2.0', 'n_tokens_mistral': 5566, 'n_tokens_neox': 5379, 'n_words': 1518, 'path': 'tutorials/corpus/CoNLL_4.ipynb', 'repo_name': 'BryanCutler/text-extensions-for-pandas', 'size': '16459'}