Repository: databricks-industry-solutions/auto-data-linkage (language: python)

Issue:
Rename to arc



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 2e2e46b..416ccf9 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -10,6 +10,8 @@ jobs:
     - name: install pandoc
       run: sudo apt-get install pandoc
     - uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
     - uses: actions/checkout@v2
       with:
         ref: main
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..8aec6ab
--- /dev/null
+++ b/README.md
@@ -0,0 +1,1 @@
+test_readme
diff --git a/arc/archive/legacy_autolinker_notebook.py b/arc/archive/legacy_autolinker_notebook.py
new file mode 100644
index 0000000..8a4d59b
--- /dev/null
+++ b/arc/archive/legacy_autolinker_notebook.py
@@ -0,0 +1,763 @@
+# Databricks notebook source
+# MAGIC %pip install --quiet splink mlflow hyperopt
+
+# COMMAND ----------
+
+from pyspark.sql import functions as F
+from pyspark.sql.types import *
+
+from splink.spark.spark_linker import SparkLinker
+from splink.databricks.enable_splink import enable_splink
+import splink.spark.spark_comparison_library as cl
+
+import hyperopt
+from hyperopt import fmin, tpe, hp, SparkTrials, STATUS_OK, Trials
+from hyperopt.pyll import scope
+
+from sklearn.mixture import GaussianMixture
+
+from pyspark.ml.evaluation import BinaryClassificationEvaluator
+
+import pandas as pd
+import numpy as np
+import itertools
+import math
+import random
+from datetime import datetime
+from pprint import pprint
+
+import mlflow
+
+# COMMAND ----------
+
+class AutoLinker:
+  """
+  Class to create object autolinker for automated data linking
+  """
+  
+  def __init__(self, schema, experiment_name):
+#     self.spark = spark
+    self.schema = schema
+    self.experiment_name = experiment_name
+    self.training_columns = None
+    self.deterministic_columns = None
+    self.df_true_positives = None
+    
+    spark.sql(f"USE {schema}")
+    mlflow.set_experiment(experiment_name)
+
+  
+  
+  def _calculate_column_entropy(self, data, column):
+    """
+    Method to calculate column entropy given a dataset, based on:
+    Entropy = -SUM(P*ln(P))
+    where P is the normalised count of the occurrence of each unique value in the column
+    Parameters
+    : data : Spark Dataframe
+    : column : name of column to calculate
+    Returns
+      - Float
+    """
+    
+    # calculate normalised value count per unique value in colunn
+    rowcount = data.count()
+    vc = data.groupBy(column).count().withColumn("norm_count", F.col("count")/rowcount)
+    
+#     _udf_entropy = F.udf(lambda x: self._calculate_p(x), DoubleType())
+    
+    # Calculate P*ln(P) per row
+    vc = vc.withColumn("entropy", F.col("norm_count")*F.log(F.col("norm_count")))
+    
+    # Entropy = -SUM(P*ln(P))
+    entropy = -vc.select(F.sum(F.col("entropy")).alias("sum_entropy")).collect()[0].sum_entropy
+
+    
+    return entropy
+  
+    
+  def _generate_rules(self, columns):
+    """
+    Method to create a list of Splink-compatible SQL statements from a list of candidate columns for rules.
+    Can be used for simple blocking rules (no AND statement), deterministic rules or for training rules
+    Parameters
+    : columns : list of str containing valid columns in the data set
+    Returns
+      - List of strings of Splink-compatible SQL statements
+    """
+
+    rules = []
+    for column in columns:
+      if "&" in column:
+        subcols = column.split("&")
+        subcols = [subcol.strip() for subcol in subcols]
+
+        rule = " AND ".join([f"l.{subcol} = r.{subcol}" for subcol in subcols])
+      else:
+        rule = f"l.{column} = r.{column}"
+
+      rules.append(rule)
+
+    return rules
+  
+  
+  def _generate_candidate_blocking_rules(self, data, attribute_columns, comparison_size_limit, max_columns_per_rule=2):
+    """
+    Method to automatically generate a list of lists of blocking rules to test, given a user-defined limit for
+    pair-wise comparison size.
+    Parameters
+    : data : Spark DataFrame containing all records to de-duplicate
+    : attribute_columns : list of strings with valid column names of data, containing all possible columns to block on
+    : comparison_size_limit : integer denoting the maximum number of pairs we want to compare - to limit hardware issues
+    : max_columns_per_rule : integer denoting the maximum number of column comparisons in a single rule to try
+    Returns
+      - nested list of lists with Splink-compatible blocking rule queries
+    """
+
+    # initialise empty list to store all combinations in
+    blocking_combinations = []
+
+    # create list of all possible rules
+    for n in range(1, max_columns_per_rule+1):
+      combs = list(itertools.combinations(attribute_columns, n))
+
+      rules = ["&".join(comb) for comb in combs]
+      blocking_combinations.extend(rules)
+
+    # generate Splink-compatible blocking rules
+    blocking_rules_to_generate_predictions = self._generate_rules(blocking_combinations)
+
+    # create dummy linker to count number of combinations
+    linker = SparkLinker(data, spark=spark)
+    settings = {
+      "retain_intermediate_calculation_columns": True,
+      "retain_matching_columns": True,
+      "link_type": "dedupe_only",
+      "unique_id_column_name": "uid",
+      "em_convergence": 0.01,
+      "blocking_rules_to_generate_predictions": blocking_rules_to_generate_predictions
+    }
+    linker.initialise_settings(settings)
+
+    # create dictionary of all possible combinations and their sizes
+    comp_size = [linker.count_num_comparisons_from_blocking_rule(i) for i in blocking_rules_to_generate_predictions]
+    comp_size_dict = {r: s for r, s in zip(blocking_rules_to_generate_predictions, comp_size)}
+
+    # loop through all combinations of combinations and save those which remain under the limit (this is a bit brute force)
+    accepted_rules = []
+    for r in range(1, len(blocking_rules_to_generate_predictions)+1):
+      for c in itertools.combinations(blocking_rules_to_generate_predictions, r):
+        if sum([v for k, v in comp_size_dict.items() if k in c])<=comparison_size_limit:
+          accepted_rules.append(list(c))
+
+    # clean up intermediate tables (not sure if needed)
+    x = spark.sql("show tables like '*__splink__*'").collect()
+    for _ in x:
+        spark.sql(f"drop table {_.tableName}") 
+
+    return accepted_rules
+
+  
+  def _create_hyperopt_space(self, data, attribute_columns, comparison_size_limit, max_columns_per_rule=2):
+    """
+    Method to create hyperopt space for comparison and blocking rule hyperparameters from list of columns.
+    Takes a given (or generated) list of columns and produces a dictionary that can be converted to a comparison list later on,
+    with function names and threhsolds, and a list (of lists) of blocking rules to test.
+    Parameters
+    : data : Spark DataFrame containing all records to de-duplicate
+    : attribute_columns : list of strings with valid column names of data, containing all possible columns to block on or compare
+    : comparison_size_limit : integer denoting the maximum number of pairs we want to compare - to limit hardware issues
+    : max_columns_per_rule : integer denoting the maximum number of column comparisons in a single rule to try
+    Returns
+      - hyperopt.pyll graph for hyperopt parameter search
+    """
+
+    # Generate candidate blocking rules
+    self.blocking_rules = self._generate_candidate_blocking_rules(
+      data=data,
+      attribute_columns=attribute_columns,
+      comparison_size_limit=comparison_size_limit,
+      max_columns_per_rule=max_columns_per_rule
+    )
+
+    # Create comparison dictionary using Hyperopt sampling
+    space = dict()
+    comparison_space = dict()
+
+    # Create hyperopt space for comparisons
+    for column in attribute_columns:
+      comparison_space.update(
+        {
+          column: {
+            "distance_function": hp.choice(
+              f"{column}|distance_function", [
+                {"distance_function": "levenshtein", "threshold": hp.quniform(f"{column}|levenshtein|threshold", 1, 5, q=1)},
+                {"distance_function": "jaccard", "threshold": hp.uniform(f"{column}|jaccard|threshold", 0.7, 0.99)},
+                {"distance_function": "jaro_winkler", "threshold": hp.uniform(f"{column}|jaro_winkler|threshold", 0.7, 0.99)}
+              ]
+            )
+          }
+        }
+      )
+
+    # Create hyperopt space for blocking rules
+    space.update({
+      "blocking_rules": hp.choice("blocking_rules", self.blocking_rules),
+      "comparisons": comparison_space
+    })
+
+    return space
+  
+  def _drop_intermediate_tables(self):
+    # Drop intermediate tables in schema for consecutive runs
+    x = spark.sql("show tables like '*__splink__*'").collect()
+    for _ in x:
+      spark.sql(f"drop table {_.tableName}") 
+        
+  
+  def _create_comparison_list(self, space):
+    """
+    Method to convert comparisons dictionary generated by hyperopt to a Splink-compatible list of
+    spark_comparison_library functions with the generated columns and thresholds.
+    Parameters
+    : comparison_dict : nested dicitonary generated by create_comparison_dict
+    Returns
+      - List of spark_comparison_library method instances
+    """
+
+    # Initialise empty list to populate
+    comparison_list = []
+
+    # Get comparison dict
+    comparison_dict = space["comparisons"]
+
+    # Loop through comparison dictionary and generate Splink-compatible comparison list
+    for column in comparison_dict.keys():
+      distance_function = comparison_dict[column]["distance_function"]["distance_function"]
+      threshold = comparison_dict[column]["distance_function"]["threshold"]
+
+      if distance_function=="levenshtein":
+        comparison_list.append(cl.levenshtein_at_thresholds(column, threshold))
+      elif distance_function=="jaccard":
+        comparison_list.append(cl.jaccard_at_thresholds(column, threshold))
+      elif distance_function=="jaro_winkler":
+        comparison_list.append(cl.jaro_winkler_at_thresholds(column, threshold))
+      else:
+        raise ValueError(f"Unknown distance function {distance_function} passed.")
+
+    return comparison_list
+  
+
+  def deduplicate_records(self, data, predictions, linker, attribute_columns, unique_id, threshold=0.9):
+    """
+    Method to deduplicate the original dataset given the predictions dataframe, its linker
+    object and an optional threshold. The clusters are grouped on and duplicates are removed (arbitrarily).
+    These are then replaced in the original dataset.
+    Parameters
+    : data : Spark Dataframe of the original data set
+    : predictions : Splink Dataframe that's a result of a linker.predict() call
+    : linker : the linker object (splink.SparkLinker instance), needed to cluster
+    : threshold : Float denoting the probability threshold above which a pair is considered a match
+    Returns
+      - Spark Dataframe with deduplicated records
+    """
+    
+    # Cluster records that are matched above threshold
+    clusters = linker.cluster_pairwise_predictions_at_threshold(predictions, threshold_match_probability=threshold)
+    df_predictions = clusters.as_spark_dataframe()
+
+    
+    # Deduplicate predictions by grouping on cluster_id
+    df_deduped_preds = df_predictions.groupBy("cluster_id").agg(
+      *(F.first(i).alias(i) for i in attribute_columns)
+    )
+
+    # Left anti-join predictions on original data to get all records outside of blocking rules
+    data_without_predictions = data.join(
+      df_predictions,
+      data.uid==df_predictions.uid,
+      how="left_anti"
+    )
+
+    # Union deduped predictions and data outside of blocking rules
+    df_deduped = df_deduped_preds \
+      .select(attribute_columns) \
+      .union(
+        data_without_predictions \
+          .select(attribute_columns)
+    )
+
+    return df_deduped
+
+
+  def calculate_entropy_delta(self, data, deduped, columns):
+    """
+    Method to calculate the change in entropy of a set of columns between two
+    datasets, used to calculate the information gain after deduping a set of records.
+    Parameters
+    : data : Spark DataFrame, the original (undeduplicated) data set of records
+    : deduped : Spark DataFrame, the new (deduplicated) data set of records
+    : columns : list of strings with valid column names to compare
+    Returns
+      - Float (average change in column entropy)
+    """
+    
+    # Initialise empty list to populate with entropy deltas
+    entropy_deltas = list()
+    
+    # For each column in given columns, calculate entropy in original dataset
+    # and the deduplicated dataset. The delta is the difference between new and old.
+    for column in columns:
+      original_entropy = self._calculate_column_entropy(data, column)
+      new_entropy = self._calculate_column_entropy(deduped, column)
+      entropy_deltas.append(new_entropy-original_entropy)
+
+
+    return np.mean(entropy_deltas)
+  
+  
+  
+  def calculate_empirical_score(self, data, predictions, threshold):
+    """
+    Method to calculate precision, recall and F1 score based on ground truth labels.
+    Assumes the data has a column with empirical IDs to assess whether two records belong
+    to the same real life entity. Will check if this has already been calculated. If not,
+    it will create it for the first (and only) time.
+    NB: this includes a bit of hard coding, but it won't make it to production anyway because
+    we won't have ground truth.
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : predictions : Splink DataFrame with predicted pairs
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    Returns
+      - 3-tuple of floats for precision, recall and F1 score
+    """
+    if not self.df_true_positives:
+      # filter original data to only contain rows where recid id appears more than once (known dupes)
+      data_recid_groupby = data.groupBy("recid").count().filter("count>1").withColumnRenamed("recid", "recid_")
+      data_tp = data.join(data_recid_groupby, data.recid==data_recid_groupby.recid_, how="inner").drop("recid_")
+
+      data_l = data_tp.withColumnRenamed("uid", "uid_l").withColumnRenamed("recid", "recid_l").select("uid_l", "recid_l")
+      data_r = data_tp.withColumnRenamed("uid", "uid_r").withColumnRenamed("recid", "recid_r").select("uid_r", "recid_r")
+
+      dt = data_l.join(data_r, data_l.uid_l!=data_r.uid_r, how="inner")
+      # create boolean col for y_true
+      df_true = dt.withColumn("match", F.when(F.col("recid_l")==F.col("recid_r"), 1).otherwise(0))
+      # only keep matches
+      df_true = df_true.filter("match=1")
+      # assign as attribute to avoid re-calculation
+      self.df_true_positives = df_true
+      
+    # convert predictions to Spark DataFrame and filter on match prob threshold - table will only contain predicted positives
+    df_pred = predictions.as_spark_dataframe().filter(f"match_probability>={threshold}")
+    
+    # Calculate TP, FP, FN, TN
+    
+    # TP is the inner join of true and predicted pairs
+    tp = df_pred.join(
+      self.df_true_positives,
+      ((df_pred.uid_l==self.df_true_positives.uid_l) & (df_pred.uid_r==self.df_true_positives.uid_r)) | ((df_pred.uid_l==self.df_true_positives.uid_r) & (df_pred.uid_r==self.df_true_positives.uid_l)),
+      how="inner"
+    ).count()
+    
+    # FN is the left anti-join of true and predicted
+    fn = self.df_true_positives.join(
+      df_pred,
+      ((df_pred.uid_l==self.df_true_positives.uid_l) & (df_pred.uid_r==self.df_true_positives.uid_r)) | ((df_pred.uid_l==self.df_true_positives.uid_r) & (df_pred.uid_r==self.df_true_positives.uid_l)),
+      how="left_anti"
+    ).count()
+    
+    # FP is the left anti-join of predicted and true
+    fp = df_pred.join(
+      self.df_true_positives,
+      ((df_pred.uid_l==self.df_true_positives.uid_l) & (df_pred.uid_r==self.df_true_positives.uid_r)) | ((df_pred.uid_l==self.df_true_positives.uid_r) & (df_pred.uid_r==self.df_true_positives.uid_l)),
+      how="left_anti"
+    ).count()
+    
+    # TN is everything else, i.e. N(N-1)-TP-FN-FP
+    N = data.count()
+    tn = (N*(N-1))-tp-fn-fp
+    
+    # Calculate precision, recall and f1
+    precision = tp/(tp+fp)
+    recall = tp/(tp+fn)
+
+    if (precision+recall)>0:
+      f1 = precision*recall/(precision+recall)
+    else:
+      f1 = 0.0
+      
+    return precision, recall, f1
+
+  def train_linker(self, data, space, attribute_columns, unique_id, deterministic_columns=None, training_columns=None):
+    """
+    Method to train a linker model
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : space : dictionary generated by hyperopt sampling
+    : attribute_columns : list of strings containing all attribute columns
+    : unique_id : string with the name of the unique ID column
+    : deterministic columns : list of strings containint columns to block on
+    : training_columns : list of strings containing training columns
+    Returns
+      - tuple of trained linker object and Splink dataframe with predictions
+    """
+
+    # Drop any intermediate tables for a clean run
+    self._drop_intermediate_tables()
+    
+    # Set up parameters
+
+    # Get blocking rules from hyperopt space
+    blocking_rules_to_generate_predictions = list(space["blocking_rules"])
+
+    # Create comparison list from hyperopt space
+    comparisons = self._create_comparison_list(space)
+
+    # if deterministic columns are not set, pick 2 at random from attribute columns
+    # and save them as an attribute so that they can remain consistent between runs
+    if not self.deterministic_columns:
+      deterministic_columns = random.sample(attribute_columns, 2)
+      self.deterministic_columns = deterministic_columns
+
+    deterministic_rules = self._generate_rules(deterministic_columns)
+
+    # if deterministic columns are not set, pick 2 at random from attribute columns
+    # and save them as an attribute so that they can remain consistent between runs
+    if not training_columns:
+      training_columns = random.sample(attribute_columns, 2)
+      self.training_columns = training_columns
+      
+    training_rules = self._generate_rules(training_columns)
+
+    # create settings dict
+    settings = {
+      "retain_intermediate_calculation_columns": True,
+      "retain_matching_columns": True,
+      "link_type": "dedupe_only",
+      "unique_id_column_name": unique_id,
+      "comparisons": comparisons,
+      "em_convergence": 0.01,
+      "blocking_rules_to_generate_predictions": blocking_rules_to_generate_predictions
+    }
+
+    # Train linker model
+    linker = SparkLinker(data, spark=spark)
+    linker.initialise_settings(settings)
+    linker.estimate_probability_two_random_records_match(deterministic_rules, recall=0.8)
+    linker.estimate_u_using_random_sampling(target_rows=1e7)
+    for training_rule in training_rules:
+      linker.estimate_parameters_using_expectation_maximisation(training_rule)
+
+    # Make predictions
+    predictions = linker.predict()
+
+
+    return linker, predictions
+  
+  def evaluate_linker(self, data, predictions, threshold, attribute_columns, unique_id, linker):
+    """
+    Method to evaluate predictions made by a trained linker model
+    Parameters
+    : data : Spark Dataframe containing the original dataset (required to establish ground truth labels)
+    : df_predictions : Spark DataFrame containing pair-wise predictions made by a linker model
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    : attribute_columns : list of strings with valid column names to compare the entropies of
+    : linker : isntance of splink.SparkLinker, used for deduplication (clustering)
+    Returns
+      - Dictionary of evaluation metrics TODO: make the metrics controllable via an argument
+    """
+
+    # Calculate empirical scores
+    precision, recall, f1 = self.calculate_empirical_score(data, predictions, threshold)
+    
+    # Deduplicate data
+    deduped = self.deduplicate_records(data, predictions, linker, attribute_columns, unique_id, threshold)
+    
+    # Calculate mean change in entropy
+    mean_entropy_change = self.calculate_entropy_delta(data, deduped, attribute_columns)
+    
+    evals = {
+      "precision": precision,
+      "recall": recall,
+      "f1": f1,
+      "mean_entropy_change": mean_entropy_change
+    }
+
+    return evals
+  
+  
+  def train_and_evaluate_linker(self, data, space, attribute_columns, unique_id, deterministic_columns=None, training_columns=None, threshold=0.9):
+    """
+    Method to train and evaluate a linker model in one go
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : space : dictionary generated by hyperopt sampling
+    : attribute_columns : list of strings containing all attribute columns
+    : unique_id : string with the name of the unique ID column
+    : deterministic columns : list of strings containint columns to block on - if None, they will be generated automatically/randomly
+    : training_columns : list of strings containing training columns - if None, they will be generated automatically/randomly
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    Returns
+      - 6-tuple of trained linker (SparkLinker instance), predictions (Splink DAtaFrame) and metrics (4x float)
+    """
+
+    # Train model
+    linker, predictions = self.train_linker(data, space, attribute_columns, unique_id, deterministic_columns, training_columns)
+
+    # Evaluate model
+    evals = self.evaluate_linker(data, predictions, threshold, attribute_columns, unique_id, linker)
+
+
+    return linker, predictions, evals
+  
+
+
+  def log_to_mlflow(self, trials):
+    """
+    Method to log parameters and loss function of each trial as an MLflow run.
+    Parameters
+    : trials : hyperopt.Trials object created after a hyperopt run
+    Returns:
+      - None
+    """
+    
+    for t in trials.trials:
+      with mlflow.start_run():
+        params = t['misc']['vals']
+        blocking_rules = t['misc']['vals']['blocking_rules'][0]
+        comparison_params = {}
+
+        for k, v in params.items():
+          if ("threshold" in k) & (len(v)>0):
+            column, function, _ = k.split('|')
+            comparison_params.update({
+              column: {
+                "function": function,
+                "threshold": v[0]
+              }
+            })
+
+        params = {
+          "blocking_rules": self.blocking_rules[blocking_rules],
+          "comparisons": comparison_params,
+          "deterministic_rules": self._generate_rules(self.deterministic_columns),
+          "training_rules": self._generate_rules(self.training_columns),
+        }
+        metrics = {k:v for k, v in t["result"].items() if not k=="status"}
+        mlflow.log_dict(params, "params.json")
+        mlflow.log_params(params)
+        mlflow.log_metrics(metrics)
+
+    return None
+  
+  
+  def auto_link(self, data, attribute_columns, unique_id, comparison_size_limit, max_evals, deterministic_columns=None, training_columns=None, threshold=0.9):
+    """
+    Method to run a series of hyperopt trials.
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : space : dictionary generated by hyperopt sampling
+    : attribute_columns : list of strings containing all attribute columns
+    : unique_id : string with the name of the unique ID column
+    : comparison_size_limit : int denoting maximum size of pairs allowed after blocking
+    : max_evals : int denoting max number of hyperopt trials to run
+    : deterministic columns : list of strings containint columns to block on - if None, they will be generated automatically/randomly
+    : training_columns : list of strings containing training columns - if None, they will be generated automatically/randomly
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    Returns
+      - None
+    """
+    
+    #set start time for measuring duration
+    start = datetime.now()
+    
+    # define objective function
+    def tune_model(space):
+      linker, predictions, evals = self.train_and_evaluate_linker(
+        data,
+        space,
+        attribute_columns,
+        unique_id="uid",
+        deterministic_columns=self.deterministic_columns,
+        training_columns=self.training_columns,
+        threshold=threshold
+      )
+      
+      loss = evals["mean_entropy_change"]
+      
+      result = {'loss': loss, 'status': STATUS_OK}
+      for k, v in evals.items():
+        result.update({k:v})
+
+      return result
+    
+    # initialise trials and create hyperopt space
+    self.trials = Trials()
+    space = self._create_hyperopt_space(data, attribute_columns, comparison_size_limit)
+
+    # run hyperopt trials
+    self.best = fmin(
+      fn=tune_model,
+      space=space,
+      algo=tpe.suggest,
+      max_evals=max_evals,
+      trials=self.trials
+    )
+    
+    # log trials to MLflow
+
+    self.log_to_mlflow(self.trials)
+    
+    # Set best params, metrics and results
+    t = self.trials.best_trial
+    self.best_metric = -t['result']['loss']
+    params = t['misc']['vals']
+    blocking_rules = t['misc']['vals']['blocking_rules'][0]
+    comparison_params = {}
+
+    for k, v in params.items():
+      if ("threshold" in k) & (len(v)>0):
+        column, function, _ = k.split('|')
+        comparison_params.update({
+          column: {
+            "function": function,
+            "threshold": v[0]
+          }
+        })
+
+    self.best_params = {
+      "blocking_rules": self.blocking_rules[blocking_rules],
+      "comparisons": comparison_params
+    }
+    
+    # TODO: move this into its own method, probably - refactor...
+    # convert best params to a format that can retrain a new linker 
+    # model based on the best params. This is necessary because the underlying tables may have been overwritten
+    best_comparisons = {}
+    for k, v in comparison_params.items():
+      best_comparisons.update({
+        k: {"distance_function": {"distance_function": v['function'], "threshold": v['threshold']}}
+      })
+      
+    best_params_for_rt = {
+      "blocking_rules": self.best_params['blocking_rules'],
+      "comparisons": best_comparisons
+    }
+    
+    self.best_linker, self.best_predictions = self.train_linker(data, best_params_for_rt, attribute_columns, unique_id, self.deterministic_columns, self.training_columns)
+    
+    #end training duration
+    end = datetime.now()
+    training_duration = (end-start).seconds
+    
+    # return succes text
+    success_text = f"""
+    ======================================================================================
+    AutoLinking completed. Carried out {max_evals} trials in {training_duration} seconds.
+    ======================================================================================
+    Deterministic rules used
+    ------------------------
+    {', '.join(self._generate_rules(self.deterministic_columns))}
+    ------------------------
+    Training rules used
+    ------------------------
+    {', '.join(self._generate_rules(self.training_columns))}
+    ------------------------
+    Best parameters
+    ------------------------
+    {self.best_params}
+    ------------------------
+    Best F1: {self.best_metric}
+    ------------------------
+    
+    You can now access the best linker model and predictions and use Splink's built-in functionality, e.g.:
+    
+    >>> linker = autolinker.best_linker
+    >>> predictions = autolinker.best_predictions
+    >>> ...
+    """
+    
+    print(success_text)
+    
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # North Carolina Voters dataset
+# MAGIC https://dbs.uni-leipzig.de/research/projects/object_matching/benchmark_datasets_for_entity_resolution
+
+# COMMAND ----------
+
+data = spark.read.table("marcell_autosplink.voters_data_sample")
+
+# COMMAND ----------
+
+autolinker = AutoLinker(
+  #spark=spark,                                                                                               # spark instance
+  schema="marcell_autosplink",                                                                               # schema to write results to
+  experiment_name="/Users/marcell.ferencz@databricks.com/autosplink/evaluate/autosplink_voters_600_rerun"  # MLflow experiment location
+)
+
+# COMMAND ----------
+
+autolinker.auto_link(
+  data=data,                                                         # dataset to dedupe
+  attribute_columns=["givenname", "surname", "suburb", "postcode"],  # columns that contain attributes to compare
+  unique_id="uid",                                                   # column name of the unique ID
+  comparison_size_limit=150000,                                      # Maximum number of pairs when blocking applied
+  max_evals=50                                                      # Maximum number of hyperopt trials to run
+)
+
+# COMMAND ----------
+
+linker = autolinker.best_linker
+
+# COMMAND ----------
+
+linker.m_u_parameters_chart()
+
+# COMMAND ----------
+
+predictions = autolinker.best_predictions
+
+predictions.as_spark_dataframe().display()
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # Music albums dataset
+# MAGIC https://dbs.uni-leipzig.de/research/projects/object_matching/benchmark_datasets_for_entity_resolution
+
+# COMMAND ----------
+
+data = spark.read.table("marcell_autosplink.music_data").sample(0.5)
+
+# COMMAND ----------
+
+data = data \
+  .withColumn("title", F.substring(F.col("title"), 0, 20)) \
+  .withColumn("artist", F.substring(F.col("artist"), 0, 20)) \
+  .withColumn("album", F.substring(F.col("album"), 0, 20))
+
+# COMMAND ----------
+
+autolinker = AutoLinker(
+  #spark=spark,                                                                                               # spark instance
+  schema="marcell_autosplink",                                                                               # schema to write results to
+  experiment_name="/Users/marcell.ferencz@databricks.com/autosplink/evaluate/autosplink_music_experiment"  # MLflow experiment location
+)
+
+# COMMAND ----------
+
+autolinker.auto_link(
+  data=data,                                                         # dataset to dedupe
+  attribute_columns=["title", "length", "artist", "album", "year", "language"],  # columns that contain attributes to compare
+  unique_id="uid",                                                   # column name of the unique ID
+  comparison_size_limit=200000,                                      # Maximum number of pairs when blocking applied
+  max_evals=100                                                      # Maximum number of hyperopt trials to run
+)
+
+# COMMAND ----------
+
+
diff --git a/arc/autolinker/__init__.py b/arc/autolinker/__init__.py
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/arc/autolinker/__init__.py
@@ -0,0 +1,1 @@
+
diff --git a/arc/autolinker/autolinker.py b/arc/autolinker/autolinker.py
new file mode 100644
index 0000000..b04df46
--- /dev/null
+++ b/arc/autolinker/autolinker.py
@@ -0,0 +1,910 @@
+from pyspark.sql import functions as F
+from pyspark.sql.types import *
+from pyspark.sql import Window
+import pyspark
+
+# from dbruntime.display import displayHTML
+
+import splink
+from splink.spark.spark_linker import SparkLinker
+from splink.databricks.enable_splink import enable_splink
+import splink.spark.spark_comparison_library as cl
+
+import hyperopt
+from hyperopt import fmin, tpe, hp, SparkTrials, STATUS_OK, Trials
+from hyperopt.pyll import scope
+
+from . import splink_mlflow
+
+import numpy as np
+import itertools
+import math
+import random
+from datetime import datetime
+
+import typing
+
+import mlflow
+
+
+class AutoLinker:
+  """
+  Class to create object arc for automated data linking.
+  
+  :param spark: Spark instance used
+  :param catalog: Unity Catalog name, if used, for intermediate tables
+  :param schema: Schema/database name for intermediate tables
+  :param experiment_name: Location or simply name of the MLflow experiment for storing trial run results
+  :param training_columns: Valid list of columns (or column names separated by '&' for AND combinations) of columns to be fixed during EM training
+  :param deterministic_columns: Valid list of columns (or column names separated by '&' for AND combinations) of columns to be used for prior estimation
+
+  Basic usage:
+  
+  .. code-block:: python
+
+  >>> arc = AutoLinker(
+  ...   catalog="splink_catalog",                # catalog name
+  ...   schema="splink_schema",                  # schema to write results to
+  ...   experiment_name="autosplink_experiment"  # MLflow experiment location
+  ... )
+  >>> arc.auto_link(
+  ...   data=data,                               # dataset to dedupe
+  ...   attribute_columns=["A", "B", "C", "D"],  # columns that contain attributes to compare
+  ...   unique_id="id",                          # column name of the unique ID
+  ...   comparison_size_limit=500000,            # Maximum number of pairs when blocking applied
+  ...   max_evals=100                            # Maximum number of hyperopt trials to run
+  ... )
+  >>> arc.best_linker.m_u_parameters_chart()     # use Splink functionality out-of-the-box
+
+  """
+  
+  
+  def __init__(
+    self,
+    spark:str=None,
+    catalog:str=None,
+    schema:str=None,
+    experiment_name:str=None,
+    training_columns:list=None,
+    deterministic_columns:list=None
+  ) -> None:
+    """
+    Initialises an AutoLinker instance to perform automated record linking.
+    """
+    self.spark = spark 
+    self.catalog = catalog# if catalog else spark.catalog.currentCatalog()
+    self.schema = schema  #. if schema else spark.catalog.currentDatabase()
+    
+    #self.username = spark.sql('select current_user() as user').collect()[0]['user']
+    self.experiment_name = experiment_name #if experiment_name else f"/Users/{self.username}/Databricks Autolinker {str(datetime.now())}"
+    self.training_columns = training_columns
+    self.deterministic_columns = deterministic_columns
+    self.best_params = None
+    self.clusters=None
+    self.cluster_threshold=None
+    self.original_entropies = dict()
+    
+
+  def __str__(self):
+    return f"AutoLinker instance working in {self.catalog}.{self.schema} and MLflow experiment {self.experiment_name}"
+  
+  
+  
+  def _calculate_column_entropy(
+    self,
+    data: pyspark.sql.DataFrame,
+    column: str
+  ) -> float:
+    """
+    Method to calculate column entropy given a dataset, based on:
+    :math:`Entropy = -SUM(P*ln(P))`
+    where P is the normalised count of the occurrence of each unique value in the column.
+    
+    :param data: input dataframe with row per record
+    :param column: (valid) name of column to calculate entropy on
+
+    """
+    
+    # calculate normalised value count per unique value in column
+    rowcount = data.count()
+    vc = data.groupBy(column).count().withColumn("norm_count", F.col("count")/rowcount)
+
+    
+    # Calculate P*ln(P) per row
+    vc = vc.withColumn("entropy", F.col("norm_count")*F.log(F.col("norm_count")))
+    
+    # Entropy = -SUM(P*ln(P))
+    entropy = -vc.select(F.sum(F.col("entropy")).alias("sum_entropy")).collect()[0].sum_entropy
+
+    
+    return entropy
+  
+  
+    
+  def _generate_rules(
+    self,
+    columns:list
+  ) -> list:
+    """
+    Method to create a list of Splink-compatible SQL statements from a list of candidate columns for rules.
+    Can be used for simple blocking rules (no AND statement), deterministic rules or for training rules.
+    
+    Returns list of strings of Splink-compatible SQL statements.
+
+    :param columns: valid attribute columns in the data set
+
+    """
+
+    rules = []
+    for column in columns:
+      if "&" in column:
+        subcols = column.split("&")
+        subcols = [subcol.strip() for subcol in subcols]
+
+        rule = " AND ".join([f"l.{subcol} = r.{subcol}" for subcol in subcols])
+      else:
+        rule = f"l.{column} = r.{column}"
+
+      rules.append(rule)
+
+    return rules
+  
+  
+  
+  def _generate_candidate_blocking_rules(
+    self,
+    data:pyspark.sql.DataFrame,
+    attribute_columns:list,
+    comparison_size_limit:int,
+    unique_id:str,
+    max_columns_per_rule:int=2
+  ) -> list:
+    """
+    Method to automatically generate a list of lists of blocking rules to test, given a user-defined limit for
+    pair-wise comparison size.
+
+    Returns a nested list of lists with Splink-compatible blocking rule queries.
+
+    :param data: input data with record-per-row
+    :param attribute_columns: valid column names of data, containing all possible columns to block on
+    :param comparison_size_limit: the maximum number of pairs we want to compare, to limit hardware issues
+    :param unique_id: the name of the unique ID column in data
+    :param max_columns_per_rule: the maximum number of column comparisons in a single rule to try
+
+    """
+
+    # initialise empty list to store all combinations in
+    blocking_combinations = []
+
+    # create list of all possible rules
+    for n in range(1, max_columns_per_rule+1):
+      combs = list(itertools.combinations(attribute_columns, n))
+      blocking_combinations.extend(combs)
+
+    # Convert them to Splink-compatible rules
+    blocking_rules = self._generate_rules(["&".join(r) for r in blocking_combinations])
+    
+    # Estimate number of pairs per blocking rule by grouping on blocking rule columns and getting the sum of squares of counts
+    comp_size_dict = dict()
+
+    for comb, rule in zip(blocking_combinations, blocking_rules):
+      num_pairs = data.groupBy(list(comb)).count().select(F.sum(F.col("count")*F.col("count"))).collect()[0]["sum((count * count))"]
+      comp_size_dict.update({rule: num_pairs})
+    
+    # loop through all combinations of combinations and save those which remain under the limit (this is a bit brute force)
+    accepted_rules = []
+    for r in range(1, len(blocking_rules)+1):
+      for c in itertools.combinations(blocking_rules, r):
+        if sum([v for k, v in comp_size_dict.items() if k in c])<=comparison_size_limit:
+          accepted_rules.append(list(c))
+
+
+    if len(accepted_rules)<1:
+      raise ValueError("No blocking rules meet the comparison size limit.")
+          
+    return accepted_rules
+
+  
+  
+  def _create_hyperopt_space(
+    self,
+    data:pyspark.sql.DataFrame,
+    attribute_columns:list,
+    comparison_size_limit:int,
+    unique_id:str,
+    max_columns_per_rule:int=2
+  ) -> dict:
+    """
+    Method to create hyperopt space for comparison and blocking rule hyperparameters from list of columns.
+    Takes a given (or generated) list of columns and produces a dictionary that can be converted to a comparison list later on,
+    with function names and threhsolds, and a list (of lists) of blocking rules to test.
+
+    Returns a dictionary for hyperopt parameter search
+
+    :param data: input data with records per row
+    :param attribute_columns: valid column names of data, containing all possible columns to compare
+    :param comparison_size_limit: maximum number of pairs we want to compare, to limit hardware issues
+    :param max_columns_per_rule: the maximum number of column comparisons in a single rule to try
+    
+    """
+
+    # Generate candidate blocking rules
+    self.blocking_rules = self._generate_candidate_blocking_rules(
+      data=data,
+      attribute_columns=attribute_columns,
+      comparison_size_limit=comparison_size_limit,
+      unique_id=unique_id,
+      max_columns_per_rule=max_columns_per_rule
+    )
+
+    # Create comparison dictionary using Hyperopt sampling
+    space = dict()
+    comparison_space = dict()
+
+    # Create hyperopt space for comparisons
+    for column in attribute_columns:
+      comparison_space.update(
+        {
+          column: {
+            "distance_function": hp.choice(
+              f"{column}|distance_function", [
+                {"distance_function": "levenshtein", "threshold": hp.quniform(f"{column}|levenshtein|threshold", 1, 5, q=1)},
+                {"distance_function": "jaccard", "threshold": hp.uniform(f"{column}|jaccard|threshold", 0.7, 0.99)},
+                {"distance_function": "jaro_winkler", "threshold": hp.uniform(f"{column}|jaro_winkler|threshold", 0.7, 0.99)}
+              ]
+            )
+          }
+        }
+      )
+
+    # Create hyperopt space for blocking rules
+    space.update({
+      "blocking_rules": hp.choice("blocking_rules", self.blocking_rules),
+      "comparisons": comparison_space
+    })
+
+    return space
+  
+  
+  
+  def _drop_intermediate_tables(self):
+    """
+    Method to drop intermediate tables for clean consecutive runs.
+    """
+    # Drop intermediate tables in schema for consecutive runs
+    tables_in_schema = self.spark.sql(f"show tables from {self.catalog}.{self.schema} like '*__splink__*'").collect()
+    for table in tables_in_schema:
+      try:
+        self.spark.sql(f"drop table {self.catalog}.{self.schema}.{table.tableName}") 
+      except:
+        self.spark.sql(f"drop table {table.tableName}")
+        
+  
+  def _clean_columns(
+    self,
+    data:pyspark.sql.DataFrame,
+    attribute_columns:list,
+    cleaning="all"):
+    """
+    Method to clean string columns (turn them to lower case and remove non-alphanumeric characters)
+    in order to help with better (and quicker) string-distance calculations. If cleaning is 'all'
+    (as is by default), it will automatically clean as much as it can.  If cleaning is 'none', it will do nothing.
+    cleaning can also be a dictionary with keys as column names and values as lists of method strings.
+    The currently available cleaning methods are turning to lower case and removing non-alphanumeric characters.
+
+    Returns a Spark DataFrame.
+
+    :param data: DataFrame containing the data to be cleaned
+    :param attribute_columns: all attribute columns
+    :param cleaning: string or dicitonary with keys as column names and values as list of valid cleaning method names.
+
+    """
+    
+    # if cleaning is set to "none", don't do anything to the data
+    if cleaning=="none":
+      return data
+    
+    # if it's set to "all", turn it into a dictionary covering all possible cases
+    if cleaning=="all":
+      cleaning = {col: ["lower", "alphanumeric_only"] for col in attribute_columns}
+      
+    for col, methods in cleaning.items():
+      # if column is not a string, skip it
+      if not data.schema[col].dataType == StringType():
+        continue
+        
+      for method in methods:
+        if method=="alphanumeric_only":
+          # replace column and only keep alphanumeric and space characters
+          data = data.withColumn(col, F.regexp_replace(F.col(col), r"[^A-Za-z0-9 ]+", ""))
+        
+        elif method=="lower":
+          data = data.withColumn(col, F.lower(F.col(col)))
+          
+    return data
+  
+
+
+  def _convert_hyperopt_to_splink(self) -> dict:
+    """
+    Method to convert hyperopt trials to a dictionary that can be used to
+    train a linker model. Used for training the best linker model at the end of an experiment.
+    Sets class attributes for best metric and parameters as well.
+
+    Returns a dictionary.
+
+    """
+    # Set best params, metrics and results
+    t = self.trials.best_trial
+    self.best_metric = t['result']['loss']
+    params = t['misc']['vals']
+    blocking_rules = t['misc']['vals']['blocking_rules'][0]
+    comparison_params = {}
+
+    for k, v in params.items():
+      if ("threshold" in k) & (len(v)>0):
+        column, function, _ = k.split('|')
+        comparison_params.update({
+          column: {
+            "function": function,
+            "threshold": v[0]
+          }
+        })
+
+    self.best_params = {
+      "blocking_rules": self.blocking_rules[blocking_rules],
+      "comparisons": comparison_params
+    }
+
+    best_comparisons = {}
+    for k, v in comparison_params.items():
+      best_comparisons.update({
+        k: {"distance_function": {"distance_function": v['function'], "threshold": v['threshold']}}
+      })
+      
+    best_params_for_rt = {
+      "blocking_rules": self.best_params['blocking_rules'],
+      "comparisons": best_comparisons
+    }
+
+    return best_params_for_rt
+
+
+    
+  def _create_comparison_list(
+    self,
+    space:dict
+  ) -> list:
+    """
+    Method to convert comparisons dictionary generated by hyperopt to a Splink-compatible list of
+    spark_comparison_library functions with the generated columns and thresholds.
+
+    Returns list of spark_comparison_library method instances.
+
+    :param space: nested dicitonary generated by create_comparison_dict
+
+    """
+
+    # Initialise empty list to populate
+    comparison_list = []
+
+    # Get comparison dict
+    comparison_dict = space["comparisons"]
+
+    # Loop through comparison dictionary and generate Splink-compatible comparison list
+    for column in comparison_dict.keys():
+      distance_function = comparison_dict[column]["distance_function"]["distance_function"]
+      threshold = comparison_dict[column]["distance_function"]["threshold"]
+
+      if distance_function=="levenshtein":
+        comparison_list.append(cl.levenshtein_at_thresholds(column, threshold))
+      elif distance_function=="jaccard":
+        comparison_list.append(cl.jaccard_at_thresholds(column, threshold))
+      elif distance_function=="jaro_winkler":
+        comparison_list.append(cl.jaro_winkler_at_thresholds(column, threshold))
+      else:
+        raise ValueError(f"Unknown distance function {distance_function} passed.")
+
+    return comparison_list
+  
+  
+
+  def deduplicate_records(
+    self,
+    predictions:splink.splink_dataframe.SplinkDataFrame,
+    linker:splink.spark.spark_linker.SparkLinker,
+    attribute_columns:list,
+    threshold:float=0.9
+  ) -> pyspark.sql.DataFrame:
+    """
+    Method to deduplicate the original dataset given the predictions dataframe, its linker
+    object and an optional threshold. The clusters are grouped on and duplicates are set to the same value (arbitrarily).
+    These are then replaced in the original dataset.
+    
+    Returns a spark Dataframe with deduplicated records.
+
+    :param predictions: Splink Dataframe that's a result of a linker.predict() call
+    :param linker: a trained linker object, needed to cluster
+    :param attribute_columns: list of column names containing attribtues in the data
+    :param threshold: the probability threshold above which a pair is considered a match
+    
+    """
+    
+    # Cluster records that are matched above threshold
+    clusters = linker.cluster_pairwise_predictions_at_threshold(predictions, threshold_match_probability=threshold)
+    # define window to aggregate over
+    window = Window.partitionBy("cluster_id")
+    df_predictions = clusters.as_spark_dataframe()
+    # loop through the attributes and standardise. we want to keep original column names.
+    for attribute_column in attribute_columns:
+      df_predictions = (
+        df_predictions
+        .withColumn(attribute_column, F.first(F.col(attribute_column)).over(window)) # standardise values
+      )
+
+    return df_predictions
+
+  
+
+  def calculate_entropy_delta(
+    self,
+    deduped:pyspark.sql.DataFrame
+  ) -> float:
+    """
+    Method to calculate the change in entropy of a set of columns between two
+    datasets, used to calculate the information gain after deduping a set of records.
+    
+    Returns average change in column entropy.
+
+    :param data: the original (undeduplicated) data set of records
+    :param deduped: the new (deduplicated) data set of records
+    :param columns: valid column names to compare
+    
+    """
+    
+    
+    # Initialise empty list to populate with entropy deltas
+    entropy_deltas = list()
+    dd_cnt=deduped.count()
+    columns = self.attribute_columns
+    if self.original_entropies == dict():
+      for column in columns:
+        self.original_entropies[column] = self._calculate_column_entropy(self.data, column, self.data_rowcount)
+    
+    # For each column in given columns, calculate entropy in original dataset
+    # and the deduplicated dataset. The delta is the difference between new and old.
+    for column in columns:
+      original_entropy = self.original_entropies[column]
+      new_entropy = self._calculate_column_entropy(deduped, column, dd_cnt)
+      entropy_deltas.append(new_entropy-original_entropy)
+
+
+    return np.mean(entropy_deltas)
+  
+  
+  
+  def _randomise_columns(
+    self,
+    attribute_columns:list
+  ) -> list:
+    """
+    Method to randomly select (combinations of) columns from attribute columns for EM training.
+    Will try to pick 2 combinations of 2 (i.e AB and BC from ABC), but will default to 2 if there are only 2.
+
+    Returns list of lists of strings.
+
+    :param attribute_columns: list of strings containing all attribute columns
+    
+    """
+    # if there are only 2 columns, return them both
+    if len(attribute_columns)<3:
+      return attribute_columns
+    
+    # pick 3 random columns
+    cols = random.sample(attribute_columns, 3)
+    
+    # this creates a list like ["A&B", "B&C", "A&C"] - it's useful to have an AND statement in the training
+    # columns, otherwise EM training tables balloon
+    training_columns = ["&".join([cols[0], cols[1]]), "&".join([cols[1], cols[2]]), "&".join([cols[0], cols[2]])]
+    
+    return training_columns
+  
+
+  def train_linker(
+    self,
+    data:pyspark.sql.DataFrame,
+    space:dict,
+    attribute_columns:list,
+    unique_id:str,
+    deterministic_columns:list=None,
+    training_columns:list=None
+    ) -> typing.Tuple[splink.spark.spark_linker.SparkLinker, splink.splink_dataframe.SplinkDataFrame]:
+    """
+    Method to train a linker model.
+
+    Returns a 2-tuple of trained linker object and Splink dataframe with predictions
+
+    :param data: Spark DataFrame containing the data to be de-duplicated
+    :param space: dictionary generated by hyperopt sampling
+    :param attribute_columns: list of strings containing all attribute columns
+    :param unique_id: string with the name of the unique ID column
+    :param deterministic columns: list of strings containint columns to block on
+    :param training_columns: list of strings containing training columns
+    
+    """
+
+
+    # Drop any intermediate tables for a clean run
+    self._drop_intermediate_tables()
+    
+    # Set up parameters
+
+    # Get blocking rules from hyperopt space
+    blocking_rules_to_generate_predictions = list(space["blocking_rules"])
+    print(blocking_rules_to_generate_predictions)
+    # Create comparison list from hyperopt space
+    comparisons = self._create_comparison_list(space)
+
+    # if deterministic columns are not set, pick 2 at random from attribute columns
+    # and save them as an attribute so that they can remain consistent between runs
+    if not self.deterministic_columns:
+      deterministic_columns = random.sample(attribute_columns, 2)
+      self.deterministic_columns = deterministic_columns
+
+    deterministic_rules = self._generate_rules(deterministic_columns)
+
+    # if deterministic columns are not set, pick 2 at random from attribute columns
+    # and save them as an attribute so that they can remain consistent between runs
+    if not training_columns:
+      training_columns = self._randomise_columns(attribute_columns)
+      self.training_columns = training_columns
+      
+    training_rules = self._generate_rules(training_columns)
+
+    # create settings dict
+    settings = {
+      "retain_intermediate_calculation_columns": True,
+      "retain_matching_columns": True,
+      "link_type": "dedupe_only",
+      "unique_id_column_name": unique_id,
+      "comparisons": comparisons,
+      "em_convergence": 0.01,
+      "blocking_rules_to_generate_predictions": blocking_rules_to_generate_predictions
+    }
+    
+    # Train linker model
+    linker = SparkLinker(data, spark=self.spark, database=self.schema, catalog=self.catalog)
+    
+    linker.initialise_settings(settings)
+    linker.estimate_probability_two_random_records_match(deterministic_rules, recall=0.8)
+    linker.estimate_u_using_random_sampling(target_rows=1e7)
+    for training_rule in training_rules:
+      linker.estimate_parameters_using_expectation_maximisation(training_rule)
+
+    # Make predictions
+    predictions = linker.predict()
+
+    return linker, predictions
+  
+  
+  
+  def evaluate_linker(
+    self,
+    data:pyspark.sql.DataFrame,
+    predictions:splink.splink_dataframe.SplinkDataFrame,
+    threshold:float,
+    attribute_columns:list,
+    unique_id:str,
+    linker:splink.spark.spark_linker.SparkLinker
+    ) -> dict:
+    """
+    Method to evaluate predictions made by a trained linker model.
+
+    Returns a dictionary of evaluation metrics.
+
+    :param data: Spark Dataframe containing the original dataset (required to establish ground truth labels)
+    :param df_predictions: Spark DataFrame containing pair-wise predictions made by a linker model
+    :param threshold: float indicating the probability threshold above which a pair is considered a match
+    :param attribute_columns: list of strings with valid column names to compare the entropies of
+    :param linker: isntance of splink.SparkLinker, used for deduplication (clustering)
+    
+    """
+
+    
+    # Deduplicate data
+    deduped = self.deduplicate_records(data, predictions, linker, attribute_columns, unique_id, threshold)
+    
+    # Calculate mean change in entropy
+    mean_entropy_change = self.calculate_entropy_delta(data, deduped, attribute_columns)
+    
+    evals = {
+      "mean_entropy_change": mean_entropy_change
+    }
+
+    return evals
+  
+  
+  
+  def train_and_evaluate_linker(
+    self,
+    data:pyspark.sql.DataFrame,
+    space:dict,
+    attribute_columns:list,
+    unique_id:str,
+    deterministic_columns:list=None,
+    training_columns:list=None,
+    threshold:float=0.9
+    ) -> typing.Tuple[
+      splink.spark.spark_linker.SparkLinker,
+      splink.splink_dataframe.SplinkDataFrame,
+      dict,
+      dict,
+      str
+    ]:
+    """
+    Method to train and evaluate a linker model in one go
+
+    Returns a tuple of trained linker, predictions, metrics, parameters and mlflow run_id
+
+    :param data: Spark DataFrame containing the data to be de-duplicated
+    :param space: dictionary generated by hyperopt sampling
+    :param attribute_columns: list of strings containing all attribute columns
+    :param unique_id: string with the name of the unique ID column
+    :param deterministic columns: list of strings containint columns to block on - if None, they will be generated automatically/randomly
+    :param training_columns: list of strings containing training columns - if None, they will be generated automatically/randomly
+    :param threshold: float indicating the probability threshold above which a pair is considered a match
+    
+    """
+
+    #set start time for measuring training duration
+    start = datetime.now()
+    
+    # Train model
+    linker, predictions = self.train_linker(data, space, attribute_columns, unique_id, deterministic_columns, training_columns)
+
+    end = datetime.now()
+
+    duration = (end-start).seconds
+    
+    # Evaluate model
+    evals = self.evaluate_linker(data, predictions, threshold, attribute_columns, unique_id, linker)
+
+    with mlflow.start_run() as run:
+      splink_mlflow.log_splink_model_to_mlflow(linker, "linker")
+      mlflow.log_metrics(evals)
+      mlflow.log_metric("training_duration", duration)
+      params = space.copy()
+      params["deterministic_columns"] = self.deterministic_columns
+      params["training_columns"] = self.training_columns
+      mlflow.log_dict(params, "model_parameters.json")
+      
+    
+    run_id = run.info.run_id
+    return linker, predictions, evals, params, run_id
+
+  
+  
+  def auto_link(
+    self,
+    data:pyspark.sql.DataFrame,
+    attribute_columns:list,
+    unique_id:str,
+    comparison_size_limit:int,
+    max_evals:int,
+    cleaning="all",
+    deterministic_columns:list=None,
+    training_columns:list=None,
+    threshold:float=0.9
+  ) -> None:
+    """
+    Method to run a series of hyperopt trials.
+    
+
+    :param data: Spark DataFrame containing the data to be de-duplicated
+    :param space: dictionary generated by hyperopt sampling
+    :param attribute_columns: list of strings containing all attribute columns
+    :param unique_id: string with the name of the unique ID column
+    :param comparison_size_limit: int denoting maximum size of pairs allowed after blocking
+    :param max_evals: int denoting max number of hyperopt trials to run
+    :param cleaning: string ("all" or "none") or dictionary with keys as column names and values as list of strings for methods (accepted are "lower" and "alphanumeric_only")
+    :param deterministic columns: list of strings containint columns to block on - if None, they will be generated automatically/randomly
+    :param training_columns: list of strings containing training columns - if None, they will be generated automatically/randomly
+    :param threshold: float indicating the probability threshold above which a pair is considered a match
+
+    
+    """
+    
+    # extract spark from input data
+
+    self.spark = data.sparkSession if not self.spark else self.spark
+    self.catalog = self.catalog if self.catalog else self.spark.catalog.currentCatalog()
+    self.schema = self.schema   if self.schema else self.spark.catalog.currentDatabase()
+    
+    self.username = self.spark.sql('select current_user() as user').collect()[0]['user']
+    self.experiment_name = self.experiment_name if self.experiment_name else f"/Users/{self.username}/Databricks Autolinker {str(datetime.now())}"
+    self.best_params = None
+    
+    mlflow.set_experiment(self.experiment_name)
+
+    # Count rows in data - doing this here so we only do it once
+    self.data_rowcount = data.count()
+    
+    # Clean the data
+    data = self._clean_columns(data, attribute_columns, cleaning)
+    
+    # define objective function
+    def tune_model(space):
+      linker, predictions, evals, params, run_id = self.train_and_evaluate_linker(
+        data,
+        space,
+        attribute_columns,
+        unique_id=unique_id,
+        deterministic_columns=self.deterministic_columns,
+        training_columns=self.training_columns,
+        threshold=threshold
+      )
+      
+      loss = evals["mean_entropy_change"]
+      
+      result = {'loss': loss, 'status': STATUS_OK, 'run_id': run_id}
+      for k, v in evals.items():
+        result.update({k:v})
+
+      return result
+    
+    # initialise trials and create hyperopt space
+    self.trials = Trials()
+    space = self._create_hyperopt_space(data, attribute_columns, comparison_size_limit, unique_id)
+    
+      # run hyperopt trials
+    self.best = fmin(
+        fn=tune_model,
+        space=space,
+        algo=tpe.suggest,
+        max_evals=max_evals,
+        trials=self.trials
+      )
+    
+    best_param_for_rt = self._convert_hyperopt_to_splink()
+    
+    
+    self.best_run_id = self.trials.best_trial["result"]["run_id"]
+    self.best_linker, self.best_predictions_df = self.train_linker(data, best_param_for_rt, attribute_columns, unique_id, self.deterministic_columns, self.training_columns)
+        
+    
+    # return succes text
+    success_text = f"""
+    ======================================================================================
+                                      AutoLinking completed.
+    ======================================================================================
+    Deterministic rules used
+    ------------------------
+    {', '.join(self._generate_rules(self.deterministic_columns))}
+    ------------------------
+    Training rules used
+    ------------------------
+    {', '.join(self._generate_rules(self.training_columns))}
+    ------------------------
+    Best parameters
+    ------------------------
+    {self.best_params}
+    ------------------------
+    Best Metric: {self.best_metric}
+    ------------------------
+    You can now access the best linker model and predictions and use Splink's built-in functionality, e.g.:
+    
+    >>> linker = arc.best_linker
+    >>> predictions = arc.best_predictions()
+    >>> ...
+    """
+    print(success_text)
+    
+    return None
+  
+  def get_best_splink_model(self):
+    """
+    Get the underlying Splink model from the MLFlow Model Registry. Useful for advanced users wishing to access Splink's
+    internal functionality not exposed as part of arc Autolinker.
+    Returns
+    -------
+    A Splink model object.
+    """
+    loaded_model = mlflow.pyfunc.load_model(f"runs:/{self.best_run_id}/linker")
+    unwrapped_model = loaded_model.unwrap_python_model()
+    return unwrapped_model
+  
+  def best_predictions(self):
+    """
+    Get the predictions from the best linker model trained.
+    Returns
+    -------
+    A spark dataframe of the pairwise predictions made by the autolinker model.
+    """
+    return self.best_predictions_df.as_spark_dataframe()
+  
+  def _do_clustering(self):
+    #factor out the clustering call as used multiple times
+    clusters = self.best_linker.cluster_pairwise_predictions_at_threshold(self.best_predictions_df, self.cluster_threshold)
+    self.clusters = clusters
+  
+  def best_clusters_at_threshold(
+    self,
+    threshold:float=0.8
+    ) -> pyspark.sql.DataFrame:
+    """
+    Convert the pairwise predictions to clusters using the connected components algorithm.
+
+    Parameters
+    :param threshold: default value=0.8.An optional parameter controlling the threshold at which records will be connected. Set it higher to produce clusters with greater precision, lower to produce clusters with greater recall.
+
+    Returns a spark dataframe of the clustered input data with a new column cluster_id prepended.
+
+    """
+    # if self.clusters is empty run clustering and log the threshold
+    if not self.clusters:
+      self.cluster_threshold=threshold
+      self._do_clustering()
+    # enter this clause if rerunning clustering. 
+    else:
+      # if the provided threshold is the same as last time, no need to recalculate
+      if threshold == self.cluster_threshold:
+        pass
+      else:
+        self.cluster_threshold=threshold
+        self._do_clustering()
+    return self.clusters.as_spark_dataframe()
+        
+  
+  def cluster_viewer(self):
+    """
+    Produce an interactive dashboard for visualising clusters. It provides examples of clusters of different sizes.
+    The shape and size of clusters can be indicative of problems with record linkage, so it provides a tool to help you
+    find potential false positive and negative links.
+    See this video for an indepth explanation of interpreting this dashboard: https://www.youtube.com/watch?v=DNvCMqjipis
+
+    Writes a HTML file to DBFS at "/dbfs/Users/{username}/scv.html"
+
+    """
+    # do clustering if not already done and no clusters provided
+    if not self.clusters:
+      raise ValueError("Pairs have not yet been clustered. Please run best_clusters_at_threshold() with an optional"
+                       "threshold argument to generate clusters. ")
+      
+    path=f"/Users/{self.username}/clusters.html"
+    self.best_linker.cluster_studio_dashboard(self.best_predictions_df, self.clusters, path, sampling_method="by_cluster_size", overwrite=True)
+    
+    # splink automatically prepends /dbfs to the output path
+    with open("/dbfs" + path, "r") as f:
+        html2=f.read()
+    
+    displayHTML(html2)
+    
+    
+  def comparison_viewer(self) -> None:
+    """
+    Produce an interactive dashboard for querying comparison details.
+    See this video for an indepth explanation of interpreting this dashboard: https://www.youtube.com/watch?v=DNvCMqjipis
+
+    Writes a HTML file to DBFS at "/dbfs/Users/{username}/scv.html"
+
+    Returns None.
+    """
+    path=f"/dbfs/Users/{self.username}/scv.html"
+
+    self.best_linker.comparison_viewer_dashboard(self.best_predictions_df, path, overwrite=True)
+
+    with open("/dbfs" + path, "r") as f:
+        html=f.read()
+
+    displayHTML(html)
+ 
+
+  def match_weights_chart(self) -> None:
+    """
+    Get the
+    Returns
+
+    """
+    return self.best_linker.match_weights_chart()
+    
+    
diff --git a/arc/autolinker/splink_mlflow.py b/arc/autolinker/splink_mlflow.py
new file mode 100644
index 0000000..bfea626
--- /dev/null
+++ b/arc/autolinker/splink_mlflow.py
@@ -0,0 +1,327 @@
+import shutil
+import mlflow
+import os
+import json
+
+from splink.charts import save_offline_chart
+import splink
+
+import mlflow
+from mlflow.models.signature import infer_signature
+from pyspark.sql import SparkSession
+from splink.spark.spark_linker import SparkLinker
+
+# ======================================================================
+# ======================================================================
+def _log_comparison_details(comparison:dict) -> None:
+    """
+    Traverse a single comparison item from the splink settings JSON and log the sql condition the comparison was made
+    under and the m and u values
+
+    :param comparison: dictionary element from a Splink settings object
+
+    """
+    output_column_name = comparison['output_column_name']
+
+    comparison_levels = comparison['comparison_levels']
+    for _ in comparison_levels:
+        sql_condition = _.get('sql_condition')
+        m_probability = _.get('m_probability')
+        u_probability = _.get('u_probability')
+        log_dict = {
+            f"output column {output_column_name} compared through condition {sql_condition} m_probability": m_probability,
+            f"output column {output_column_name} compared through condition {sql_condition} u_probability": u_probability
+
+        }
+        if m_probability: # this means we don't log the null comparison as it doesnt have a m and u value
+            mlflow.log_params(log_dict)
+
+
+# ======================================================================
+# ======================================================================
+def _log_comparisons(splink_model_json:dict) -> None:
+    """
+    Traverse the comparisons part of the splink settings and extract the learned values.
+    This allows you to easily compare values across different training conditions.
+
+    :param splink_model_json: the settings json from a splink model.
+
+    """
+    comparisons = splink_model_json['comparisons']
+    for _ in comparisons:
+        _log_comparison_details(_)
+
+
+# ======================================================================
+# ======================================================================
+def _log_hyperparameters(splink_model_json:dict) -> None:
+    """
+    Simple method for extracting parameters from the splink settings and logging as parameters in MLFlow
+    
+    :param splink_model_json: the settings json from a splink model.
+
+    """
+    hyper_param_keys = ['link_type', 'probability_two_random_records_match', 'sql_dialect', 'unique_id_column_name',
+                        'em_convergence', 'retain_matching_columns', 'blocking_rules_to_generate_predictions']
+    for key in hyper_param_keys:
+        mlflow.log_param(key, splink_model_json[key])
+
+
+# ======================================================================
+# ======================================================================
+def _log_splink_model_json(splink_model_json:dict) -> None:
+    """
+    Log the splink settings json in its entirery under the name "linker.json"
+
+    :param splink_model_json: the settings json from a splink model.
+
+    """
+    mlflow.log_dict(splink_model_json, "linker.json")
+
+
+# ======================================================================
+# ======================================================================
+class splinkSparkMLFlowWrapper(mlflow.pyfunc.PythonModel):
+    """
+    Wrapper class necessary for storing and retrieving the splink model JSON from the MLFlow model repository.
+    """
+
+    def load_context(self, context:mlflow.pyfunc.PythonModelContext) -> None:
+        """
+        Loads linker settings json into context
+
+        :param context: Python Model context instance
+
+        """
+        # this simply stores the json with the MLFLow model.
+        self.linker_json = context.artifacts['linker_json_path']
+
+
+
+# ======================================================================
+# ======================================================================
+def _save_splink_model_to_mlflow(
+    linker:splink.spark.spark_linker.SparkLinker,
+    model_name:str) -> None:
+    """
+    Simple method for logging a splink model
+    
+    :param linker: Splink model object
+    :param model_name: name to save model under
+
+    """
+    path = "linker.json"
+    if os.path.isfile(path):
+        os.remove(path)
+    linker.save_settings_to_json(path)
+    artifacts = {"linker_json_path": path}
+    mlflow.pyfunc.log_model(model_name, python_model=splinkSparkMLFlowWrapper(), artifacts=artifacts)
+    os.remove(path)
+
+
+# ======================================================================
+# ======================================================================
+def log_splink_model_to_mlflow(
+    linker:splink.spark.spark_linker.SparkLinker,
+    model_name:str,
+    log_parameters_charts:bool=True,
+    log_profiling_charts:bool=False,
+    params:dict=None,
+    metrics:dict=None,
+    artifacts:dict=None
+    ) -> None:
+    """
+    Comprehensive logging of Splink attributes, parameters, charts and model JSON to MLFlow to provide easy
+    reproducability and tracking during model development and deployment. This will create a new run under which
+    to log everything. Additional parameters and metrics can be logged through the params and metrics arguments.
+
+    For details on MLFlow please visit https://mlflow.org/.
+
+    This will log
+    - All the m and u values trained along with their training condition
+    - All the hyperparameters set during training (blocking rules, link type, sql_dialect etc)
+
+    :param run: an mlflow.start_run() instance
+    :param linker: a trained Splink linkage model
+    :param model_name: str, the name to log the model under
+    :param log_parameters_charts: boolean, whether to log parameter charts or not. Default to True as this is a quick operation.
+    :param log_profiling_charts: boolean, whether to log data profiling charts or not. Default to False as this requires running Spark jobs and can take time.
+    :param params: dict[str: str]. Dictionary is of param_name :param_value. Optional argument for logging arbitrary parameters
+    :param metrics: dict[str: double or int]. Dictionary is of metric_name :metric_value. Optional argument for logging arbitrary metrics
+    :param artifacts: dict[str: str]. Dictionary is of artifact_name :artifact_filepath. Optional argument for logging arbitrary artifacts (NB - artifacts must be written to disk before logging).
+
+
+    """
+
+    splink_model_json = linker._settings_obj.as_dict()
+
+   
+    _log_splink_model_json(splink_model_json)
+    _log_hyperparameters(splink_model_json)
+    _log_comparisons(splink_model_json)
+    #_save_splink_model_to_mlflow(linker, model_name)
+    model = SplinkMLFlowWrapper()
+    model.set_spark_linker(linker)
+    mlflow.pyfunc.log_model(model_name, python_model=model)
+    if log_profiling_charts or log_parameters_charts:
+        _log_linker_charts(linker, log_parameters_charts, log_profiling_charts)
+    if params:
+        mlflow.log_params(params)
+    if metrics:
+        mlflow.log_metrics(metrics)
+    if artifacts:
+        mlflow.log_artifacts(artifacts)
+
+
+
+# ======================================================================
+# ======================================================================
+def _log_linker_charts(
+    linker:splink.spark.spark_linker.SparkLinker,
+    log_parameters_charts:bool,
+    log_profiling_charts:bool
+    ) -> None:
+    '''
+    Log all the non-data related charts to MLFlow
+
+    :param linker: a Splink linker object
+    :param log_parameters_charts: boolean, whether to log parameter charts or not
+    :param log_profiling_charts: boolean, whether to log data profiling charts or not
+
+
+    '''
+
+    charts_dict = {}
+    if log_parameters_charts:
+        weights_chart = linker.match_weights_chart()
+        mu_chart = linker.m_u_parameters_chart()
+        compare_chart = linker.parameter_estimate_comparisons_chart()
+        charts_dict["weights_chart"] = weights_chart
+        charts_dict["mu_chart"] = mu_chart
+        charts_dict["compare_chart"] = compare_chart
+
+    if log_profiling_charts:
+        missingness_chart = linker.missingness_chart()
+        # completeness = linker.completeness_chart() # Note to Splink team - this method behaves differently to the others
+        unlinkables = linker.unlinkables_chart()
+        blocking_rules_chart = linker.cumulative_num_comparisons_from_blocking_rules_chart()
+        charts_dict["missingness_chart"] = missingness_chart
+       # charts_dict["completeness"] = completeness
+        charts_dict["unlinkables"] = unlinkables
+        charts_dict["blocking_rules_chart"] = blocking_rules_chart
+
+    for name, chart in charts_dict.items():
+        _log_chart(name, chart)
+
+
+# ======================================================================
+# ======================================================================
+def _log_chart(chart_name:str, chart):
+    '''
+    Save a chart to MLFlow. This writes the chart out temporarily for MLFlow to pick up and save.
+
+    :param chart_name: the name the chart will be given in MLFlow
+    :param chart: chart object from Splink
+
+    '''
+    path = f"{chart_name}.html"
+    if os.path.isfile(path):
+        os.remove(path)
+    save_offline_chart(chart.spec, path)
+    mlflow.log_artifact(path)
+    os.remove(path)
+
+# ======================================================================
+# ======================================================================
+def get_model_json(artifact_uri:str) -> dict:
+    """
+
+    Returns a Splink model JSON from a model tracked with MLFlow. This can be retrieved either from a run or
+    from the MLFlow Model Registry, if and only if the model was logged to MLFlow using the method
+    splink.mlflow.log_splink_model_to_mlflow().
+
+    Extracting the linker's JSON specification requires temporarily writing to /tmp on the local file system. All
+    written files are deleted before the function returns.
+
+    Returns a splink setting json as a python dict.
+
+    :param artifact_uri: URI pointing to the artifacts, such as "runs:/500cf58bee2b40a4a82861cc31a617b1/my_model.pkl", "models:/my_model/Production", or "s3://my_bucket/my/file.txt".
+
+    """
+    temp_file_path = f"/tmp/{artifact_uri.split('/')[1]}/splink_mlflow_artifacts_download"
+    os.makedirs(temp_file_path)
+    mlflow.artifacts.download_artifacts(
+        artifact_uri=artifact_uri
+        , dst_path=temp_file_path
+    )
+
+    with open(f"{temp_file_path}/artifacts/linker.json", "r") as f:
+        linker_json = json.load(f)
+    shutil.rmtree(temp_file_path)
+
+    return linker_json
+
+# ======================================================================
+# ======================================================================
+class SplinkMLFlowWrapper(mlflow.pyfunc.PythonModel):
+    """
+    MLFlow compatibility wrapper for splink linker.
+
+    """
+ 
+    def __init__(self, **kwargs):
+        self.settings = {}
+        self.linker = None
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+            
+    def __getstate__(self):
+        json_result = self.get_settings_object()
+        return json_result
+            
+    def __setstate__(self, json_dict):
+        self.settings = json_dict
+    
+    def load_context(self, context):
+        # this simply stores the json with the MLFLow model.
+        self.context = context
+        
+    def clear_context(self):
+        # model cannot be pickled, so disposing it
+        self.linker.spark = None
+        
+    def set_should_evaluate(self, flag):
+        self.should_evaluate = flag
+        
+    def spark_linker(self, data):
+        spark = SparkSession.builder.getOrCreate()
+        self.linker = SparkLinker(data, spark=spark)
+        spark = None
+        return self.linker
+        
+    def set_spark_linker(self, linker: SparkLinker):
+        self.linker = linker
+    
+    def set_settings(self, settings):
+        self.settings = settings
+        if self.linker:
+            self.linker.initialise_settings(self.settings)
+        
+    def get_linker(self):
+        return self.linker
+    
+    def get_settings(self):
+        return self.settings
+   
+    def get_settings_object(self):
+        return self.get_linker()._settings_obj.as_dict()
+    
+ 
+    def predict(self, context, X):
+        """
+        Predict labels on provided data
+        """
+        self.linker = self.spark_linker(X)
+        self.linker.initialise_settings(self.settings)
+        result = self.linker.predict()
+        return result
diff --git a/arc/notebooks/autolinker_notebook.py b/arc/notebooks/autolinker_notebook.py
new file mode 100644
index 0000000..1dd8ede
--- /dev/null
+++ b/arc/notebooks/autolinker_notebook.py
@@ -0,0 +1,837 @@
+# Databricks notebook source
+# MAGIC %pip install --quiet splink mlflow hyperopt
+
+# COMMAND ----------
+
+# disable AQE
+# Repartition data to 400 & cache
+spark.conf.set("spark.databricks.optimizer.adaptive.enabled", "false")
+
+# COMMAND ----------
+
+from pyspark.sql import functions as F
+from pyspark.sql.types import *
+
+from splink.spark.spark_linker import SparkLinker
+from splink.databricks.enable_splink import enable_splink
+import splink.spark.spark_comparison_library as cl
+
+import hyperopt
+from hyperopt import fmin, tpe, hp, SparkTrials, STATUS_OK, Trials
+from hyperopt.pyll import scope
+
+from sklearn.mixture import GaussianMixture
+
+from pyspark.ml.evaluation import BinaryClassificationEvaluator
+
+import pandas as pd
+import numpy as np
+import itertools
+import math
+import random
+from datetime import datetime
+from pprint import pprint
+
+import mlflow
+
+# COMMAND ----------
+
+class AutoLinker:
+  """
+  Class to create object arc for automated data linking
+  """
+  
+  def __init__(self, catalog, schema, experiment_name, training_columns=None, deterministic_columns=None):
+
+    self.catalog = catalog
+    self.schema = schema
+    self.experiment_name = experiment_name
+    self.training_columns = training_columns
+    self.deterministic_columns = deterministic_columns
+    self.df_true_positives = None
+    
+    mlflow.set_experiment(experiment_name)
+
+  def __str__(self):
+    return f"AutoLinker instance working in {self.catalog}.{self.schema} and MLflow experiment {experiment_name}"
+  
+  
+  def _calculate_column_entropy(self, data, column):
+    """
+    Method to calculate column entropy given a dataset, based on:
+    Entropy = -SUM(P*ln(P))
+    where P is the normalised count of the occurrence of each unique value in the column
+    Parameters
+    : data : Spark Dataframe
+    : column : name of column to calculate
+    Returns
+      - Float
+    """
+    
+    # calculate normalised value count per unique value in colunn
+    vc = data.groupBy(column).count().withColumn("norm_count", F.col("count")/self.data_rowcount)
+
+    
+    # Calculate P*ln(P) per row
+    vc = vc.withColumn("entropy", F.col("norm_count")*F.log(F.col("norm_count")))
+    
+    # Entropy = -SUM(P*ln(P))
+    entropy = -vc.select(F.sum(F.col("entropy")).alias("sum_entropy")).collect()[0].sum_entropy
+
+    
+    return entropy
+  
+    
+  def _generate_rules(self, columns):
+    """
+    Method to create a list of Splink-compatible SQL statements from a list of candidate columns for rules.
+    Can be used for simple blocking rules (no AND statement), deterministic rules or for training rules
+    Parameters
+    : columns : list of str containing valid columns in the data set
+    Returns
+      - List of strings of Splink-compatible SQL statements
+    """
+
+    rules = []
+    for column in columns:
+      if "&" in column:
+        subcols = column.split("&")
+        subcols = [subcol.strip() for subcol in subcols]
+
+        rule = " AND ".join([f"l.{subcol} = r.{subcol}" for subcol in subcols])
+      else:
+        rule = f"l.{column} = r.{column}"
+
+      rules.append(rule)
+
+    return rules
+  
+  
+  def _generate_candidate_blocking_rules(self, data, attribute_columns, comparison_size_limit, unique_id, max_columns_per_rule=2):
+    """
+    Method to automatically generate a list of lists of blocking rules to test, given a user-defined limit for
+    pair-wise comparison size.
+    Parameters
+    : data : Spark DataFrame containing all records to de-duplicate
+    : attribute_columns : list of strings with valid column names of data, containing all possible columns to block on
+    : comparison_size_limit : integer denoting the maximum number of pairs we want to compare - to limit hardware issues
+    : max_columns_per_rule : integer denoting the maximum number of column comparisons in a single rule to try
+    Returns
+      - nested list of lists with Splink-compatible blocking rule queries
+    """
+
+    # initialise empty list to store all combinations in
+    blocking_combinations = []
+
+    # create list of all possible rules
+    for n in range(1, max_columns_per_rule+1):
+      combs = list(itertools.combinations(attribute_columns, n))
+      blocking_combinations.extend(combs)
+
+    # Convert them to Splink-compatible rules
+    blocking_rules = self._generate_rules(["&".join(r) for r in blocking_combinations])
+    
+    # Estimate number of pairs per blocking rule by grouping on blocking rule columns and getting the sum of squares of counts
+    comp_size_dict = dict()
+
+    for comb, rule in zip(blocking_combinations, blocking_rules):
+      num_pairs = data.groupBy(list(comb)).count().select(F.sum(F.col("count")*F.col("count"))).collect()[0]["sum((count * count))"]
+      comp_size_dict.update({rule: num_pairs})
+    
+    # loop through all combinations of combinations and save those which remain under the limit (this is a bit brute force)
+    accepted_rules = []
+    for r in range(1, len(blocking_rules)+1):
+      for c in itertools.combinations(blocking_rules, r):
+        if sum([v for k, v in comp_size_dict.items() if k in c])<=comparison_size_limit:
+          accepted_rules.append(list(c))
+
+
+    if len(accepted_rules)<1:
+      raise ValueError("No blocking rules meet the comparison size limit.")
+          
+    return accepted_rules
+
+  
+  def _create_hyperopt_space(self, data, attribute_columns, comparison_size_limit, unique_id, max_columns_per_rule=2):
+    """
+    Method to create hyperopt space for comparison and blocking rule hyperparameters from list of columns.
+    Takes a given (or generated) list of columns and produces a dictionary that can be converted to a comparison list later on,
+    with function names and threhsolds, and a list (of lists) of blocking rules to test.
+    Parameters
+    : data : Spark DataFrame containing all records to de-duplicate
+    : attribute_columns : list of strings with valid column names of data, containing all possible columns to block on or compare
+    : comparison_size_limit : integer denoting the maximum number of pairs we want to compare - to limit hardware issues
+    : max_columns_per_rule : integer denoting the maximum number of column comparisons in a single rule to try
+    Returns
+      - hyperopt.pyll graph for hyperopt parameter search
+    """
+
+    # Generate candidate blocking rules
+    self.blocking_rules = self._generate_candidate_blocking_rules(
+      data=data,
+      attribute_columns=attribute_columns,
+      comparison_size_limit=comparison_size_limit,
+      unique_id=unique_id,
+      max_columns_per_rule=max_columns_per_rule
+    )
+
+    # Create comparison dictionary using Hyperopt sampling
+    space = dict()
+    comparison_space = dict()
+
+    # Create hyperopt space for comparisons
+    for column in attribute_columns:
+      comparison_space.update(
+        {
+          column: {
+            "distance_function": hp.choice(
+              f"{column}|distance_function", [
+                {"distance_function": "levenshtein", "threshold": hp.quniform(f"{column}|levenshtein|threshold", 1, 5, q=1)},
+                {"distance_function": "jaccard", "threshold": hp.uniform(f"{column}|jaccard|threshold", 0.7, 0.99)},
+                {"distance_function": "jaro_winkler", "threshold": hp.uniform(f"{column}|jaro_winkler|threshold", 0.7, 0.99)}
+              ]
+            )
+          }
+        }
+      )
+
+    # Create hyperopt space for blocking rules
+    space.update({
+      "blocking_rules": hp.choice("blocking_rules", self.blocking_rules),
+      "comparisons": comparison_space
+    })
+
+    return space
+  
+  def _drop_intermediate_tables(self):
+    # Drop intermediate tables in schema for consecutive runs
+    tables_in_schema = spark.sql(f"show tables from {self.catalog}.{self.schema} like '*__splink__*'").collect()
+    for table in tables_in_schema:
+      try:
+        spark.sql(f"drop table marcell_splink.marcell_autosplink.{table.tableName}") 
+      except:
+        spark.sql(f"drop table {table.tableName}")
+        
+        
+  def _clean_columns(self, data, attribute_columns, cleaning):
+    """
+    Method to clean string columns (turn them to lower case and remove non-alphanumeric characters)
+    in order to help with better (and quicker) string-distance calculations. If cleaning is "all"
+    (as is by default), it will automatically clean as much as it can.  If cleaning is "none", it will do nothing.
+    cleaning can also be a dictionary with keys as column names and values as lists of method strings.
+    The currently available cleaning methods are turning to lower case and removing non-alphanumeric characters.
+    Parameters
+    : data : Spark DataFrame containing the data to be cleaned
+    : attribute_columns : list of strings containing all attribute columns
+    : cleaning : string or dicitonary with keys as column names and values as list of valid cleaning method names.
+    Returns
+      - Spark dataframe
+    """
+    
+    # if cleaning is set to "none", don't do anything to the data
+    if cleaning=="none":
+      return data
+    
+    # if it's set to "all", turn it into a dictionary covering all possible cases
+    if cleaning=="all":
+      cleaning = {col: ["lower", "alphanumeric_only"] for col in attribute_columns}
+      
+    for col, methods in cleaning.items():
+      # if column is not a string, skip it
+      if not data.schema[col].dataType.typeName == "string":
+        continue
+        
+      for method in methods:
+        if method=="alphanumeric_only":
+          # replace column and only keep alphanumeric and space characters
+          data = data.withColumn(col, F.regexp_replace(F.col(col), r"[^A-Za-z0-9 ]+", ""))
+        
+        elif method=="lower":
+          data = data.withColumn(col, F.lower(F.col(col)))
+          
+    return data
+  
+  
+  def _create_comparison_list(self, space):
+    """
+    Method to convert comparisons dictionary generated by hyperopt to a Splink-compatible list of
+    spark_comparison_library functions with the generated columns and thresholds.
+    Parameters
+    : comparison_dict : nested dicitonary generated by create_comparison_dict
+    Returns
+      - List of spark_comparison_library method instances
+    """
+
+    # Initialise empty list to populate
+    comparison_list = []
+
+    # Get comparison dict
+    comparison_dict = space["comparisons"]
+
+    # Loop through comparison dictionary and generate Splink-compatible comparison list
+    for column in comparison_dict.keys():
+      distance_function = comparison_dict[column]["distance_function"]["distance_function"]
+      threshold = comparison_dict[column]["distance_function"]["threshold"]
+
+      if distance_function=="levenshtein":
+        comparison_list.append(cl.levenshtein_at_thresholds(column, threshold))
+      elif distance_function=="jaccard":
+        comparison_list.append(cl.jaccard_at_thresholds(column, threshold))
+      elif distance_function=="jaro_winkler":
+        comparison_list.append(cl.jaro_winkler_at_thresholds(column, threshold))
+      else:
+        raise ValueError(f"Unknown distance function {distance_function} passed.")
+
+    return comparison_list
+  
+
+  def deduplicate_records(self, data, predictions, linker, attribute_columns, unique_id, threshold=0.9):
+    """
+    Method to deduplicate the original dataset given the predictions dataframe, its linker
+    object and an optional threshold. The clusters are grouped on and duplicates are removed (arbitrarily).
+    These are then replaced in the original dataset.
+    Parameters
+    : data : Spark Dataframe of the original data set
+    : predictions : Splink Dataframe that's a result of a linker.predict() call
+    : linker : the linker object (splink.SparkLinker instance), needed to cluster
+    : threshold : Float denoting the probability threshold above which a pair is considered a match
+    Returns
+      - Spark Dataframe with deduplicated records
+    """
+    
+    # Cluster records that are matched above threshold
+    clusters = linker.cluster_pairwise_predictions_at_threshold(predictions, threshold_match_probability=threshold)
+    df_predictions = clusters.as_spark_dataframe().select(unique_id, "cluster_id")
+
+    # Get max cluster id to use as lower bound of monotonically increasing id
+    max_cluster_id = df_predictions.select(F.max("cluster_id")).collect()[0]["max(cluster_id)"] + 1
+    
+    # Assign monotonically increasing id greater than max cluster_id
+    data = data.withColumn("tmp_cluster_id", F.monotonically_increasing_id()+max_cluster_id)
+    
+    # Left join predictions onto input data
+    data = data.join(df_predictions, on=unique_id, how="left")
+    
+    # Coalesce
+    data = data.withColumn("cluster_id", F.coalesce(F.col("cluster_id"), F.col("tmp_cluster_id")))
+    
+    # Group by cluster ID
+    df_deduped = data.groupBy("cluster_id").agg(
+      *(F.first(i).alias(i) for i in attribute_columns)
+    )
+
+    return df_deduped
+
+
+  def calculate_entropy_delta(self, data, deduped, columns):
+    """
+    Method to calculate the change in entropy of a set of columns between two
+    datasets, used to calculate the information gain after deduping a set of records.
+    Parameters
+    : data : Spark DataFrame, the original (undeduplicated) data set of records
+    : deduped : Spark DataFrame, the new (deduplicated) data set of records
+    : columns : list of strings with valid column names to compare
+    Returns
+      - Float (average change in column entropy)
+    """
+    
+    # Initialise empty list to populate with entropy deltas
+    entropy_deltas = list()
+    
+    # For each column in given columns, calculate entropy in original dataset
+    # and the deduplicated dataset. The delta is the difference between new and old.
+    for column in columns:
+      original_entropy = self._calculate_column_entropy(data, column)
+      new_entropy = self._calculate_column_entropy(deduped, column)
+      entropy_deltas.append(new_entropy-original_entropy)
+
+
+    return np.mean(entropy_deltas)
+  
+  
+  
+  def calculate_empirical_score(self, data, predictions, unique_id, threshold):
+    """
+    Method to calculate precision, recall and F1 score based on ground truth labels.
+    Assumes the data has a column with empirical IDs to assess whether two records belong
+    to the same real life entity. Will check if this has already been calculated. If not,
+    it will create it for the first (and only) time.
+    NB: this includes a bit of hard coding, but it won't make it to production anyway because
+    we won't have ground truth.
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : predictions : Splink DataFrame with predicted pairs
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    Returns
+      - 3-tuple of floats for precision, recall and F1 score
+    """
+    if not self.df_true_positives:
+      # filter original data to only contain rows where recid id appears more than once (known dupes)
+      data_recid_groupby = data.groupBy("recid").count().filter("count>1").withColumnRenamed("recid", "recid_")
+      data_tp = data.join(data_recid_groupby, data.recid==data_recid_groupby.recid_, how="inner").drop("recid_")
+
+      data_l = data_tp.withColumnRenamed(unique_id, f"{unique_id}_l").withColumnRenamed("recid", "recid_l").select(f"{unique_id}_l", "recid_l")
+      data_r = data_tp.withColumnRenamed(unique_id, f"{unique_id}_r").withColumnRenamed("recid", "recid_r").select(f"{unique_id}_r", "recid_r")
+
+      dt = data_l.join(data_r, data_l[f"{unique_id}_l"]!=data_r[f"{unique_id}_r"], how="inner")
+      # create boolean col for y_true
+      df_true = dt.withColumn("match", F.when(F.col("recid_l")==F.col("recid_r"), 1).otherwise(0))
+      # only keep matches
+      df_true = df_true.filter("match=1")
+      # assign as attribute to avoid re-calculation
+      self.df_true_positives = df_true
+      
+    # convert predictions to Spark DataFrame and filter on match prob threshold - table will only contain predicted positives
+    df_pred = predictions.as_spark_dataframe().filter(f"match_probability>={threshold}")
+    
+    # Calculate TP, FP, FN, TN
+    
+    # TP is the inner join of true and predicted pairs
+    tp = df_pred.join(
+      self.df_true_positives,
+      ((df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_r"]) & (df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_r"])) | ((df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_r"]) & (df_pred[f"{unique_id}_r"]==self.df_true_positives[f"{unique_id}_l"])),
+      how="inner"
+    ).count()
+    
+    # FN is the left anti-join of true and predicted
+    fn = self.df_true_positives.join(
+      df_pred,
+      ((df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_l"]) & (df_pred[f"{unique_id}_r"]==self.df_true_positives[f"{unique_id}_r"])) | ((df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_r"]) & (df_pred[f"{unique_id}_r"]==self.df_true_positives[f"{unique_id}_l"])),
+      how="left_anti"
+    ).count()
+    
+    # FP is the left anti-join of predicted and true
+    fp = df_pred.join(
+      self.df_true_positives,
+      ((df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_l"]) & (df_pred[f"{unique_id}_r"]==self.df_true_positives[f"{unique_id}_r"])) | ((df_pred[f"{unique_id}_l"]==self.df_true_positives[f"{unique_id}_r"]) & (df_pred[f"{unique_id}_r"]==self.df_true_positives[f"{unique_id}_l"])),
+      how="left_anti"
+    ).count()
+    
+    # TN is everything else, i.e. N(N-1)-TP-FN-FP
+    N = data.count()
+    tn = (N*(N-1))-tp-fn-fp
+    
+    # Calculate precision, recall and f1
+    if tp+fp>0.0:
+      precision = tp/(tp+fp)
+    else:
+      precision: 0.9
+      
+    if tp+fn>0.0:
+      recall = tp/(tp+fn)
+    else:
+      recall = 0.0
+
+    if (precision+recall)>0:
+      f1 = precision*recall/(precision+recall)
+    else:
+      f1 = 0.0
+      
+    return precision, recall, f1
+  
+  
+  def _randomise_columns(self, attribute_columns):
+    """
+    Method to randomly select (combinations of) columns from attribute columns for EM training.
+    Will try to pick 2 combinations of 2 (i.e AB and BC from ABC), but will default to 2 if there are only 2.
+    Parameters
+    : attribute_columns : list of strings containing all attribute columns
+    Returns
+      - List of lists of strings
+    """
+    # if there are only 2 columns, return them both
+    if len(attribute_columns)<3:
+      return attribute_columns
+    
+    # pick 3 random columns
+    cols = random.sample(attribute_columns, 3)
+    
+    # this creates a list like ["A&B", "B&C", "A&C"] - it's useful to have an AND statement in the training
+    # columns, otherwise EM training tables balloon
+    training_columns = ["&".join([cols[0], cols[1]]), "&".join([cols[1], cols[2]]), "&".join([cols[0], cols[2]])]
+    
+    return training_columns
+  
+
+  def train_linker(self, data, space, attribute_columns, unique_id, deterministic_columns=None, training_columns=None):
+    """
+    Method to train a linker model
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : space : dictionary generated by hyperopt sampling
+    : attribute_columns : list of strings containing all attribute columns
+    : unique_id : string with the name of the unique ID column
+    : deterministic columns : list of strings containint columns to block on
+    : training_columns : list of strings containing training columns
+    Returns
+      - tuple of trained linker object and Splink dataframe with predictions
+    """
+
+    # Drop any intermediate tables for a clean run
+    self._drop_intermediate_tables()
+    
+    # Set up parameters
+
+    # Get blocking rules from hyperopt space
+    blocking_rules_to_generate_predictions = list(space["blocking_rules"])
+    print(blocking_rules_to_generate_predictions)
+    # Create comparison list from hyperopt space
+    comparisons = self._create_comparison_list(space)
+
+    # if deterministic columns are not set, pick 2 at random from attribute columns
+    # and save them as an attribute so that they can remain consistent between runs
+    if not self.deterministic_columns:
+      deterministic_columns = random.sample(attribute_columns, 2)
+      self.deterministic_columns = deterministic_columns
+
+    deterministic_rules = self._generate_rules(deterministic_columns)
+
+    # if deterministic columns are not set, pick 2 at random from attribute columns
+    # and save them as an attribute so that they can remain consistent between runs
+    if not training_columns:
+      training_columns = self._randomise_columns(attribute_columns)
+      self.training_columns = training_columns
+      
+    training_rules = self._generate_rules(training_columns)
+
+    # create settings dict
+    settings = {
+      "retain_intermediate_calculation_columns": True,
+      "retain_matching_columns": True,
+      "link_type": "dedupe_only",
+      "unique_id_column_name": unique_id,
+      "comparisons": comparisons,
+      "em_convergence": 0.01,
+      "blocking_rules_to_generate_predictions": blocking_rules_to_generate_predictions
+    }
+
+    # Train linker model
+    linker = SparkLinker(data, spark=spark, database=self.schema, catalog=self.catalog)
+    linker.initialise_settings(settings)
+    linker.estimate_probability_two_random_records_match(deterministic_rules, recall=0.8)
+    linker.estimate_u_using_random_sampling(target_rows=1e7)
+    for training_rule in training_rules:
+      linker.estimate_parameters_using_expectation_maximisation(training_rule)
+
+    # Make predictions
+    predictions = linker.predict()
+
+
+    return linker, predictions
+  
+  def evaluate_linker(self, data, predictions, threshold, attribute_columns, unique_id, linker):
+    """
+    Method to evaluate predictions made by a trained linker model
+    Parameters
+    : data : Spark Dataframe containing the original dataset (required to establish ground truth labels)
+    : df_predictions : Spark DataFrame containing pair-wise predictions made by a linker model
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    : attribute_columns : list of strings with valid column names to compare the entropies of
+    : linker : isntance of splink.SparkLinker, used for deduplication (clustering)
+    Returns
+      - Dictionary of evaluation metrics TODO: make the metrics controllable via an argument
+    """
+
+    # Calculate empirical scores
+#     precision, recall, f1 = self.calculate_empirical_score(data, predictions, unique_id, threshold)
+    
+    # Deduplicate data
+    deduped = self.deduplicate_records(data, predictions, linker, attribute_columns, unique_id, threshold)
+    
+    # Calculate mean change in entropy
+    mean_entropy_change = self.calculate_entropy_delta(data, deduped, attribute_columns)
+    
+    evals = {
+#       "precision": precision,
+#       "recall": recall,
+#       "f1": f1,
+      "mean_entropy_change": mean_entropy_change
+    }
+
+    return evals
+  
+  
+  def train_and_evaluate_linker(self, data, space, attribute_columns, unique_id, deterministic_columns=None, training_columns=None, threshold=0.9):
+    """
+    Method to train and evaluate a linker model in one go
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : space : dictionary generated by hyperopt sampling
+    : attribute_columns : list of strings containing all attribute columns
+    : unique_id : string with the name of the unique ID column
+    : deterministic columns : list of strings containint columns to block on - if None, they will be generated automatically/randomly
+    : training_columns : list of strings containing training columns - if None, they will be generated automatically/randomly
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    Returns
+      - 6-tuple of trained linker (SparkLinker instance), predictions (Splink DAtaFrame) and metrics (4x float)
+    """
+
+    # Train model
+    linker, predictions = self.train_linker(data, space, attribute_columns, unique_id, deterministic_columns, training_columns)
+
+    # Evaluate model
+    evals = self.evaluate_linker(data, predictions, threshold, attribute_columns, unique_id, linker)
+
+    # Log to MLflow
+    with mlflow.start_run():
+      mlflow.log_metrics(evals)
+      params = space.copy()
+      params["deterministic_columns"] = self.deterministic_columns
+      params["training_columns"] = self.training_columns
+      mlflow.log_dict(params, "model_parameters.json")
+      
+    
+    
+    return linker, predictions, evals
+  
+  
+  
+  def auto_link(self, data, attribute_columns, unique_id, comparison_size_limit, max_evals, cleaning="all", deterministic_columns=None, training_columns=None, threshold=0.9):
+    """
+    Method to run a series of hyperopt trials.
+    Parameters
+    : data : Spark DataFrame containing the data to be de-duplicated
+    : space : dictionary generated by hyperopt sampling
+    : attribute_columns : list of strings containing all attribute columns
+    : unique_id : string with the name of the unique ID column
+    : comparison_size_limit : int denoting maximum size of pairs allowed after blocking
+    : max_evals : int denoting max number of hyperopt trials to run
+    : deterministic columns : list of strings containint columns to block on - if None, they will be generated automatically/randomly
+    : training_columns : list of strings containing training columns - if None, they will be generated automatically/randomly
+    : threshold : float indicating the probability threshold above which a pair is considered a match
+    Returns
+      - None
+    """
+    
+
+    #set start time for measuring duration
+    start = datetime.now()
+
+    # Count rows in data - doing this here so we only do it once
+    self.data_rowcount = data.count()
+    
+    # Clean the data
+    data = self._clean_columns(data, attribute_columns, cleaning)
+    
+    # define objective function
+    def tune_model(space):
+      linker, predictions, evals = self.train_and_evaluate_linker(
+        data,
+        space,
+        attribute_columns,
+        unique_id=unique_id,
+        deterministic_columns=self.deterministic_columns,
+        training_columns=self.training_columns,
+        threshold=threshold
+      )
+      
+      loss = evals["mean_entropy_change"]
+      
+      result = {'loss': loss, 'status': STATUS_OK}
+      for k, v in evals.items():
+        result.update({k:v})
+
+      return result
+    
+    # initialise trials and create hyperopt space
+    self.trials = Trials()
+    space = self._create_hyperopt_space(data, attribute_columns, comparison_size_limit, unique_id)
+
+    # run hyperopt trials
+    self.best = fmin(
+      fn=tune_model,
+      space=space,
+      algo=tpe.suggest,
+      max_evals=max_evals,
+      trials=self.trials
+    )
+    
+    # Set best params, metrics and results
+    t = self.trials.best_trial
+    self.best_metric = -t['result']['loss']
+    params = t['misc']['vals']
+    blocking_rules = t['misc']['vals']['blocking_rules'][0]
+    comparison_params = {}
+
+    for k, v in params.items():
+      if ("threshold" in k) & (len(v)>0):
+        column, function, _ = k.split('|')
+        comparison_params.update({
+          column: {
+            "function": function,
+            "threshold": v[0]
+          }
+        })
+
+    self.best_params = {
+      "blocking_rules": self.blocking_rules[blocking_rules],
+      "comparisons": comparison_params
+    }
+    
+    # TODO: move this into its own method, probably - refactor...
+    # convert best params to a format that can retrain a new linker 
+    # model based on the best params. This is necessary because the underlying tables may have been overwritten
+    best_comparisons = {}
+    for k, v in comparison_params.items():
+      best_comparisons.update({
+        k: {"distance_function": {"distance_function": v['function'], "threshold": v['threshold']}}
+      })
+      
+    best_params_for_rt = {
+      "blocking_rules": self.best_params['blocking_rules'],
+      "comparisons": best_comparisons
+    }
+    
+    self.best_linker, self.best_predictions = self.train_linker(data, best_params_for_rt, attribute_columns, unique_id, self.deterministic_columns, self.training_columns)
+    
+    #end training duration
+    end = datetime.now()
+    training_duration = (end-start).seconds
+    
+    # return succes text
+    success_text = f"""
+    ======================================================================================
+    AutoLinking completed. Carried out {max_evals} trials in {training_duration} seconds.
+    ======================================================================================
+    Deterministic rules used
+    ------------------------
+    {', '.join(self._generate_rules(self.deterministic_columns))}
+    ------------------------
+    Training rules used
+    ------------------------
+    {', '.join(self._generate_rules(self.training_columns))}
+    ------------------------
+    Best parameters
+    ------------------------
+    {self.best_params}
+    ------------------------
+    Best F1: {self.best_metric}
+    ------------------------
+    
+    You can now access the best linker model and predictions and use Splink's built-in functionality, e.g.:
+    
+    >>> linker = arc.best_linker
+    >>> predictions = arc.best_predictions
+    >>> ...
+    """
+    
+    print(success_text)
+    
+    
+    
+    
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # North Carolina Voters dataset
+# MAGIC https://dbs.uni-leipzig.de/research/projects/object_matching/benchmark_datasets_for_entity_resolution
+
+# COMMAND ----------
+
+data = spark.read.table("marcell_autosplink.voters_data_sample")
+
+# COMMAND ----------
+
+autolinker = AutoLinker(
+  #spark=spark,                                                                                               # spark instance
+  schema="marcell_autosplink",                                                                               # schema to write results to
+  experiment_name="/Users/marcell.ferencz@databricks.com/autosplink/evaluate/autosplink_voters_test"  # MLflow experiment location
+)
+
+# COMMAND ----------
+
+autolinker.auto_link(
+  data=data,                                                         # dataset to dedupe
+  attribute_columns=["givenname", "surname", "suburb", "postcode"],  # columns that contain attributes to compare
+  unique_id="uid",                                                   # column name of the unique ID
+  comparison_size_limit=100000,                                      # Maximum number of pairs when blocking applied
+  max_evals=3                                                      # Maximum number of hyperopt trials to run
+)
+
+# COMMAND ----------
+
+linker = autolinker.best_linker
+
+# COMMAND ----------
+
+linker.m_u_parameters_chart()
+
+# COMMAND ----------
+
+predictions = autolinker.best_predictions
+
+predictions.as_spark_dataframe().display()
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # Music albums dataset
+# MAGIC https://dbs.uni-leipzig.de/research/projects/object_matching/benchmark_datasets_for_entity_resolution
+
+# COMMAND ----------
+
+# DBTITLE 1,Sample dataset to reduce to 100k rows
+data = spark.read.table("marcell_autosplink.music_data")
+
+# COMMAND ----------
+
+# DBTITLE 1,Truncate strings to 20 chars max
+data = data \
+  .withColumn("title", F.substring(F.col("title"), 0, 20)) \
+  .withColumn("artist", F.substring(F.col("artist"), 0, 20)) \
+  .withColumn("album", F.substring(F.col("album"), 0, 20))
+
+# COMMAND ----------
+
+data = data.repartition(400).cache()
+
+# COMMAND ----------
+
+autolinker = AutoLinker(
+  #spark=spark,                                                                                               # spark instance
+  schema="marcell_autosplink",                                                                               # schema to write results to
+  experiment_name="/Users/marcell.ferencz@databricks.com/autosplink/evaluate/autosplink_music_experiment_full_3"  # MLflow experiment location
+)
+
+# COMMAND ----------
+
+autolinker.auto_link(
+  data=data,                                                         # dataset to dedupe
+  attribute_columns=["title", "length", "artist", "album", "year", "language"],  # columns that contain attributes to compare
+  unique_id="uid",                                                   # column name of the unique ID
+  comparison_size_limit=100000,                                      # Maximum number of pairs when blocking applied
+  max_evals=100                                                      # Maximum number of hyperopt trials to run
+)
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # Beer data
+# MAGIC https://github.com/ZJU-DAILY/CollaborEM/tree/main/data/ER-Magellan/Structured/BeerAdvo-RateBeer
+
+# COMMAND ----------
+
+data = spark.read.table(f"marcell_autosplink.beer_data")
+
+# COMMAND ----------
+
+autolinker = AutoLinker(
+  #spark=spark,                                                                                               # spark instance
+  schema="marcell_autosplink",                                                                               # schema to write results to
+  experiment_name="/Users/marcell.ferencz@databricks.com/autosplink/evaluate/autosplink_beer_experiment"  # MLflow experiment location
+)
+
+# COMMAND ----------
+
+autolinker.auto_link(
+  data=data,                                                         # dataset to dedupe
+  attribute_columns=["Beer_Name", "Brew_Factory_Name", "Style", "ABV"],  # columns that contain attributes to compare
+  unique_id="id",                                                   # column name of the unique ID
+  comparison_size_limit=100000,                                      # Maximum number of pairs when blocking applied
+  max_evals=100                                                      # Maximum number of hyperopt trials to run
+)
+
+# COMMAND ----------
+
+
diff --git a/arc/notebooks/example.py b/arc/notebooks/example.py
new file mode 100644
index 0000000..f7716e7
--- /dev/null
+++ b/arc/notebooks/example.py
@@ -0,0 +1,76 @@
+# Databricks notebook source
+# MAGIC %pip install --quiet splink # mlflow hyperopt
+
+# COMMAND ----------
+
+from arc.autolinker.autolinker import AutoLinker
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # Read Data
+# MAGIC Data downloaded and sampled from https://dbs.uni-leipzig.de/research/projects/object_matching/benchmark_datasets_for_entity_resolution
+
+# COMMAND ----------
+
+data = spark.read.table("marcell_autosplink.voters_data_sample")
+
+# COMMAND ----------
+
+# this is here to set defaults - Autolinker() will use the spark default catalog and database if not provided as args
+spark.sql("use catalog robert_whiffin_uc")
+spark.sql("use database autosplink")
+
+# COMMAND ----------
+
+data.display()
+
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC # Run Auto-Linking
+
+# COMMAND ----------
+
+autolinker = AutoLinker()
+
+# COMMAND ----------
+
+autolinker.auto_link(
+  data=data,                                                             # dataset to dedupe
+  attribute_columns=["givenname", "surname", "suburb", "postcode"],      # columns that contain attributes to compare
+  unique_id="uid",                                                       # column name of the unique ID
+  comparison_size_limit=200000,                                          # Maximum number of pairs when blocking applied
+  max_evals=10                                                            # Maximum number of hyperopt trials to run
+)
+
+# COMMAND ----------
+
+display(
+  autolinker.best_predictions()
+)
+
+# COMMAND ----------
+
+import pyspark.sql.functions as F
+from pyspark.sql import Window
+
+display(
+  autolinker.best_clusters_at_threshold()# default=0.8
+  .withColumn("size", F.count("*").over(Window.partitionBy("cluster_id")))
+  .orderBy(-F.col("size"))
+)
+
+# COMMAND ----------
+
+autolinker.cluster_viewer()
+
+# COMMAND ----------
+
+autolinker.comparison_viewer()
+
+# COMMAND ----------
+
+autolinker.match_weights_chart()
+
diff --git a/autolinker/.DS_Store b/autolinker/.DS_Store
new file mode 100644
index 0000000..151f39d
Binary files /dev/null and b/autolinker/.DS_Store differ
diff --git a/docs/docs-requirements.txt b/docs/docs-requirements.txt
index 7e676ff..f9715ba 100644
--- a/docs/docs-requirements.txt
+++ b/docs/docs-requirements.txt
@@ -6,4 +6,8 @@ ipython==8.0.1
 sphinxcontrib-fulltoc==1.2.0
 livereload==2.6.3
 autodocsumm==0.2.7
-sphinx-tabs==3.2.0
\ No newline at end of file
+sphinx-tabs==3.2.0
+mlflow>=1.8.0
+hyperopt>=0.2.0
+pyspark>=3.0.0
+splink>=3.0.0
diff --git a/docs/source/conf.py b/docs/source/conf.py
index fcb45e3..20bc5e7 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -13,8 +13,8 @@
 import os
 import sys
 
-sys.path.insert(0, os.path.abspath('../../python'))
 
+sys.path.insert(0, os.path.abspath('../../arc'))
 # -- Project information -----------------------------------------------------
 
 project = 'ARC'
@@ -35,12 +35,15 @@
     "sphinx_tabs.tabs",
     "sphinx.ext.githubpages",
     "sphinx.ext.autosectionlabel",
-    "sphinx.ext.todo"
+    "sphinx.ext.todo",
+    "sphinx.ext.autodoc"
 ]
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['_templates']
 
+autodoc_typehints = "description"
+
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 9595131..151658e 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -15,6 +15,52 @@ Documentation
    :caption: Contents:
 
 
+AutoLinker
+----------
+.. autoclass:: autolinker.autolinker.AutoLinker
+    :members:
+    :undoc-members:
+    :private-members:
+
+
+MLflow utilities
+----------------
+
+Wrapper classes
+~~~~~~~~~~~~~~~
+.. autoclass:: autolinker.splink_mlflow.splinkSparkMLFlowWrapper
+   :members:
+   :undoc-members:
+   :private-members:
+
+.. autoclass:: autolinker.splink_mlflow.SplinkMLFlowWrapper
+   :members:
+   :undoc-members:
+   :private-members:
+
+
+Utility functions
+~~~~~~~~~~~~~~~~~
+.. autofunction:: autolinker.splink_mlflow._log_comparison_details
+
+.. autofunction:: autolinker.splink_mlflow._log_comparisons
+
+.. autofunction:: autolinker.splink_mlflow._log_hyperparameters
+
+.. autofunction:: autolinker.splink_mlflow._log_splink_model_json
+
+.. autofunction:: autolinker.splink_mlflow._save_splink_model_to_mlflow
+
+.. autofunction:: autolinker.splink_mlflow.log_splink_model_to_mlflow
+
+.. autofunction:: autolinker.splink_mlflow._log_linker_charts
+
+.. autofunction:: autolinker.splink_mlflow._log_chart
+
+.. autofunction:: autolinker.splink_mlflow.get_model_json
+
+
+
 Indices and tables
 ==================
 


Test patch:
diff --git a/arc/tests/autolinker_test.py b/arc/tests/autolinker_test.py
new file mode 100644
index 0000000..617e671
--- /dev/null
+++ b/arc/tests/autolinker_test.py
@@ -0,0 +1,336 @@
+import pytest
+from arc.autolinker.autolinker import AutoLinker
+from pyspark.sql import SparkSession
+import splink.spark.spark_comparison_library as cl
+import hyperopt.pyll.stochastic
+
+
+
+def test_calculate_column_entropy():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+  # mock data
+  df_mock = spark.createDataFrame([
+    (0, "a"),
+    (1, "a"),
+    (2, "b"),
+    (3, "b"),
+    (4, "b"),
+    (5, "c"),
+    (6, "c"),
+    (7, "c"),
+    (8, "c"),
+    (9, "c"),
+  ], ["id", "dummy"])
+
+  autolinker = AutoLinker()
+
+  # set row count manually
+  autolinker.data_rowcount = df_mock.count()
+
+  actual_entropy = autolinker._calculate_column_entropy(df_mock, "dummy")
+
+  expected_entropy = 1.0296530140645737
+
+  assert actual_entropy==expected_entropy
+
+
+def test_generate_rules():
+  mock_columns = [
+    ["a"],
+    ["a", "b"],
+    ["a", "b", "c"],
+    ["a", "a"]
+  ]
+
+  autolinker = AutoLinker()
+
+  expected_output = [
+    ['l.a = r.a'],
+    ['l.a = r.a', 'l.b = r.b'],
+    ['l.a = r.a', 'l.b = r.b', 'l.c = r.c'],
+    ['l.a = r.a', 'l.a = r.a']
+  ]
+
+  actual_output = [autolinker._generate_rules(c) for c in mock_columns]
+
+  assert expected_output==actual_output
+
+
+def test_generate_candidate_blocking_rules():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+  # mock data
+  df_mock = spark.createDataFrame([
+    (0, "a"),
+    (1, "a"),
+    (2, "b"),
+    (3, "b"),
+    (4, "b"),
+    (5, "c"),
+    (6, "c"),
+    (7, "c"),
+    (8, "c"),
+    (9, "c"),
+  ], ["id", "dummy"])
+
+  autolinker = AutoLinker()
+
+  expected_output = [['l.dummy = r.dummy']]
+
+  actual_output = autolinker._generate_candidate_blocking_rules(df_mock, ["dummy"], 100, "id", max_columns_per_rule=2)
+
+  assert expected_output==actual_output
+
+
+def test_clean_columns():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+  # mock data
+  df_mock = spark.createDataFrame([
+    (0, "aaaa"),
+    (1, "a123"),
+    (2, "aaa aaa"),
+    (3, " a 1 a 1"),
+    (4, "A 1 A 999"),
+    (5, "a---[a][a][]"),
+    (6, " [ ] [ ]"),
+    (7, "___#aasdf"),
+    (8, "___")
+  ], ["id", "dummy"])
+
+  autolinker = AutoLinker()
+
+  expected_output = spark.createDataFrame([
+    (0, "aaaa"),
+    (1, "a123"),
+    (2, "aaa aaa"),
+    (3, " a 1 a 1"),
+    (4, "a 1 a 999"),
+    (5, "aaa"),
+    (6, "    "),
+    (7, "aasdf"),
+    (8, "")
+  ], ["id", "dummy"])
+
+  actual_output = autolinker._clean_columns(df_mock, ["dummy"], cleaning="all")
+
+  assert expected_output.schema==actual_output.schema and expected_output.collect()==actual_output.collect()
+
+
+def test_create_comparison_list():
+  mock_space = {"comparisons": {
+    "a": {
+      "distance_function": {
+        "distance_function": "jaccard",
+        "threshold": 0.5
+      }
+    },
+    "b": {
+      "distance_function": {
+        "distance_function": "levenshtein",
+        "threshold": 1
+      }
+    },
+    "c": {
+      "distance_function": {
+        "distance_function": "jaro_winkler",
+        "threshold": 0.5
+      }
+    }
+  },
+  "blocking_rules": ["l.a = r.a"]
+  }
+
+  autolinker = AutoLinker()
+
+  expected_output = [
+    cl.jaccard_at_thresholds("a", 0.5),
+    cl.levenshtein_at_thresholds("b", 1),
+    cl.jaro_winkler_at_thresholds("c", 0.5)
+  ]
+
+  actual_output = autolinker._create_comparison_list(mock_space)
+
+
+  assert all([expected._comparison_dict["comparison_description"]==actual._comparison_dict["comparison_description"] for expected, actual in zip(expected_output, actual_output)])
+
+
+def test_calculate_entropy_delta():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+  # mock data
+  df_mock = spark.createDataFrame([
+    (0, "a"),
+    (1, "a"),
+    (2, "b"),
+    (3, "b"),
+    (4, "b"),
+    (5, "c"),
+    (6, "c"),
+    (7, "c"),
+    (8, "c"),
+    (9, "c"),
+  ], ["id", "dummy"])
+  
+  df_mock_dedupe = spark.createDataFrame([
+    (0, "a"),
+    (3, "b"),
+    (4, "b"),
+    (8, "c"),
+    (9, "c"),
+  ], ["id", "dummy"])
+
+  expected_output = -0.15561933979152887
+
+  autolinker = AutoLinker()
+
+  # set row count manually
+  autolinker.data_rowcount = df_mock.count()
+  actual_output = autolinker.calculate_entropy_delta(df_mock, df_mock_dedupe, ["dummy"])
+
+  assert expected_output==actual_output
+
+
+def test_create_hyperopt_space():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+  # mock data
+  df_mock = spark.createDataFrame([
+    (0, "a"),
+    (1, "a"),
+    (2, "b"),
+    (3, "b"),
+    (4, "b"),
+    (5, "c"),
+    (6, "c"),
+    (7, "c"),
+    (8, "c"),
+    (9, "c"),
+  ], ["id", "dummy"])
+
+  autolinker = AutoLinker()
+
+  space = autolinker._create_hyperopt_space(df_mock, ["dummy"], 1000, "id", max_columns_per_rule=2)
+
+  sample_space = hyperopt.pyll.stochastic.sample(space)
+
+  # keys exist
+  keys = ["blocking_rules", "comparisons"]
+  keys_exist = all([key in sample_space.keys() for key in keys])
+
+  #blocking rule type
+  blocking_rule_type = type(sample_space["blocking_rules"]) == tuple
+
+  # correct function
+  distance_functions = ["jaccard", "levenshtein", "jaro_winkler"]
+  correct_function = any([sample_space["comparisons"]["dummy"]["distance_function"]["distance_function"] in df for df in distance_functions])
+
+  # correct threshold
+  correct_threshold = sample_space["comparisons"]["dummy"]["distance_function"]["threshold"]>0.0
+
+  assert all([keys_exist, blocking_rule_type, correct_function, correct_threshold])
+
+
+def test_drop_intermediate_tables():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+
+  # create dummy tables in test catalog.schema
+  spark.sql("create or replace table marcell_splink.test_tables.test__splink__abc as select 'a' as test_col") # table to delete
+  spark.sql("create or replace table marcell_splink.test_tables.__splink__abc as select 'a' as test_col") # table to delete
+  spark.sql("create or replace table marcell_splink.test_tables.test__splink__ as select 'a' as test_col") # table to delete
+  spark.sql("create or replace table marcell_splink.test_tables.test_dont_delete as select 'a' as test_col") # table to ignore
+
+  tables_before = spark.sql("show tables from marcell_splink.test_tables").collect()
+
+  # check if __splink__ tables exist in the database before
+  splink_exists_before = any(["__splink__" in table.tableName for table in tables_before])
+
+  autolinker = AutoLinker(
+    catalog="marcell_splink",
+    schema="test_tables"
+  )
+  autolinker.spark = spark
+
+  # drop tables
+  autolinker._drop_intermediate_tables()
+
+  # collect tables again
+  tables_after = spark.sql("show tables from marcell_splink.test_tables").collect()
+
+  # check that __splink__ tables have been deleted
+  splink_tables_deleted = all(["__splink__" not in table.tableName for table in tables_after])
+
+  # check if the dummy table is still there
+  dummy_table_exists = any(["test_dont_delete" == table.tableName for table in tables_after])
+
+  assert all([splink_exists_before, splink_tables_deleted, dummy_table_exists])
+
+
+def test_convert_hyperopt_to_splink():
+  spark = SparkSession.builder.appName('abc').getOrCreate()
+  df_mock = spark.createDataFrame([
+    (0, "a", "1"),
+    (1, "a", "1"),
+    (2, "b", "2"),
+    (3, "b", "2"),
+    (4, "b", "3"),
+    (5, "c", "6"),
+    (6, "c", "7"),
+    (7, "c", "8"),
+    (8, "c", "9"),
+    (9, "c", "0"),
+  ], ["id", "dummy1", "dummy2"])
+
+  autolinker = AutoLinker(
+    catalog="marcell_splink",
+    schema="test_tables"
+  )
+  autolinker.spark = spark
+
+  # train autolinker
+  autolinker.auto_link(
+    df_mock,
+    attribute_columns=["dummy1", "dummy2"],
+    unique_id="id",
+    comparison_size_limit=1000,
+    max_evals=1
+  )
+
+  # convert best trial to dictionary
+  output = autolinker._convert_hyperopt_to_splink()
+
+  # keys exist
+  keys = ["blocking_rules", "comparisons"]
+  keys_exist = all([key in output.keys() for key in keys])
+
+  #blocking rule type
+  blocking_rule_type = type(output["blocking_rules"]) == list
+
+  # correct functions
+  distance_functions = ["jaccard", "levenshtein", "jaro_winkler"]
+  correct_function_dummy1 = any([output["comparisons"]["dummy1"]["distance_function"]["distance_function"] in df for df in distance_functions])
+  correct_function_dummy2 = any([output["comparisons"]["dummy2"]["distance_function"]["distance_function"] in df for df in distance_functions])
+  # correct threshold
+  correct_threshold_dummy1 = output["comparisons"]["dummy1"]["distance_function"]["threshold"]>0.0
+  correct_threshold_dummy2 = output["comparisons"]["dummy2"]["distance_function"]["threshold"]>0.0
+
+  assert all([keys_exist, blocking_rule_type, correct_function_dummy1, correct_function_dummy2, correct_threshold_dummy1, correct_threshold_dummy2])
+
+
+def test_randomise_columns():
+  autolinker = AutoLinker()
+
+  # test with 2 columns
+  test_attribute_columns_2 = ["a", "b"]
+  expected_randomised_columns_2 = ["a", "b"]
+
+  actual_randomised_columns_2 = autolinker._randomise_columns(test_attribute_columns_2)
+
+  two_columns_test = expected_randomised_columns_2==actual_randomised_columns_2
+
+  # test with 3 columns
+  test_attribute_columns_3 = ["a", "b", "c"]
+  expected_randomised_columns_3 = ["a&b", "a&c", "b&c", "b&a", "c&a", "c&b"]
+
+  actual_randomised_columns_3 = autolinker._randomise_columns(test_attribute_columns_3)
+
+  three_columns_test = all([col in expected_randomised_columns_3 for col in actual_randomised_columns_3])
+
+
+  assert all([two_columns_test, three_columns_test])
diff --git a/arc/tests/run_unit_testing.py b/arc/tests/run_unit_testing.py
new file mode 100644
index 0000000..53dc034
--- /dev/null
+++ b/arc/tests/run_unit_testing.py
@@ -0,0 +1,59 @@
+# Databricks notebook source
+# MAGIC %pip install --quiet pytest splink mlflow hyperopt
+
+# COMMAND ----------
+
+import pytest
+import os
+import sys
+import datetime
+
+# COMMAND ----------
+
+class ResultsCollector:
+    def __init__(self):
+        self.reports = []
+
+    @pytest.hookimpl(hookwrapper=True)
+    def pytest_runtest_makereport(self, item, call):
+        outcome = yield
+        report = outcome.get_result()
+        if report.when == 'call':
+            self.reports.append(report)
+
+# COMMAND ----------
+
+TIME_NOW = datetime.datetime.now()
+TIME_NOW
+
+# COMMAND ----------
+
+user = spark.sql("select current_user()").collect()[0]["current_user()"]
+
+# COMMAND ----------
+
+# Prepare to run pytest from the repo.
+os.chdir(f"/Workspace/Repos/{user}/splink-public-sector-hackathon/arc/tests")
+
+sys.dont_write_bytecode = True
+
+
+
+# COMMAND ----------
+
+collector = ResultsCollector()
+r = pytest.main(args=[".", "-vv", "-p","no:cacheprovider"], plugins=[collector])
+
+# COMMAND ----------
+
+test_report = spark.createDataFrame([
+  (TIME_NOW, report.head_line, report.duration, report.outcome, report.passed, report.failed) for report in collector.reports
+], ["test_run_timestamp", "test_name", "test_duration", "outcome", "passed", "failed"])
+
+# COMMAND ----------
+
+test_report.display()
+
+# COMMAND ----------
+
+