!pip install torch
!pip install git+https://github.com/nebw/mabe.git

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from oauth2client.client import GoogleCredentials

auth.authenticate_user()
gauth = GoogleAuth()
gauth.credentials = GoogleCredentials.get_application_default()
drive = GoogleDrive(gauth)

file_id = '1X7RGaoBtFwIvedfoAioZ8Ft1d7PvxbP0'

downloaded = drive.CreateFile({'id': file_id})
downloaded.GetContentFile('pretrained.pt')

import torch
import numpy as np
import pandas as pd

import mabe
import mabe.model

training_result = torch.load('pretrained.pt')[0]

# usually we would load these hyperparameters directly from the data, but here 
# we define them manually to avoid having to download and preprocess the 
# training data. they can be ignored in the context of this notebook.
num_extra_features = 4
num_extra_clf_tasks = 7

num_annotators = 6
annotators = np.arange(num_annotators)

# instantiate a new MultiAnnotatorLogisticRegressionHead object
logreg = mabe.model.MultiAnnotatorLogisticRegressionHead(
    training_result.config.num_context,
    num_annotators,
    num_extra_features,
    num_extra_clf_tasks,
)

def load_annotator_embeddings(logreg):
  # load pretrained annotator embeddings from model and store them in a pandas
  # dataframe for convenient analysis
  annotator_embedding = logreg.embedding.cpu().data.numpy().T
  annotator_embedding = pd.DataFrame(
      annotator_embedding, columns=[f"Annotator {a}" for a in range(6)]
  )
  return annotator_embedding

load_annotator_embeddings(logreg)

# load parameters from epoch with best validation loss during training
cpc_state_dict, logreg_state_dict = training_result.best_params[0]
logreg.load_state_dict(logreg_state_dict)
_ = logreg.eval()

annotator_embedding = load_annotator_embeddings(logreg)
annotator_embedding

import io
import seaborn as sns
import scipy
import scipy.special
import scipy.spatial as sp
import scipy.cluster.hierarchy as hc
import matplotlib.pyplot as plt

# perform and plot hierachical clustering on the annotator embeddings. similar
# embeddings (with a low euclidean distance) will be assigned to the same 
# cluster
cm = sns.clustermap(
    annotator_embedding, row_cluster=False, method="ward", 
    cbar_kws={"label": "Learned value"}
)
cm.ax_heatmap.set_ylabel("Embedding dimension")
plt.setp(cm.ax_heatmap.get_xticklabels(), rotation=90)

# sequence IDs from all test videos
sequences = list(training_result.test_predictions.keys())

distances = []

# iterate over all test video sequence IDs
for sequence_id in sequences:
  # this will store the probabilites the model assigns to each class for each 
  # frame, for each annotator
  predictions_by_annotator = []

  # iterate over all annotators
  for annotator_id in annotators:
    # the results are stored as compressed numpy arrays, so we need to load
    # them into a BytesIO object first s.t. numpy.load can read them like a 
    # normal file
    compressed_results = io.BytesIO(
        training_result.test_logits[sequence_id][annotator_id]
    )

    # we can now load the model predictions for the current video and 
    # annotator. We apply the softmax function to turn the model outputs
    # (logits) to probabilities
    logits = np.load(compressed_results)['arr_0']
    predictions = scipy.special.softmax(logits, -1)

    predictions_by_annotator.append(predictions)

  # we can now stack and flatten the predictions for each annotator
  predictions_by_annotator = np.stack(predictions_by_annotator)
  predictions_flat = predictions_by_annotator.reshape(num_annotators, -1)

  # loaded all predictions from all videos at once and calculating the distance 
  # matrix in one step would be quite memory intensive, so we simply calculate
  # the distance matrix for each video seperately and sum them up at the end.
  # note that pdist returns a `condensed distance matrix` which works out fine
  # in this case, because that's the format that scipy's linkage function
  # expects anyways.
  sequence_distances = scipy.spatial.distance.pdist(
      predictions_flat, 
      lambda u, v: ((u - v) ** 2).sum()
  )

  distances.append(sequence_distances)

distances = np.stack(distances)
# sum up the precomputed distance matrices and apply sqrt to get the final 
# euclidean distance matrix.
distances = np.sqrt(np.sum(distances, axis=0))

# finally, we can use scipy to perform and visualize the hierachical clustering.
Z = scipy.cluster.hierarchy.linkage(distances, 'ward')
dn = scipy.cluster.hierarchy.dendrogram(Z)
plt.xlabel('Annotator')