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

# # Train a ML Classifier to Link FEBRL People Data

# <a href="https://colab.research.google.com/github/rachhouse/intro-to-data-linking/blob/main/tutorial_notebooks/02_Link_FEBRL_Data_with_ML_Classifier.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>

# In this tutorial, we'll train a machine learning classifier to score candidate pairs for linking, using supervised learning. We will use the same training dataset as the SimSum classification tutorial, as well as the same augmentation, blocking, and comparing functions. The functions have been included in a separate `.py` file for re-use and convenience, so we can focus on code unique to this tutorial.
#
# The SimSum classification tutorial included a more detailed walkthrough of augmentation, blocking, and comparing, and since we're using the same functions within this tutorial, details will be light for those steps. Please see the SimSum tutorial if you need a refresher.

# ## Google Colab Setup

# +
# Check if we're running locally, or in Google Colab.
try:
    import google.colab
    COLAB = True
except ModuleNotFoundError:
    COLAB = False
    
# If we're running in Colab, download the tutorial functions file 
# to the Colab session local directory, and install required libraries.
if COLAB:
    import requests
    
    tutorial_functions_url = "https://raw.githubusercontent.com/rachhouse/intro-to-data-linking/main/tutorial_notebooks/linking_tutorial_functions.py"
    r = requests.get(tutorial_functions_url)
    
    with open("linking_tutorial_functions.py", "w") as fh:
        fh.write(r.text)
    
    # !pip install -q recordlinkage jellyfish altair
# -

# ## Imports

# +
import itertools

import altair as alt
import pandas as pd

from sklearn.ensemble import AdaBoostClassifier
from sklearn.model_selection import train_test_split
# -

# Grab the linking functions file from github and save locally for Colab.
# We'll import our previously used linking functions from this file.
import linking_tutorial_functions as tutorial

# ## Load Training Data and Ground Truth Labels

df_A, df_B, df_ground_truth = tutorial.load_febrl_training_data(COLAB)

# ## Data Augmentation

for df in [df_A, df_B]:
    df = tutorial.augment_data(df)

# ## Blocking

candidate_links = tutorial.block(df_A, df_B)

# ## Comparing

# +
# %%time

features = tutorial.compare(candidate_links, df_A, df_B)
# -

features.head()

# ## Add Labels to Feature Vectors

# We've augmented, blocked, and compared, so now we're ready to train a classification model which can score candidate record pairs on how likely it is that they are a link. As we did when classifying links via SimSum, we'll append our ground truth values to the features DataFrame.

# +
df_ground_truth["ground_truth"] = df_ground_truth["ground_truth"].apply(lambda x: 1.0 if x else 0.0)

df_labeled_features = pd.merge(
    features,
    df_ground_truth,
    on=["person_id_A", "person_id_B"],
    how="left"
)

df_labeled_features["ground_truth"].fillna(0, inplace=True)
df_labeled_features.head()
# -

# ## Separate Candidate Links into Train/Test

# Next, we'll separate our features DataFrame into a train and test set.

# +
X = df_labeled_features.drop("ground_truth", axis=1)
y = df_labeled_features["ground_truth"]

X_train, X_test, y_train, y_test = train_test_split(
    X, y, stratify=y, test_size=0.2
)
# -

# ## Train ML Classifier

# Though we're using a very simple machine learning model here, the important takeaway is to think of the classification step as a black box that produces a score indicating how likely the model thinks a given candidate record pair is a link. There must be an output score, but *how* that score is generated provides a lot of flexibility. Perhaps you just want to use SimSum, which could be considered an extremely simple "model". Maybe you want to build a neural net to ingest the comparison vectors and produce a score. Generally, in linking, the classification model is the simplest piece, and much more work will go into your blockers and comparators.

classifier = AdaBoostClassifier(n_estimators=64, learning_rate=0.5)

classifier.fit(X_train, y_train)

# ## Predict Using ML Classifier

# Here, we'll generate scores for our test set, and format those predictions in a form useful for evaluation.

y_pred = classifier.predict_proba(X_test)[:,1]

df_predictions = X_test.copy()
df_predictions["model_score"] = y_pred
df_predictions["ground_truth"] = y_test

# ## Choosing a Linking Model Score Threshold

# As with SimSum, we're able to examine the resulting score distribution and precision/recall vs. model score threshold plot to determine where the cutoff should be set.

# ### Model Score Distribution

tutorial.plot_model_score_distribution(df_predictions)

# ### Precision and Recall vs. Model Score

df_eval = tutorial.evaluate_linking(df_predictions)

tutorial.plot_precision_recall_vs_threshold(df_eval)

tutorial.plot_f1_score_vs_threshold(df_eval)

# ### Top Scoring Non-Links

# +
display_cols = [
    "first_name", "surname",
    "street_number", "address_1", "address_2", "suburb", "postcode", "state",
    "date_of_birth", "age", "phone_number", "soc_sec_id",
    "soundex_surname", "soundex_firstname",
    "nysiis_surname", "nysiis_firstname",
]

display_cols = [[f"{col}_A", f"{col}_B"] for col in display_cols]
display_cols = list(itertools.chain.from_iterable(display_cols))

# +
df_top_scoring_negatives = df_predictions[
    df_predictions["ground_truth"] == False
][["model_score", "ground_truth"]].sort_values("model_score", ascending=False).head(n=10)

df_top_scoring_negatives = tutorial.augment_scored_pairs(df_top_scoring_negatives, df_A, df_B)

with pd.option_context('display.max_columns', None):
    display(df_top_scoring_negatives[["model_score", "ground_truth"] + display_cols])
# -

# ### Lowest Scoring True Links

# +
df_lowest_scoring_positives = df_predictions[
    df_predictions["ground_truth"] == True
][["model_score", "ground_truth"]].sort_values("model_score").head(n=10)

df_lowest_scoring_positives = tutorial.augment_scored_pairs(df_lowest_scoring_positives, df_A, df_B)

with pd.option_context('display.max_columns', None):
    display(df_lowest_scoring_positives[["model_score", "ground_truth"] + display_cols])
# -



StarcoderdataJupyterScriptsDedupFiltered