upload_trial

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+models/met_joint_f78b0fa9c1d7718b9ed703ddcf621ec9_lf_sd_train_gen_4/ filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+models/
+models_orig/
+baseline_src/wandb
+data/raw_data
+**/wandb/
+**/trash/
+**/.env
+**/__pycache__/
+**/.hydra/
+**/*result*.jsonl
+**/*nohup.out**
+**/extras/
+models_7_6_24/**
+results_old/**

app.py

@@ -0,0 +1,119 @@
+import time
+import spacy
+import json
+
+import gradio as gr
+from spacy.tokens import Doc, Span
+from spacy import displacy
+import matplotlib.pyplot as plt
+from matplotlib.colors import to_hex
+
+from inference.model_inference import Inference
+from configs import *
+
+
+def get_MEIRa_clusters(doc_name, text, model_type):
+    model_str = MODELS[model_type]
+    model = Inference(model_str)
+
+    output_dict = model.perform_coreference(text, doc_name)
+    return output_dict
+
+
+def coref_visualizer(doc_name, text, model_type):
+    coref_output = get_MEIRa_clusters(doc_name, text, model_type)
+
+    tokens = coref_output["tokenized_doc"]
+    clusters = coref_output["clusters"]
+    labels = coref_output["representative_names"]
+
+    ## Get a pastel palette
+    color_palette = {
+        label: to_hex(plt.cm.get_cmap("tab20", len(labels))(i))
+        for i, label in enumerate(labels)
+    }
+
+    nlp = spacy.blank("en")
+    doc = Doc(nlp.vocab, words=tokens)
+
+    print("Tokens:", tokens, flush=True)
+    # print("Doc:", doc, flush=True)
+    print(color_palette)
+
+    spans = []
+    for cluster_ind, cluster in enumerate(clusters[:-1]):
+        label = labels[cluster_ind]
+        for (start, end), mention in cluster:
+            span = Span(doc, start, end + 1, label=label)
+            spans.append(span)
+
+    doc.spans["coref_spans"] = spans
+
+    print("Rendering the visualization...")
+
+    # color_map = {label: color_palette[i] for i, label in enumerate(labels)}
+    # Generate the HTML output
+    html = displacy.render(
+        doc,
+        style="span",
+        options={
+            "spans_key": "coref_spans",
+            "colors": color_palette,
+        },
+        jupyter=False,
+    )
+
+    ## Create a hash based on time and doc_name
+    time_hash = hash(str(time.time()) + doc_name)
+    html_file = f"gradio_outputs/output_{time_hash}.html"
+    json_file = f"gradio_outputs/output_{time_hash}.json"
+
+    with open(html_file, "w") as f:
+        f.write(html)
+
+    with open(json_file, "w") as f:
+        json.dump(coref_output, f)
+
+    return (
+        html_file,
+        json_file,
+        gr.DownloadButton(value=html_file, visible=True),
+        gr.DownloadButton(value=json_file, visible=True),
+    )
+
+
+def download_html():
+    return gr.DownloadButton(visible=False)
+
+
+def download_json():
+    return gr.DownloadButton(visible=False)
+
+
+options = ["static", "hybrid"]
+
+with gr.Blocks() as demo:
+    html_file = gr.File(visible=False)
+    json_file = gr.File(visible=False)
+    html_button = gr.DownloadButton("Download HTML", visible=False)
+    json_button = gr.DownloadButton("Download JSON", visible=False)
+
+    html_button.click()
+    json_button.click()
+    iface = gr.Interface(
+        fn=coref_visualizer,
+        inputs=[
+            gr.Textbox(lines=1, placeholder="Enter document name:"),
+            gr.Textbox(lines=100, placeholder="Enter text for coreference resolution:"),
+            gr.Radio(choices=options, label="Select an Option"),
+        ],
+        outputs=[
+            html_file,
+            json_file,
+            html_button,
+            json_button,
+        ],
+        title="Coreference Resolution Visualizer",
+    )
+
+demo.launch(debug=True)

conf/config.yaml

@@ -0,0 +1,37 @@
+metrics: ['MUC', 'Bcub', 'CEAFE']
+keep_singletons: True
+seed: 45
+train: True
+use_wandb: True
+desc: "Major Entity Tracking"
+override_encoder: False
+log_vals: False
+# Useful for testing models with different memory architecture than the one trained on
+override_memory: False
+log_dir_add: ""
+device: "cuda:0"
+key: ""
+
+defaults:
+  - _self_
+  - datasets: litbank
+  - model: model
+  - optimizer: adam
+  - trainer: train
+  - infra: local
+  - experiment: debug
+
+paths:
+  resource_dir: "../data/"
+  base_data_dir: ${paths.resource_dir}/raw_data
+  conll_scorer: ${paths.resource_dir}/reference-coreference-scorers/scorer.pl
+  base_model_dir: ${infra.work_dir}/../models ## remove /../
+  model_dir: null
+  best_model_dir: null
+  model_filename: 'model.pth'
+  model_name: null
+  model_name_prefix: 'met_'
+  model_path: null
+  best_model_path: null
+  doc_encoder_dirname: 'doc_encoder'
+

conf/datasets/aft.yaml

@@ -0,0 +1,5 @@
+aft:
+  name: "aft"
+  targeted_eval: False
+  num_test_docs: 3
+  has_conll: True

conf/datasets/aft_increase.yaml

@@ -0,0 +1,5 @@
+aft_increase:
+  name: "aft_increase"
+  targeted_eval: False
+  num_test_docs: 23
+  has_conll: False

conf/datasets/aft_ind.yaml

@@ -0,0 +1,5 @@
+aft_ind:
+  name: "aft_ind"
+  targeted_eval: False
+  num_test_docs: 23
+  has_conll: False

conf/datasets/all.yaml

@@ -0,0 +1,11 @@
+defaults:
+  - litbank
+  - fantasy
+  - aft
+  - aft_increase
+  - aft_ind
+  - animal
+  - pride
+  - movie
+  - movie_cased
+  - avengers

conf/datasets/animal.yaml

@@ -0,0 +1,5 @@
+animal:
+  name: "animal"
+  targeted_eval: False
+  num_test_docs: 3
+  has_conll: True

conf/datasets/avengers.yaml

@@ -0,0 +1,5 @@
+avengers:
+  name: "avengers"
+  targeted_eval: False
+  num_test_docs: 1
+  has_conll: True

conf/datasets/fantasy.yaml

@@ -0,0 +1,7 @@
+fantasy:
+  name: "fantasy"
+  targeted_eval: False
+  num_train_docs: 171
+  num_dev_docs: 20
+  num_test_docs: 20
+  has_conll: True

conf/datasets/joint_lf.yaml

@@ -0,0 +1,3 @@
+defaults:
+  - litbank
+  - fantasy

conf/datasets/litbank.yaml

@@ -0,0 +1,8 @@
+litbank:
+  name: "LitBank"
+  cross_val_split: 0
+  targeted_eval: False
+  num_train_docs: 80
+  num_dev_docs: 10
+  num_test_docs: 10
+  has_conll: True

conf/datasets/movie.yaml

@@ -0,0 +1,5 @@
+movie:
+  name: "movie"
+  targeted_eval: False
+  num_test_docs: 6
+  has_conll: False

conf/datasets/movie_cased.yaml

@@ -0,0 +1,5 @@
+movie_cased:
+  name: "movie_cased"
+  targeted_eval: False
+  num_test_docs: 6
+  has_conll: False

conf/datasets/ontonotes.yaml

@@ -0,0 +1,12 @@
+ontonotes:
+  name: "OntoNotes"
+  targeted_eval: False
+  num_train_docs: 2802
+  num_dev_docs: 343
+  num_test_docs: 348
+  has_conll: True
+  # OntoNotes specific attributes
+#  use_genre_feature: False  # Whether to use document genre as a feature or not
+#  default_genre: "nw"
+#  genres: [ "bc", "bn", "mz", "nw", "pt", "tc", "wb" ]
+  singleton_file: null  # File path with pseudo-singletons

conf/datasets/preco.yaml

@@ -0,0 +1,6 @@
+preco:
+  name: "PreCo"
+  targeted_eval: False
+  num_train_docs: 3000
+  num_dev_docs: 500
+  num_test_docs: 500

conf/experiment/eval_all.yaml

@@ -0,0 +1,13 @@
+# @package _global_
+
+# Evaluate all models
+
+
+defaults:
+  - override /datasets: all
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+    
+train: False

conf/experiment/lf_coref_id.yaml

@@ -0,0 +1,26 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+
+# Model name in CRAC 2021: longdoc^S Joint
+
+
+defaults:
+  - override /datasets: joint_lf
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 500
+  max_evals: 20
+  eval_per_k_steps: null
+  patience: 10
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+  memory:
+    pseudo_dist: False
+
+log_dir_add: "coref_id"

conf/experiment/lf_eval.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+
+# Model name in CRAC 2021: longdoc^S Joint
+
+
+defaults:
+  - override /datasets: joint_lf
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 500
+  max_evals: 25
+  eval_per_k_steps: null
+  patience: 10
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+

conf/experiment/lf_extment.yaml

@@ -0,0 +1,31 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+
+# Model name in CRAC 2021: longdoc^S Joint
+
+
+defaults:
+  - override /datasets: joint_lf
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 500
+  max_evals: 20
+  eval_per_k_steps: null
+  patience: 10
+
+datasets:
+  litbank:
+    external_md_file: "litbank/longformer_speaker/0/mentions_ment_model_litbank_eval.jsonl"
+  fantasy:
+    external_md_file: "fantasy/longformer_speaker/mentions_ment_model_fantasy_eval.jsonl"
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+  mention_params:
+    ext_ment: True
+

conf/experiment/lf_hybrid.yaml

@@ -0,0 +1,25 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+
+# Model name in CRAC 2021: longdoc^S Joint
+
+
+defaults:
+  - override /datasets: joint_lf
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 500
+  max_evals: 25
+  eval_per_k_steps: null
+  patience: 10
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+  memory:
+    type: hybrid
+

conf/experiment/lf_static.yaml

@@ -0,0 +1,25 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+# OntoNotes also uses pseudo-singletons
+
+# Model name in CRAC 2021: longdoc^S Joint + PS 30K
+
+
+defaults:
+  - override /datasets: joint_lf
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 500
+  max_evals: 25
+  eval_per_k_steps: null
+  patience: 10
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+  memory:
+    type: "static"

conf/experiment/litbank.yaml

@@ -0,0 +1,21 @@
+# @package _global_
+
+# Vanilla LitBank configuration
+
+# Model name in CRAC 2021: longdoc LB_0
+
+
+defaults:
+  - override /datasets: litbank
+  - override /trainer: train.yaml
+
+trainer:
+  log_frequency: 10
+  max_evals: 40
+  patience: 20
+  eval_per_k_steps: null
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+

conf/experiment/onto_pseudo_hybrid.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+
+# Model name in CRAC 2021: longdoc^S Joint
+
+
+defaults:
+  - override /datasets: ontonotes
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 250
+  max_evals: 20
+  eval_per_k_steps: null
+  patience: 10
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+  memory:
+    type: hybrid
+
+datasets:
+  ontonotes:
+    singleton_file: ontonotes/ment_singletons_longformer_speaker/30.jsonlines
+

conf/experiment/onto_pseudo_static.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+
+# This configuration trains a joint model for Ontonotes, Litbank, and Preco.
+# Note that OntoNotes and Preco are downsampled in this configuration.
+# OntoNotes also uses pseudo-singletons
+
+# Model name in CRAC 2021: longdoc^S Joint + PS 30K
+
+
+defaults:
+  - override /datasets: ontonotes
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 500
+  max_evals: 20
+  eval_per_k_steps: null
+  patience: 10
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+  memory:
+    type: "static"
+
+datasets:
+  ontonotes:
+    singleton_file: ontonotes/ment_singletons_longformer_speaker/30.jsonlines

conf/experiment/ontonotes.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+
+# Vanilla ontonotes configuration which doesn't assume any upstream features
+# of speaker and document genre
+
+# Model name in CRAC 2021: longdoc ON
+
+defaults:
+  - override /datasets: ontonotes
+  - override /trainer: train.yaml
+
+trainer:
+  log_frequency: 250
+  patience: 10
+  eval_per_k_steps: 5000
+
+

conf/experiment/ontonotes_pseudo.yaml

@@ -0,0 +1,27 @@
+# @package _global_
+
+# This configuration trains ontonotes using the speaker information and pseudo singletons.
+# This is the best OntoNotes configuration in our CRAC 2021 work.
+# Note that this configuration doesn't assume other upstream features such as document genre
+
+# Model name in CRAC 2021: longdoc^S ON + PS 60K
+
+
+defaults:
+  - override /datasets: ontonotes
+  - override /trainer: train.yaml
+  - override /model: model.yaml
+
+trainer:
+  log_frequency: 250
+  max_evals: 20
+  patience: 10
+  eval_per_k_steps: null
+
+model:
+  doc_encoder:
+    add_speaker_tokens: True
+
+datasets:
+  ontonotes:
+    singleton_file: ontonotes/ment_singletons_longformer_speaker/30.jsonlines

conf/infra/local.yaml

@@ -0,0 +1,8 @@
+is_local: True
+work_dir: "./"
+
+
+#hydra:
+#  run:
+#    dir:
+#      "~/Research/fast-coref/models"

conf/infra/slurm.yaml

@@ -0,0 +1,10 @@
+is_local: False
+job_time: 14280
+job_id: null
+work_dir: "./"
+
+
+#hydra:
+#  run:
+#    dir:
+#      /share/data/speech/shtoshni/research/fast-coref/slurm_scripts/thesis/${job_id}.log

conf/model/doc_encoder/transformer/longformer_large.yaml

@@ -0,0 +1,5 @@
+name: 'longformer'
+model_size: 'large'
+model_str: 'allenai/longformer-large-4096'
+max_encoder_segment_len: 4096
+max_segment_len: 4096

conf/model/doc_encoder/transformer_encoder.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - transformer: longformer_large
+
+chunking: independent
+finetune: true  # Add logic of finetuning depending on the training logic
+add_speaker_tokens: true  # Change this value depending on the dataset
+speaker_start: '[SPEAKER_START]'
+speaker_end: '[SPEAKER_END]'
+
+

conf/model/memory/mem_type/unbounded.yaml

@@ -0,0 +1,3 @@
+name: unbounded
+max_ents: null
+eval_max_ents: null

conf/model/memory/memory.yaml

@@ -0,0 +1,15 @@
+defaults:
+ - mem_type: unbounded
+
+emb_size: 20
+mlp_size: 3000
+mlp_depth: 1
+sim_func: hadamard
+entity_rep: wt_avg
+num_feats: 2 ## Change this to remove position information.
+thresh: 0.0
+rep_pos: "learned"
+pseudo_dist: True
+num_embeds: 10
+type: "dyn"
+batch_size: 64

conf/model/model.yaml

@@ -0,0 +1,24 @@
+defaults:
+ - doc_encoder: transformer_encoder
+ - memory: memory
+
+mention_params:
+  max_span_width: 20
+  ment_emb: attn
+  use_gold_ments: false
+  ext_ment: false
+  use_topk: false
+  top_span_ratio: 0.4
+  emb_size: 20
+  mlp_size: 3000
+  mlp_depth: 1
+  ment_emb_to_size_factor:
+    attn: 3
+    endpoint: 2
+    max: 1
+  ignore_non_gold: True
+
+metadata_params:
+  use_genre_feature: False
+  default_genre: "nw"
+  genres: [ "bc", "bn", "mz", "nw", "pt", "tc", "wb" ]

conf/optimizer/adam.yaml

@@ -0,0 +1,4 @@
+init_lr: 3e-4
+fine_tune_lr: 1e-5
+max_gradient_norm: 1.0
+lr_decay: linear

conf/trainer/train.yaml

@@ -0,0 +1,13 @@
+dropout_rate: 0.3
+label_smoothing_wt: 0.1
+ment_loss_mode: 'all'
+normalize_loss: False
+ment_loss_incl: True
+max_evals: 20
+to_save_model: False
+log_frequency: 500
+patience: 10
+eval_per_k_steps: null
+num_training_steps: null
+max_training_segments: 1
+generalise: True

configs.py

@@ -0,0 +1,4 @@
+MODELS = {
+    "static": "models/met_joint_efbc65248a6aedce066a04a9f4f40084_lf_s_train_gen_5",
+    "hybrid": "models/met_joint_f78b0fa9c1d7718b9ed703ddcf621ec9_lf_sd_train_gen_4",
+}

coref_utils/conll.py

@@ -0,0 +1,126 @@
+import re
+import subprocess
+import operator
+import collections
+
+BEGIN_DOCUMENT_REGEX = re.compile(r"#begin document \((.*)\); part (\d+)")
+COREF_RESULTS_REGEX = re.compile(
+    r".*Coreference: Recall: \([0-9.]+ / [0-9.]+\) ([0-9.]+)%\tPrecision: \([0-9.]+ / [0-9.]+\) "
+    r"([0-9.]+)%\tF1: ([0-9.]+)%.*",
+    re.DOTALL,
+)
+
+
+def get_doc_key(doc_id, part):
+    return "{}_{}".format(doc_id, int(part))
+
+
+def output_conll(input_file, output_file, predictions, subtoken_map):
+    prediction_map = {}
+    for doc_key, clusters in predictions.items():
+        start_map = collections.defaultdict(list)
+        end_map = collections.defaultdict(list)
+        word_map = collections.defaultdict(list)
+        for cluster_id, mentions in enumerate(clusters):
+            for start, end in mentions:
+                start, end = subtoken_map[doc_key][start], subtoken_map[doc_key][end]
+                if start == end:
+                    word_map[start].append(cluster_id)
+                else:
+                    start_map[start].append((cluster_id, end))
+                    end_map[end].append((cluster_id, start))
+        for k, v in start_map.items():
+            start_map[k] = [
+                cluster_id
+                for cluster_id, end in sorted(
+                    v, key=operator.itemgetter(1), reverse=True
+                )
+            ]
+        for k, v in end_map.items():
+            end_map[k] = [
+                cluster_id
+                for cluster_id, start in sorted(
+                    v, key=operator.itemgetter(1), reverse=True
+                )
+            ]
+        prediction_map[doc_key] = (start_map, end_map, word_map)
+
+    word_index = 0
+    for line in input_file.readlines():
+        row = line.split()
+        if len(row) == 0:
+            output_file.write("\n")
+        elif row[0].startswith("#"):
+            begin_match = re.match(BEGIN_DOCUMENT_REGEX, line)
+            if begin_match:
+                doc_key = get_doc_key(begin_match.group(1), begin_match.group(2))
+                start_map, end_map, word_map = prediction_map[doc_key]
+                word_index = 0
+            output_file.write(line)
+            # output_file.write("\n")
+        else:
+            assert get_doc_key(row[0], row[1]) == doc_key
+            coref_list = []
+            if word_index in end_map:
+                for cluster_id in end_map[word_index]:
+                    coref_list.append("{})".format(cluster_id))
+            if word_index in word_map:
+                for cluster_id in word_map[word_index]:
+                    coref_list.append("({})".format(cluster_id))
+            if word_index in start_map:
+                for cluster_id in start_map[word_index]:
+                    coref_list.append("({}".format(cluster_id))
+
+            if len(coref_list) == 0:
+                row[-1] = "-"
+            else:
+                row[-1] = "|".join(coref_list)
+
+            output_file.write("   ".join(row))
+            output_file.write("\n")
+            word_index += 1
+
+
+def official_conll_eval(
+    conll_scorer, gold_path, predicted_path, metric, official_stdout=False
+):
+    cmd = [conll_scorer, metric, gold_path, predicted_path, "none"]
+    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+    stdout, stderr = process.communicate()
+    process.wait()
+
+    stdout = stdout.decode("utf-8")
+    if stderr is not None:
+        print(stderr)
+
+    if official_stdout:
+        print("Official result for {}".format(metric))
+        print(stdout)
+
+    coref_results_match = re.match(COREF_RESULTS_REGEX, stdout)
+    recall = float(coref_results_match.group(1))
+    precision = float(coref_results_match.group(2))
+    f1 = float(coref_results_match.group(3))
+    return {"r": recall, "p": precision, "f": f1}
+
+
+def evaluate_conll(
+    conll_scorer,
+    gold_path,
+    predictions,
+    subtoken_maps,
+    prediction_path,
+    all_metrics=False,
+    official_stdout=False,
+):
+    with open(prediction_path, "w") as prediction_file:
+        with open(gold_path, "r") as gold_file:
+            output_conll(gold_file, prediction_file, predictions, subtoken_maps)
+
+    result = {
+        metric: official_conll_eval(
+            conll_scorer, gold_file.name, prediction_file.name, metric, official_stdout
+        )
+        for metric in ("muc", "bcub", "ceafe")
+    }
+    return result

coref_utils/metrics.py

@@ -0,0 +1,198 @@
+import numpy as np
+from collections import Counter
+from scipy.optimize import linear_sum_assignment
+
+
+def f1(p_num, p_den, r_num, r_den, beta=1):
+    p = 0 if p_den == 0 else p_num / float(p_den)
+    r = 0 if r_den == 0 else r_num / float(r_den)
+    return 0 if p + r == 0 else (1 + beta * beta) * p * r / (beta * beta * p + r)
+
+
+class CorefEvaluator(object):
+    def __init__(self):
+        self.evaluators = [Evaluator(m) for m in (muc, b_cubed, ceafe)]
+
+    def update(self, predicted, gold, mention_to_predicted, mention_to_gold):
+        for e in self.evaluators:
+            e.update(predicted, gold, mention_to_predicted, mention_to_gold)
+
+    def get_f1(self):
+        return sum(e.get_f1() for e in self.evaluators) / len(self.evaluators)
+
+    def get_recall(self):
+        return sum(e.get_recall() for e in self.evaluators) / len(self.evaluators)
+
+    def get_precision(self):
+        return sum(e.get_precision() for e in self.evaluators) / len(self.evaluators)
+
+    def get_prf(self):
+        return self.get_precision(), self.get_recall(), self.get_f1()
+
+
+class F1Evaluator(object):
+    def __init__(self):
+        self.f1_macro_sum = 0.0
+        self.f1_micro_sum = 0.0
+        self.macro_support = 0
+        self.micro_support = 0
+
+    def update(self, predicted, gold):
+        if gold:
+            for cluster_ind, cluster in enumerate(gold):
+                predicted_set = set(predicted[cluster_ind])
+                correct = set(cluster).intersection(set(predicted_set))
+                num_correct = len(correct)
+                num_predicted = len(predicted_set)
+                num_gt = len(cluster)
+                precision = num_correct / num_predicted if num_predicted > 0 else 0
+                recall = num_correct / num_gt if num_gt > 0 else 0
+                f1_score = (
+                    2 * precision * recall / (precision + recall)
+                    if precision + recall > 0
+                    else 0
+                )
+                support_entity_micro = num_gt
+                support_entity_macro = 1
+                self.f1_macro_sum += f1_score * support_entity_macro
+                self.f1_micro_sum += f1_score * support_entity_micro
+                self.macro_support += support_entity_macro
+                self.micro_support += support_entity_micro
+
+    def get_numbers(self):
+        f1_macro = (
+            (self.f1_macro_sum / self.macro_support) * 100
+            if self.macro_support > 0
+            else 0
+        )
+        f1_micro = (
+            (self.f1_micro_sum / self.micro_support) * 100
+            if self.micro_support > 0
+            else 0
+        )
+        return f1_macro, f1_micro
+
+
+class Evaluator(object):
+    def __init__(self, metric, beta=1):
+        self.p_num = 0
+        self.p_den = 0
+        self.r_num = 0
+        self.r_den = 0
+        self.metric = metric
+        self.beta = beta
+
+    def update(self, predicted, gold, mention_to_predicted, mention_to_gold):
+        if self.metric == ceafe:
+            pn, pd, rn, rd = self.metric(predicted, gold)
+        else:
+            pn, pd = self.metric(predicted, mention_to_gold)
+            rn, rd = self.metric(gold, mention_to_predicted)
+        self.p_num += pn
+        self.p_den += pd
+        self.r_num += rn
+        self.r_den += rd
+
+    def get_f1(self):
+        return f1(self.p_num, self.p_den, self.r_num, self.r_den, beta=self.beta)
+
+    def get_recall(self):
+        return 0 if self.r_num == 0 else self.r_num / float(self.r_den)
+
+    def get_precision(self):
+        return 0 if self.p_num == 0 else self.p_num / float(self.p_den)
+
+    def get_prf(self):
+        return self.get_precision(), self.get_recall(), self.get_f1()
+
+    def get_counts(self):
+        return self.p_num, self.p_den, self.r_num, self.r_den
+
+    def get_prf_str(self):
+        perf_str = (
+            f"Recall: {self.get_recall() * 100}, Precision: {self.get_precision() * 100}, "
+            f"F-score: {self.get_f1() * 100}\n"
+        )
+
+        return perf_str
+
+
+def evaluate_documents(documents, metric, beta=1):
+    evaluator = Evaluator(metric, beta=beta)
+    for document in documents:
+        evaluator.update(document)
+    return evaluator.get_precision(), evaluator.get_recall(), evaluator.get_f1()
+
+
+def b_cubed(clusters, mention_to_gold):
+    num, dem = 0, 0
+
+    for c in clusters:
+        gold_counts = Counter()
+        correct = 0
+        for m in c:
+            if m in mention_to_gold:
+                gold_counts[tuple(mention_to_gold[m])] += 1
+        for c2, count in gold_counts.items():
+            correct += count * count
+
+        num += correct / float(len(c))
+        dem += len(c)
+
+    return num, dem
+
+
+def muc(clusters, mention_to_gold):
+    tp, p = 0, 0
+    for c in clusters:
+        p += len(c) - 1
+        tp += len(c)
+        linked = set()
+        for m in c:
+            if m in mention_to_gold:
+                linked.add(mention_to_gold[m])
+            else:
+                tp -= 1
+        tp -= len(linked)
+    return tp, p
+
+
+def phi4(c1, c2):
+    return 2 * len([m for m in c1 if m in c2]) / float(len(c1) + len(c2))
+
+
+def ceafe(clusters, gold_clusters):
+    scores = np.zeros((len(gold_clusters), len(clusters)))
+    for i in range(len(gold_clusters)):
+        for j in range(len(clusters)):
+            scores[i, j] = phi4(gold_clusters[i], clusters[j])
+    matching = linear_sum_assignment(-scores)
+    matching = np.asarray(matching)
+    matching = np.transpose(matching)
+
+    similarity = sum(scores[matching[:, 0], matching[:, 1]])
+    return similarity, len(clusters), similarity, len(gold_clusters)
+
+
+def lea(clusters, mention_to_gold):
+    num, dem = 0, 0
+
+    for c in clusters:
+        if len(c) == 1:
+            continue
+
+        common_links = 0
+        all_links = len(c) * (len(c) - 1) / 2.0
+        for i, m in enumerate(c):
+            if m in mention_to_gold:
+                for m2 in c[i + 1 :]:
+                    if (
+                        m2 in mention_to_gold
+                        and mention_to_gold[m] == mention_to_gold[m2]
+                    ):
+                        common_links += 1
+
+        num += len(c) * common_links / float(all_links)
+        dem += len(c)
+
+    return num, dem

coref_utils/utils.py

@@ -0,0 +1,43 @@
+from typing import List, Dict, Tuple
+
+
+def filter_clusters(clusters: List, threshold: int = 1) -> List:
+    """Filter clusters with mentions less than the specified threshold."""
+
+    return [
+        tuple(tuple(mention) for mention in cluster)
+        for cluster_ind,cluster in enumerate(clusters)
+        if len(cluster) >= threshold and cluster_ind != len(clusters) - 1   # last cluster is always removed.
+    ]
+
+
+def get_mention_to_cluster(clusters: List) -> Dict:
+    """Get mention to cluster mapping."""
+
+    clusters = [tuple(tuple(mention) for mention in cluster) for cluster in clusters]
+    mention_to_cluster_dict = {}
+    for cluster in clusters:
+        for mention in cluster:
+            mention_to_cluster_dict[mention] = cluster
+    return mention_to_cluster_dict
+
+
+def get_mention_to_cluster_idx(clusters: List) -> Dict:
+    """Get mention to cluster idx mapping while filtering clustering."""
+
+    clusters = [tuple(tuple(mention) for mention in cluster) for cluster in clusters]
+    mention_to_cluster_dict = {}
+    for cluster_idx, cluster in enumerate(clusters):
+        for mention in cluster:
+            mention_to_cluster_dict[mention] = cluster_idx
+    return mention_to_cluster_dict
+
+
+def is_aligned(span1: Tuple[int, int], span2: Tuple[int, int]) -> bool:
+    """Return true if one of the span is a substring of the other span."""
+
+    if span1[0] >= span2[0] and span1[1] <= span2[1]:
+        return True
+    if span2[0] >= span1[0] and span2[1] <= span1[1]:
+        return True
+    return False

data_utils/tensorize_dataset.py

@@ -0,0 +1,76 @@
+import torch
+from typing import List, Dict, Union
+from transformers import PreTrainedTokenizerFast
+from torch import Tensor
+
+
+class TensorizeDataset:
+    def __init__(
+        self, tokenizer: PreTrainedTokenizerFast, remove_singletons: bool = False
+    ) -> None:
+        self.tokenizer = tokenizer
+        self.remove_singletons = remove_singletons
+        self.device = torch.device("cpu")
+
+    def tensorize_data(
+        self, split_data: List[Dict], training: bool = False
+    ) -> List[Dict]:
+        tensorized_data = []
+        for document in split_data:
+            tensorized_data.append(
+                self.tensorize_instance_independent(document, training=training)
+            )
+
+        return tensorized_data
+
+    def process_segment(self, segment: List) -> List:
+        if self.tokenizer.sep_token_id is None:
+            # print("SentencePiece Tokenizer")
+            return [self.tokenizer.bos_token_id] + segment + [self.tokenizer.eos_token_id]
+        else:
+            # print("WordPiece Tokenizer")
+            return [self.tokenizer.cls_token_id] + segment + [self.tokenizer.sep_token_id]
+    
+    def tensorize_instance_independent(
+        self, document: Dict, training: bool = False
+    ) -> Dict:
+        segments: List[List[int]] = document["sentences"]
+        clusters: List = document.get("clusters", [])
+        ext_predicted_mentions: List = document.get("ext_predicted_mentions", [])
+        sentence_map: List[int] = document["sentence_map"]
+        subtoken_map: List[int] = document["subtoken_map"]
+        representatives: List = document.get("representatives", [])
+        representative_embs: List = document.get("representative_embs", [])
+        
+        tensorized_sent: List[Tensor] = [
+            torch.unsqueeze(
+                torch.tensor(self.process_segment(sent), device=self.device), dim=0
+            )
+            for sent in segments
+        ]
+
+        sent_len_list = [len(sent) for sent in segments]
+        output_dict = {
+            "tensorized_sent": tensorized_sent,
+            "sentences": segments,
+            "sent_len_list": sent_len_list,
+            "doc_key": document.get("doc_key", None),
+            "clusters": clusters,
+            "ext_predicted_mentions": ext_predicted_mentions,
+            "subtoken_map": subtoken_map,
+            "sentence_map": torch.tensor(sentence_map, device=self.device),
+            "representatives": representatives,
+            "representative_embs": representative_embs,
+        }
+
+        # Pass along other metadata
+        for key in document:
+            if key not in output_dict:
+                output_dict[key] = document[key]
+
+        if self.remove_singletons:
+            output_dict["clusters"] = [
+                cluster for cluster in output_dict["clusters"] if len(cluster) > 1
+            ]
+
+        return output_dict

data_utils/utils.py

@@ -0,0 +1,95 @@
+import json
+from os import path
+from typing import Dict
+import jsonlines
+
+
+def get_data_file(data_dir: str, split: str, max_segment_len: int) -> str:
+    jsonl_file = path.join(
+        data_dir, "{}.{}.met.jsonlines".format(split, max_segment_len)
+    )
+    print("File access: ", jsonl_file)
+    if path.exists(jsonl_file):
+        return jsonl_file
+    else:
+        jsonl_file = path.join(data_dir, "{}.met.jsonlines".format(split))
+        if path.exists(jsonl_file):
+            return jsonl_file
+
+
+def load_dataset(
+    data_dir: str,
+    singleton_file: str = None,
+    max_segment_len: int = 2048,
+    num_train_docs: int = None,
+    num_dev_docs: int = None,
+    num_test_docs: int = None,
+    dataset_name: str = None,
+) -> Dict:
+    all_splits = []
+    for split in ["train", "dev", "test"]:
+        jsonl_file = get_data_file(data_dir, split, max_segment_len)
+        if jsonl_file is None:
+            raise ValueError(f"No relevant files at {data_dir}")
+        split_data = []
+        with open(jsonl_file) as f:
+            for line in f:
+                load_dict = json.loads(line.strip())
+                load_dict["dataset_name"] = dataset_name
+                split_data.append(load_dict)
+        all_splits.append(split_data)
+
+    train_data, dev_data, test_data = all_splits
+
+    if singleton_file is not None and path.exists(singleton_file):
+        num_singletons = 0
+        with open(singleton_file) as f:
+            singleton_data = json.loads(f.read())
+
+        for instance in train_data:
+            doc_key = instance["doc_key"]
+            if doc_key in singleton_data:
+                if len(instance["clusters"]) != 0:
+                    num_singletons += len(singleton_data[doc_key])
+                    instance["clusters"][-1].extend(
+                        [cluster[0] for cluster in singleton_data[doc_key]]
+                    )
+
+        print("Added %d singletons" % num_singletons)
+
+    return {
+        "train": train_data[:num_train_docs],
+        "dev": dev_data[:num_dev_docs],
+        "test": test_data[:num_test_docs],
+    }
+
+
+def load_eval_dataset(
+    data_dir: str, external_md_file: str, max_segment_len: int, dataset_name: str = None
+) -> Dict:
+    data_dict = {}
+    for split in ["dev", "test"]:
+        jsonl_file = get_data_file(data_dir, split, max_segment_len)
+        if jsonl_file is not None:
+            split_data = []
+            with open(jsonl_file) as f:
+                for line in f:
+                    load_dict = json.loads(line.strip())
+                    load_dict["dataset_name"] = dataset_name
+                    split_data.append(load_dict)
+
+            data_dict[split] = split_data
+
+    if external_md_file is not None and path.exists(external_md_file):
+        predicted_mentions = {}
+        with jsonlines.open(external_md_file, mode="r") as reader:
+            for line in reader:
+                predicted_mentions[line["doc_key"]] = line
+        for split in ["dev", "test"]:
+            for instance in data_dict[split]:
+                doc_key = instance["doc_key"]
+                if doc_key in predicted_mentions:
+                    instance["ext_predicted_mentions"] = sorted(
+                        predicted_mentions[doc_key]["pred_mentions"]
+                    )
+    return data_dict

error_analysis/missing_clusters.py

@@ -0,0 +1,99 @@
+import argparse
+import os
+import logging
+import json
+import numpy as np
+from coref_utils.metrics import CorefEvaluator
+from coref_utils.utils import get_mention_to_cluster
+
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+logging.basicConfig(format="%(message)s", level=logging.INFO)
+logger = logging.getLogger()
+
+
+def process_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser()
+
+    # Add arguments to parser
+    parser.add_argument("log_file", help="Log file", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def singleton_analysis(data):
+    max_length = 0
+    max_doc_id = ""
+    max_cluster = []
+
+    for instance in data:
+
+        gold_clusters, gold_mentions_to_cluster = get_mention_to_cluster(
+            instance["clusters"]
+        )
+        pred_clusters, pred_mentions_to_cluster = get_mention_to_cluster(
+            instance["predicted_clusters"]
+        )
+
+        for cluster in gold_clusters:
+            all_mention_unseen = True
+            for mention in cluster:
+                if mention in pred_mentions_to_cluster:
+                    all_mention_unseen = False
+                    break
+
+            if all_mention_unseen:
+                if len(cluster) > max_length:
+                    max_length = len(cluster)
+                    max_doc_id = instance["doc_key"]
+                    max_cluster = cluster
+
+    print(max_doc_id)
+    print(max_length, max_cluster)
+
+
+def reverse_analysis(data):
+    max_length = 0
+    max_doc_id = ""
+    max_cluster = []
+
+    for instance in data:
+
+        gold_clusters, gold_mentions_to_cluster = get_mention_to_cluster(
+            instance["clusters"]
+        )
+        pred_clusters, pred_mentions_to_cluster = get_mention_to_cluster(
+            instance["predicted_clusters"]
+        )
+
+        for cluster in pred_clusters:
+            all_mention_unseen = True
+            for mention in cluster:
+                if mention in gold_mentions_to_cluster:
+                    all_mention_unseen = False
+                    break
+
+            if all_mention_unseen:
+                if len(cluster) > max_length:
+                    max_length = len(cluster)
+                    max_doc_id = instance["doc_key"]
+                    max_cluster = cluster
+
+    print(max_doc_id)
+    print(max_length, max_cluster)
+
+
+def main():
+    args = process_args()
+    data = []
+    with open(args.log_file) as f:
+        for line in f:
+            data.append(json.loads(line))
+    singleton_analysis(data)
+    reverse_analysis(data)
+
+
+if __name__ == "__main__":
+    main()

error_analysis/singleton_analysis.py

@@ -0,0 +1,120 @@
+import argparse
+import os
+import logging
+import json
+import numpy as np
+from coref_utils.metrics import CorefEvaluator
+from coref_utils.utils import get_mention_to_cluster, filter_clusters
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+logging.basicConfig(format="%(message)s", level=logging.INFO)
+logger = logging.getLogger()
+
+
+def process_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser()
+
+    # Add arguments to parser
+    parser.add_argument("log_file", help="Log file", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def singleton_analysis(data):
+    gold_singletons = 0
+    pred_singletons = 0
+
+    # singleton_evaluator = CorefEvaluator()
+    non_singleton_evaluator = CorefEvaluator()
+
+    gold_cluster_lens = []
+    pred_cluster_lens = []
+
+    overlap_sing = 0
+    total_sing = 0
+    pred_sing = 0
+
+    for instance in data:
+        # Singleton performance
+        gold_clusters = set(
+            [tuple(cluster[0]) for cluster in instance["clusters"] if len(cluster) == 1]
+        )
+        pred_clusters = set(
+            [
+                tuple(cluster[0])
+                for cluster in instance["predicted_clusters"]
+                if len(cluster) == 1
+            ]
+        )
+
+        total_sing += len(gold_clusters)
+        pred_sing += len(pred_clusters)
+        overlap_sing += len(gold_clusters.intersection(pred_clusters))
+
+        gold_singletons += len(gold_clusters)
+        pred_singletons += len(pred_clusters)
+
+        # predicted_clusters, mention_to_predicted = get_mention_to_cluster(pred_clusters, threshold=1)
+        # gold_clusters, mention_to_gold = get_mention_to_cluster(gold_clusters, threshold=1)
+        # singleton_evaluator.update(predicted_clusters, gold_clusters, mention_to_predicted, mention_to_gold)
+
+        # Non-singleton performance
+        gold_clusters = filter_clusters(instance["clusters"], threshold=2)
+        pred_clusters = filter_clusters(instance["predicted_clusters"], threshold=2)
+
+        gold_cluster_lens.extend([len(cluster) for cluster in instance["clusters"]])
+        pred_cluster_lens.extend(
+            [len(cluster) for cluster in instance["predicted_clusters"]]
+        )
+
+        # gold_clusters = filter_clusters(gold_clusters, threshold=1)
+        # pred_clusters = filter_clusters(pred_clusters, threshold=1)
+
+        mention_to_predicted = get_mention_to_cluster(pred_clusters)
+        mention_to_gold = get_mention_to_cluster(gold_clusters)
+        non_singleton_evaluator.update(
+            pred_clusters, gold_clusters, mention_to_predicted, mention_to_gold
+        )
+
+    logger.info(
+        "\nGT singletons: %d, Pred singletons: %d\n"
+        % (gold_singletons, pred_singletons)
+    )
+    recall_sing = overlap_sing / total_sing
+    pred_sing = overlap_sing / pred_sing
+    f_sing = 2 * recall_sing * pred_sing / (recall_sing + pred_sing)
+    logger.info(
+        f"\nSingletons - Recall: {recall_sing * 100}, Pred: {pred_sing * 100}, "
+        f"F1: {f_sing * 100}\n"
+    )
+    logger.info(
+        f"\nNon-singleton cluster lengths, Gold: {np.mean(gold_cluster_lens):.2f}, "
+        f"Pred: {np.mean(pred_cluster_lens)}\n"
+    )
+
+    for evaluator, evaluator_str in zip([non_singleton_evaluator], ["Non-singleton"]):
+        perf_str = ""
+        indv_metrics_list = ["MUC", "BCub", "CEAFE"]
+        for indv_metric, indv_evaluator in zip(indv_metrics_list, evaluator.evaluators):
+            # perf_str += ", " + indv_metric + ": {:.1f}".format(indv_evaluator.get_f1() * 100)
+            perf_str += "{} - {}".format(indv_metric, indv_evaluator.get_prf_str())
+
+        fscore = evaluator.get_f1() * 100
+        perf_str += "{} ".format(fscore)
+        perf_str = perf_str.strip(", ")
+        logger.info("\n%s\n%s\n" % (evaluator_str, perf_str))
+
+
+def main():
+    args = process_args()
+    data = []
+    with open(args.log_file) as f:
+        for line in f:
+            data.append(json.loads(line))
+    singleton_analysis(data)
+
+
+if __name__ == "__main__":
+    main()

experiment.py

@@ -0,0 +1,1052 @@
+import sys
+import os
+import time
+import logging
+import torch
+
+## Uncomment the following line to make the code deterministic and use CUBLAS_WORKSPACE_CONFIG=:4096:8
+torch.use_deterministic_algorithms(True)
+import json
+import numpy as np
+import random
+import wandb
+
+from omegaconf import OmegaConf, open_dict
+from os import path
+from collections import OrderedDict, defaultdict
+from transformers import get_linear_schedule_with_warmup
+from transformers import AutoModel, AutoTokenizer
+
+from data_utils.utils import load_dataset, load_eval_dataset
+import pytorch_utils.utils as utils
+from torch.profiler import profile, record_function, ProfilerActivity
+
+from model.entity_ranking_model import EntityRankingModel
+from model.mention_proposal import MentionProposalModule
+from data_utils.tensorize_dataset import TensorizeDataset
+from pytorch_utils.optimization_utils import get_inverse_square_root_decay
+
+from utils_evaluate import coref_evaluation
+
+from typing import Dict, Union, List, Optional
+from omegaconf import DictConfig
+import copy
+
+logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)
+logger = logging.getLogger()
+
+loss_acc_template_dict = {
+    "total": 0.0,
+    "ment_loss": 0.0,
+    "coref": 0.0,
+    "mention_count": 0.0,
+    "processed_docs": 0.0,
+    "ment_correct": 0.0,
+    "ment_total": 0.0,
+    "ment_tp": 0.0,
+    "ment_pp": 0.0,
+    "ment_ap": 0.0,
+}
+
+
+class Experiment:
+    """Class for training and evaluating coreference models."""
+
+    def __init__(self, config: DictConfig):
+        self.config = config
+
+        print("Seeded: ", config.seed)
+        print("Cuda Available: ", torch.cuda.is_available())
+
+        # Whether to train or not
+        self.eval_model: bool = not self.config.train
+
+        # Initialize dictionary to track key training variables
+        self.train_info = {
+            "val_perf": 0.0,
+            "global_steps": 0,
+            "num_stuck_evals": 0,
+            "peak_memory": 0.0,
+        }
+
+        self.wandbdata = {}
+
+        # Initialize model path attributes
+        self.model_path = self.config.paths.model_path
+        self.best_model_path = self.config.paths.best_model_path
+
+        if not self.eval_model:
+            # Step 1 - Initialize model
+            self._build_model()
+            # Step 2 - Load Data - Data processing choices such as tokenizer will depend on the model
+            self._load_data()
+            # Step 3 - Resume training
+            self._setup_training()
+            # Step 4 - Loading the checkpoint also restores the training metadata
+            self._load_previous_checkpoint()
+
+            # All set to resume training
+            # But first check if training is remaining
+            if self._is_training_remaining():
+                self.train()
+
+        # Perform final evaluation
+        if path.exists(self.best_model_path):
+            # Step 1 - Initialize model
+            self._initialize_best_model()
+            # Step 2 - Load evaluation data
+            self._load_data()
+            # Step 3 - Perform evaluation
+            self.perform_final_eval()
+        else:
+            logger.info("No model accessible!")
+            sys.exit(1)
+
+    def _build_model(self) -> None:
+        """Constructs the model with given config."""
+
+        model_params: DictConfig = self.config.model
+        train_config: DictConfig = self.config.trainer
+
+        self.model = EntityRankingModel(
+            model_config=model_params, train_config=train_config
+        )
+
+        if torch.cuda.is_available():
+            self.model.cuda(device=self.config.device)
+
+        # Print model
+        utils.print_model_info(self.model)
+        sys.stdout.flush()
+
+    def _load_data(self):
+        """Loads and processes the training and evaluation data.
+
+        Loads the data concerning all the specified datasets for training and eval.
+        The first part of this method loads all the data from the preprocessed jsonline files.
+        In the second half, the loaded data is tensorized for consumption by the model.
+
+        Apart from loading and processing the data, the method also populates important
+        attributes such as:
+                num_train_docs_map (dict): Dictionary to maintain the number of training
+                        docs per dataset which is useful for implementing sampling in joint training.
+                num_training_steps (int): Number of total training steps.
+                eval_per_k_steps (int): Number of gradient updates before each evaluation.
+        """
+
+        self.data_iter_map, self.conll_data_dir, self.num_split_docs_map = (
+            {},
+            {},
+            {"train": {}, "dev": {}, "test": {}},
+        )
+        raw_data_map = {}
+
+        max_segment_len: int = self.config.model.doc_encoder.transformer.max_segment_len
+        model_name: str = self.config.model.doc_encoder.transformer.name
+        add_speaker_tokens: bool = self.config.model.doc_encoder.add_speaker_tokens
+        base_data_dir: str = path.abspath(self.config.paths.base_data_dir)
+
+        # Load data
+        for dataset_name, attributes in self.config.datasets.items():
+            num_train_docs: Optional[int] = attributes.get("num_train_docs", None)
+            num_dev_docs: Optional[int] = attributes.get("num_dev_docs", None)
+            num_test_docs: Optional[int] = attributes.get("num_test_docs", None)
+            singleton_file: Optional[str] = attributes.get("singleton_file", None)
+            external_md_file: Optional[str] = attributes.get("external_md_file", None)
+
+            if singleton_file is not None:
+                singleton_file = path.join(base_data_dir, singleton_file)
+                if path.exists(singleton_file):
+                    logger.info(f"Singleton file found: {singleton_file}")
+
+            if external_md_file is not None:
+                external_md_file = path.join(base_data_dir, external_md_file)
+                if path.exists(external_md_file):
+                    logger.info(
+                        f"External mention detector file found: {external_md_file}"
+                    )
+
+            # Data directory is a function of dataset name and tokenizer used
+            data_dir = path.join(path.join(base_data_dir, dataset_name), model_name)
+            # Check if speaker tokens are added
+            if add_speaker_tokens:
+                pot_data_dir = path.join(
+                    path.join(path.join(base_data_dir, dataset_name)),
+                    model_name + "_speaker",
+                )
+                if path.exists(pot_data_dir):
+                    data_dir = pot_data_dir
+
+            # Datasets such as litbank have cross validation splits
+            if attributes.get("cross_val_split", None) is not None:
+                data_dir = path.join(data_dir, str(attributes.get("cross_val_split")))
+
+            logger.info("Data directory: %s" % data_dir)
+
+            # CoNLL data dir
+            if attributes.get("has_conll", False):
+                conll_dir = path.join(
+                    path.join(path.join(base_data_dir, dataset_name)), "conll"
+                )
+                if attributes.get("cross_val_split", None) is not None:
+                    # LitBank like datasets have cross validation splits
+                    conll_dir = path.join(
+                        conll_dir, str(attributes.get("cross_val_split"))
+                    )
+
+                if path.exists(conll_dir):
+                    self.conll_data_dir[dataset_name] = conll_dir
+
+            self.num_split_docs_map["train"][dataset_name] = num_train_docs
+            self.num_split_docs_map["dev"][dataset_name] = num_dev_docs
+            self.num_split_docs_map["test"][dataset_name] = num_test_docs
+
+            if self.eval_model:
+                print("In Eval Model DataLoader")
+                raw_data_map[dataset_name] = load_eval_dataset(
+                    data_dir,
+                    external_md_file=external_md_file,
+                    max_segment_len=max_segment_len,
+                    dataset_name=dataset_name,
+                )
+            else:
+                raw_data_map[dataset_name] = load_dataset(
+                    data_dir,
+                    singleton_file=singleton_file,
+                    num_dev_docs=num_dev_docs,
+                    num_test_docs=num_test_docs,
+                    max_segment_len=max_segment_len,
+                    dataset_name=dataset_name,
+                )
+
+        # Tensorize data
+        data_processor = TensorizeDataset(
+            self.model.get_tokenizer(),
+            remove_singletons=(not self.config.keep_singletons),
+        )
+
+        if self.eval_model:
+            for split in ["dev", "test"]:
+                self.data_iter_map[split] = {}
+
+            for dataset in raw_data_map:
+                for split in raw_data_map[dataset]:
+                    self.data_iter_map[split][dataset] = data_processor.tensorize_data(
+                        raw_data_map[dataset][split], training=False
+                    )
+        else:
+            # Training
+            for split in ["train", "dev", "test"]:
+                self.data_iter_map[split] = {}
+                training = split == "train"
+                for dataset in raw_data_map:
+                    self.data_iter_map[split][dataset] = data_processor.tensorize_data(
+                        raw_data_map[dataset][split], training=training
+                    )
+
+            # Estimate number of training steps
+            if self.config.trainer.eval_per_k_steps is None:
+                # Eval steps is 1 epoch (with subsampling) of all the datasets used in joint training
+                self.config.trainer.eval_per_k_steps = sum(
+                    self.num_split_docs_map["train"].values()
+                )
+
+            self.config.trainer.num_training_steps = (
+                self.config.trainer.eval_per_k_steps * self.config.trainer.max_evals
+            )
+            logger.info(
+                f"Number of training steps: {self.config.trainer.num_training_steps}"
+            )
+
+            logger.info(f"Eval per k steps: {self.config.trainer.eval_per_k_steps}")
+
+    def _load_previous_checkpoint(self):
+        """Loads the last checkpoint or best checkpoint."""
+
+        # Resume training
+        print("Model Path: ", self.model_path)
+        print("Model Initialised:", torch.cuda.memory_summary(self.config.device))
+        if path.exists(self.model_path):
+            self.load_model(self.model_path, last_checkpoint=True)
+            logger.info("Model loaded\n")
+            print(
+                "Loaded Model Returned:", torch.cuda.memory_summary(self.config.device)
+            )
+        else:
+            # Starting training
+            logger.info("Model initialized\n")
+            sys.stdout.flush()
+
+    def _is_training_remaining(self):
+        """Check if training is done or remaining.
+
+        There are two cases where we don't resume training:
+        (a) The dev performance has not improved for the allowed patience parameter number of evaluations.
+        (b) Number of gradient updates is already >= Total training steps.
+
+        Returns:
+                bool: If true, we resume training. Otherwise do final evaluation.
+        """
+
+        if self.train_info["num_stuck_evals"] >= self.config.trainer.patience:
+            return False
+        if self.train_info["global_steps"] >= self.config.trainer.num_training_steps:
+            return False
+
+        return True
+
+    def _setup_training(self):
+        """Initialize optimizer and bookkeeping variables for training."""
+
+        # Dictionary to track key training variables
+        self.train_info = {
+            "val_perf": 0.0,
+            "global_steps": 0,
+            "num_stuck_evals": 0,
+            "peak_memory": 0.0,
+            "max_mem": 0.0,
+        }
+
+        # Initialize optimizers
+        self._initialize_optimizers()
+
+    def _initialize_optimizers(self):
+        """Initialize model + optimizer(s). Check if there's a checkpoint in which case we resume from there."""
+
+        optimizer_config: DictConfig = self.config.optimizer
+        train_config: DictConfig = self.config.trainer
+        self.optimizer, self.optim_scheduler = {}, {}
+
+        if torch.cuda.is_available():
+            # Gradient scaler required for mixed precision training
+            self.scaler = torch.GradScaler("cuda")
+        else:
+            self.scaler = None
+
+        # Optimizer for clustering params
+        self.optimizer["mem"] = torch.optim.Adam(
+            self.model.get_params()[1], lr=optimizer_config.init_lr, eps=1e-6
+        )
+
+        if optimizer_config.lr_decay == "inv":
+            self.optim_scheduler["mem"] = get_inverse_square_root_decay(
+                self.optimizer["mem"], num_warmup_steps=0
+            )
+        else:
+            # No warmup steps for model params
+            self.optim_scheduler["mem"] = get_linear_schedule_with_warmup(
+                self.optimizer["mem"],
+                num_warmup_steps=0,
+                num_training_steps=train_config.num_training_steps,
+            )
+
+        if self.config.model.doc_encoder.finetune:
+            # Optimizer for document encoder
+            no_decay = [
+                "bias",
+                "LayerNorm.weight",
+            ]  # No weight decay for bias and layernorm weights
+            encoder_params = self.model.get_params(named=True)[0]
+            grouped_param = [
+                {
+                    "params": [
+                        p
+                        for n, p in encoder_params
+                        if not any(nd in n for nd in no_decay)
+                    ],
+                    "lr": optimizer_config.fine_tune_lr,
+                    "weight_decay": 1e-2,
+                },
+                {
+                    "params": [
+                        p for n, p in encoder_params if any(nd in n for nd in no_decay)
+                    ],
+                    "lr": optimizer_config.fine_tune_lr,
+                    "weight_decay": 0.0,
+                },
+            ]
+
+            self.optimizer["doc"] = torch.optim.AdamW(
+                grouped_param, lr=optimizer_config.fine_tune_lr, eps=1e-6
+            )
+
+            # Scheduler for document encoder
+            num_warmup_steps = int(0.1 * train_config.num_training_steps)
+            if optimizer_config.lr_decay == "inv":
+                self.optim_scheduler["doc"] = get_inverse_square_root_decay(
+                    self.optimizer["doc"], num_warmup_steps=num_warmup_steps
+                )
+            else:
+                self.optim_scheduler["doc"] = get_linear_schedule_with_warmup(
+                    self.optimizer["doc"],
+                    num_warmup_steps=num_warmup_steps,
+                    num_training_steps=train_config.num_training_steps,
+                )
+
+    def agg(self, datadepdict):
+        agg_dict = defaultdict(float)
+        for dataset in datadepdict:
+            for key in datadepdict[dataset]:
+                agg_dict[key] += datadepdict[dataset][key]
+
+        agg_dict["loss_norm"] = (
+            agg_dict["coref"] / agg_dict["mention_count"]
+            + agg_dict["ment_loss"] / agg_dict["ment_total"]
+            if agg_dict["mention_count"] > 0
+            else 0
+        )
+        agg_dict["ment_acc"] = agg_dict["ment_correct"] / agg_dict["ment_total"]
+        agg_dict["ment_prec"] = (
+            agg_dict["ment_tp"] / agg_dict["ment_pp"] if agg_dict["ment_pp"] > 0 else 0
+        )
+        agg_dict["ment_rec"] = (
+            agg_dict["ment_tp"] / agg_dict["ment_ap"] if agg_dict["ment_ap"] > 0 else 0
+        )
+        agg_dict["ment_f1"] = (
+            2
+            * (agg_dict["ment_prec"] * agg_dict["ment_rec"])
+            / (agg_dict["ment_prec"] + agg_dict["ment_rec"])
+            if (agg_dict["ment_prec"] + agg_dict["ment_rec"]) > 0
+            else 0
+        )
+
+        return agg_dict
+
+    def train(self) -> None:
+        """Method for training the model.
+
+        This method implements the training loop.
+        Within the training loop, the model is periodically evaluated on the dev set(s).
+        """
+
+        model, optimizer, scheduler, scaler = (
+            self.model,
+            self.optimizer,
+            self.optim_scheduler,
+            self.scaler,
+        )
+        model.train()
+
+        optimizer_config, train_config = self.config.optimizer, self.config.trainer
+
+        start_time = time.time()
+        eval_time = {"total_time": 0, "num_evals": 0}
+        print("Started Training..")
+        while True:
+            logger.info("Steps done %d" % (self.train_info["global_steps"]))
+
+            train_data = self.runtime_load_dataset("train")
+            np.random.shuffle(train_data)
+            logger.info("Per epoch training steps: %d" % len(train_data))
+            logger.info("Per epoch training steps: %d" % len(train_data))
+
+            encoder_params, task_params = model.get_params()
+            stat_per_dataset = defaultdict(
+                lambda: copy.deepcopy(loss_acc_template_dict)
+            )
+            agg_stat = self.agg
+
+            # Training "epoch" -> May not correspond to actual epoch
+            for cur_document in train_data:
+
+                def handle_example(document: Dict) -> Union[None, float]:
+                    self.train_info["global_steps"] += 1
+                    for key in optimizer:
+                        optimizer[key].zero_grad()
+                    loss_dict: Dict = model.forward_training(document)
+
+                    total_loss = loss_dict["total"]
+
+                    if total_loss is None or torch.isnan(total_loss):
+                        print("Problem with Loss. Should not occur often")
+                        return None
+
+                    total_loss.backward()
+
+                    # Gradient clipping
+                    try:
+                        for name_ind, param_group in enumerate(
+                            [encoder_params, task_params]
+                        ):
+                            torch.nn.utils.clip_grad_norm_(
+                                param_group,
+                                optimizer_config.max_gradient_norm,
+                                error_if_nonfinite=True,
+                            )
+                    except RuntimeError:
+                        print("Non Finite Gradient")
+                        return None
+
+                    for key in optimizer:
+                        self.wandbdata[key + "_lr"] = scheduler[key].get_last_lr()[0]
+
+                    for key in optimizer:
+                        optimizer[key].step()
+                        scheduler[key].step()
+
+                    loss_dict_items = {}
+                    for key in loss_dict:
+                        loss_dict_items[key] = loss_dict[key].item()
+
+                    dataset_name = document["dataset_name"]
+                    # print(f"Total loss {cur_document['doc_key']}: {total_loss.item()}")
+
+                    for key in loss_dict_items:
+                        stat_per_dataset[dataset_name][key] += loss_dict_items[key]
+
+                    stat_per_dataset[dataset_name]["processed_docs"] += 1
+
+                    return total_loss.item()
+
+                loss = handle_example(cur_document)
+
+                if self.train_info["global_steps"] % train_config.log_frequency == 0:
+                    max_mem = (
+                        (
+                            torch.cuda.max_memory_allocated(self.config.device)
+                            / (1024**3)
+                        )
+                        if torch.cuda.is_available()
+                        else 0.0
+                    )
+                    if self.train_info.get("max_mem", 0.0) < max_mem:
+                        self.train_info["max_mem"] = max_mem
+
+                    if loss is not None:
+                        logger.info(
+                            "{} {:.3f} Max mem {:.1f} GB".format(
+                                cur_document["doc_key"],
+                                loss,
+                                max_mem,
+                            )
+                        )
+                    sys.stdout.flush()
+                    if torch.cuda.is_available():
+                        torch.cuda.reset_peak_memory_stats()
+
+                if train_config.eval_per_k_steps and (
+                    self.train_info["global_steps"] % train_config.eval_per_k_steps == 0
+                ):
+                    print("Eval needs to be done here")
+                    coref_dict = {}
+                    print(stat_per_dataset)
+                    if self.config.use_wandb:
+                        self._wandb_log(
+                            split="train",
+                            stat_per_dataset=stat_per_dataset,
+                            agg_stat=agg_stat,
+                            coref_dict=coref_dict,
+                            step=self.train_info["global_steps"]
+                            // train_config.eval_per_k_steps,
+                        )
+
+                    stat_per_dataset = defaultdict(
+                        lambda: copy.deepcopy(loss_acc_template_dict)
+                    )
+
+                    macro_fscore = self.periodic_model_eval()
+
+                    model.train()
+                    # Get elapsed time
+                    elapsed_time = time.time() - start_time
+
+                    start_time = time.time()
+                    logger.info(
+                        "Steps: %d, Micro F1: %.1f, Max Micro F1: %.1f, Time: %.2f"
+                        % (
+                            self.train_info["global_steps"],
+                            macro_fscore,
+                            self.train_info["val_perf"],
+                            elapsed_time,
+                        )
+                    )
+
+                    # Check stopping criteria
+                    if not self._is_training_remaining():
+                        break
+
+            # Check stopping criteria
+            if not self._is_training_remaining():
+                break
+
+            logger.handlers[0].flush()
+
+    def runtime_load_dataset(self, split):
+        # Shuffle and load the training data
+        data = []
+        for dataset, dataset_data in self.data_iter_map[split].items():
+            np.random.shuffle(
+                dataset_data
+            )  ### Commenting this so that we can have a deterministic training
+            if self.num_split_docs_map[split].get(dataset, None) is not None:
+                # Subsampling the data - This is useful in joint training
+                logger.info(
+                    f"{dataset}: Subsampled {self.num_split_docs_map[split].get(dataset)}"
+                )
+                random_indices = range(self.num_split_docs_map[split].get(dataset))
+                data += [dataset_data[idx] for idx in random_indices]
+            else:
+                data += dataset_data
+        return data
+
+    def _wandb_log(self, split, stat_per_dataset, agg_stat, coref_dict, step=None):
+        for dataset_name in stat_per_dataset:
+            for metric_vals in stat_per_dataset[dataset_name]:
+                wandb.log(
+                    data={
+                        f"{split}/{dataset_name}/{metric_vals}": stat_per_dataset[
+                            dataset_name
+                        ][metric_vals]
+                    },
+                    step=step,
+                )
+            if stat_per_dataset[dataset_name]["mention_count"] > 0.0:
+                ment_prec = (
+                    stat_per_dataset[dataset_name]["ment_tp"]
+                    / stat_per_dataset[dataset_name]["ment_pp"]
+                    if stat_per_dataset[dataset_name]["ment_pp"] > 0
+                    else 0
+                )
+                ment_rec = (
+                    stat_per_dataset[dataset_name]["ment_tp"]
+                    / stat_per_dataset[dataset_name]["ment_ap"]
+                    if stat_per_dataset[dataset_name]["ment_ap"] > 0
+                    else 0
+                )
+                ment_f1 = (
+                    2 * (ment_prec * ment_rec) / (ment_prec + ment_rec)
+                    if (ment_prec + ment_rec) > 0
+                    else 0
+                )
+                wandb.log(
+                    data={
+                        f"{split}/{dataset_name}/loss_norm": stat_per_dataset[
+                            dataset_name
+                        ]["coref"]
+                        / stat_per_dataset[dataset_name]["mention_count"]
+                        + stat_per_dataset[dataset_name]["ment_loss"]
+                        / stat_per_dataset[dataset_name]["ment_total"],
+                        f"{split}/{dataset_name}/ment_acc": stat_per_dataset[
+                            dataset_name
+                        ]["ment_correct"]
+                        / stat_per_dataset[dataset_name]["ment_total"],
+                        f"{split}/{dataset_name}/ment_prec": ment_prec,
+                        f"{split}/{dataset_name}/ment_rec": ment_rec,
+                        f"{split}/{dataset_name}/ment_f1": ment_f1,
+                    },
+                    step=step,
+                )
+            else:
+                print("No mentions processed. Should not occur many times.")
+
+        if agg_stat:
+            for metric in agg_stat(stat_per_dataset):
+                wandb.log(
+                    data={f"{split}/{metric}": agg_stat(stat_per_dataset)[metric]},
+                    step=step,
+                )
+
+        for dataset in coref_dict:
+            for key in coref_dict[dataset]:
+                # Log result for individual metrics
+                if isinstance(coref_dict[dataset][key], dict):
+                    wandb.log(
+                        data={
+                            f"{split}/{dataset}/{key}": coref_dict[dataset][key].get(
+                                "fscore", 0.0
+                            )
+                        },
+                        step=step,
+                    )
+
+            # Log the overall F-score
+            wandb.log(
+                data={
+                    f"{split}/{dataset}/CoNLL": coref_dict[dataset].get("fscore", 0.0)
+                },
+                step=step,
+            )
+
+            wandb.log(
+                data={
+                    f"{split}/{dataset}/Micro-F1": coref_dict[dataset].get(
+                        "f1_micro", 0.0
+                    )
+                },
+                step=step,
+            )
+
+            wandb.log(
+                data={
+                    f"{split}/{dataset}/Macro-F1": coref_dict[dataset].get(
+                        "f1_macro", 0.0
+                    )
+                },
+                step=step,
+            )
+
+        wandb.log(data=self.wandbdata, step=step)
+
+    @torch.no_grad()
+    def periodic_model_eval(self) -> float:
+        """Method for evaluating and saving the model during the training loop.
+
+        Returns:
+                float: Average CoNLL F-score over all the development sets of datasets.
+        """
+
+        self.model.eval()
+
+        ## Dev Loss Calculations:
+        dev_data = self.runtime_load_dataset("dev")
+        np.random.shuffle(dev_data)
+        stat_per_dataset = defaultdict(lambda: copy.deepcopy(loss_acc_template_dict))
+        agg_stat = self.agg
+
+        for cur_document in dev_data:
+
+            def handle_example(document: Dict) -> Union[None, float]:
+                loss_dict: Dict = self.model.forward_training(document)
+                total_loss = loss_dict["total"]
+                if total_loss is None or torch.isnan(total_loss):
+                    print("Problem with Loss. Should not occur many times")
+                    return None
+
+                loss_dict_items = {}
+                for key in loss_dict:
+                    loss_dict_items[key] = loss_dict[key].item()
+
+                dataset_name = document["dataset_name"]
+
+                for key in loss_dict_items:
+                    stat_per_dataset[dataset_name][key] += loss_dict_items[key]
+
+                stat_per_dataset[dataset_name]["processed_docs"] += 1
+                return total_loss.item()
+
+            loss = handle_example(cur_document)
+            if loss is None:
+                continue
+
+        # Dev performance
+        coref_dict = {}
+        train_config = self.config.trainer
+        for dataset in self.data_iter_map["dev"]:
+            for go in [False]:
+                for tf in [False]:
+                    result_dict = coref_evaluation(
+                        self.config,
+                        self.model,
+                        self.data_iter_map,
+                        dataset,
+                        teacher_force=tf,
+                        gold_mentions=go,
+                        _iter="_"
+                        + str(
+                            self.train_info["global_steps"]
+                            // train_config.eval_per_k_steps
+                        ),
+                        conll_data_dir=self.conll_data_dir,
+                    )
+
+                    coref_dict[dataset] = result_dict
+
+        if self.config.use_wandb:
+            self._wandb_log(
+                split="dev",
+                stat_per_dataset=stat_per_dataset,
+                agg_stat=agg_stat,
+                coref_dict=coref_dict,
+                step=self.train_info["global_steps"] // train_config.eval_per_k_steps,
+            )
+
+        # Calculate Mean F-score
+        fscore = sum([coref_dict[dataset]["fscore"] for dataset in coref_dict]) / len(
+            coref_dict
+        )
+        micro_fscore = sum(
+            [coref_dict[dataset]["f1_micro"] for dataset in coref_dict]
+        ) / len(coref_dict)
+        macro_fscore = sum(
+            [coref_dict[dataset]["f1_macro"] for dataset in coref_dict]
+        ) / len(coref_dict)
+
+        logger.info(
+            "Avg Macro F1: %.1f, Max Micro F1: %.1f"
+            % (macro_fscore, self.train_info["val_perf"])
+        )
+        logger.info("Avg Macro F1: %.1f" % (macro_fscore))
+
+        # Update model if dev performance improves
+        if macro_fscore > self.train_info["val_perf"]:
+            # Update training bookkeeping variables
+            self.train_info["num_stuck_evals"] = 0
+            self.train_info["val_perf"] = macro_fscore
+
+            # Save the best model
+            logger.info("Saving best model")
+            self.save_model(self.best_model_path, last_checkpoint=False)
+        else:
+            self.train_info["num_stuck_evals"] += 1
+
+        # Save model
+        if self.config.trainer.to_save_model:
+            self.save_model(self.model_path, last_checkpoint=True)
+
+        # Go back to training mode
+        self.model.train()
+        return macro_fscore
+
+    @torch.no_grad()
+    def perform_final_eval(self) -> None:
+        """Method to evaluate the model after training has finished."""
+
+        self.model.eval()
+        base_output_dict = OmegaConf.to_container(self.config)
+        perf_summary = {"best_perf": self.train_info["val_perf"]}
+        if self.config.paths.model_dir:
+            perf_summary["model_dir"] = path.normpath(self.config.paths.model_dir)
+
+        logger.info(
+            "Max training memory: %.1f GB" % self.train_info.get("max_mem", 0.0)
+        )
+
+        logger.info("Validation performance: %.1f" % self.train_info["val_perf"])
+
+        perf_file_dict = {}
+        dataset_output_dict = {}
+
+        for split in ["dev", "test"]:
+            perf_summary[split] = {}
+            logger.info("\n")
+            logger.info("%s" % split.capitalize())
+            coref_dict = {}
+            for dataset in self.data_iter_map.get(split, []):
+                dataset_dir = path.join(self.config.paths.model_dir, dataset)
+                if not path.exists(dataset_dir):
+                    os.makedirs(dataset_dir)
+
+                if dataset not in dataset_output_dict:
+                    dataset_output_dict[dataset] = {}
+                if dataset not in perf_file_dict:
+                    perf_file_dict[dataset] = path.join(dataset_dir, f"perf.json")
+
+                print("Dataset Name:", self.config.datasets[dataset].name)
+                logger.info("Dataset: %s\n" % self.config.datasets[dataset].name)
+
+                for go in [False]:
+                    for tf in [False]:
+                        result_dict = coref_evaluation(
+                            self.config,
+                            self.model,
+                            self.data_iter_map,
+                            dataset=dataset,
+                            split=split,
+                            teacher_force=tf,
+                            gold_mentions=go,
+                            final_eval=True,
+                            conll_data_dir=self.conll_data_dir,
+                        )
+                        coref_dict[dataset] = result_dict
+                        dataset_output_dict[dataset][split] = result_dict
+                        perf_summary[split][dataset] = result_dict["f1_micro"]
+
+            if self.config.use_wandb:
+                self._wandb_log(
+                    split=split,
+                    stat_per_dataset={},
+                    agg_stat=None,
+                    coref_dict=coref_dict,
+                    step=None,
+                )
+
+            sys.stdout.flush()
+
+        for dataset, output_dict in dataset_output_dict.items():
+            perf_file = perf_file_dict[dataset]
+            json.dump(output_dict, open(perf_file, "w"), indent=2)
+            logger.info("Final performance summary at %s" % path.abspath(perf_file))
+
+        summary_file = path.join(self.config.paths.model_dir, "perf.json")
+        json.dump(perf_summary, open(summary_file, "w"), indent=2)
+        logger.info("Performance summary file: %s" % path.abspath(summary_file))
+
+    def _initialize_best_model(self):
+        checkpoint = torch.load(
+            self.best_model_path,
+            map_location="cpu",
+        )
+        config = checkpoint["config"]
+
+        ## Due to version changes -- these changes are necessary
+        # if
+
+        if self.config.get("override_encoder", False):
+            model_config = config.model
+            print(type(self.config.model.doc_encoder.transformer))
+            print(self.config.model.doc_encoder.transformer)
+            model_config.doc_encoder.transformer = (
+                self.config.model.doc_encoder.transformer
+            )
+
+        # Override memory
+        # For e.g., can test with a different bounded memory size
+        if self.config.get("override_memory", False):
+            model_config = config.model
+            model_config.memory = self.config.model.memory
+
+        with open_dict(config):
+            print("Config change")
+            config.model.mention_params.ext_ment = (
+                self.config.model.mention_params.ext_ment
+            )
+            config = utils.fill_missing_configs(config, self.config)
+            print("Type: ", config.model.memory.type)
+
+        self.config.model = config.model
+
+        self.train_info = checkpoint["train_info"]
+
+        if self.config.model.doc_encoder.finetune:
+            # Load the document encoder params if encoder is finetuned
+            doc_encoder_dir = path.join(
+                path.dirname(self.best_model_path),
+                self.config.paths.doc_encoder_dirname,
+            )
+            if path.exists(doc_encoder_dir):
+                logger.info(
+                    "Loading document encoder from %s" % path.abspath(doc_encoder_dir)
+                )
+                config.model.doc_encoder.transformer.model_str = doc_encoder_dir
+
+        self.model = EntityRankingModel(config.model, config.trainer)
+
+        # Document encoder parameters will be loaded via the huggingface initialization
+        self.model.load_state_dict(checkpoint["model"], strict=False)
+
+        if torch.cuda.is_available():
+            self.model.cuda(device=self.config.device)
+
+    def load_model(self, location: str, last_checkpoint=True) -> None:
+        """Load model from given location.
+
+        Args:
+                location: str
+                        Location of checkpoint
+                last_checkpoint: bool
+                        Whether the checkpoint is the last one saved or not.
+                        If false, don't load optimizers, schedulers, and other training variables.
+        """
+
+        checkpoint = torch.load(location, map_location="cpu")
+        logger.info("Loading model from %s" % path.abspath(location))
+
+        # self.config = checkpoint["config"] ## Commented out so that it does not load the config of the trained model. Removed comment
+
+        self.model.load_state_dict(
+            checkpoint["model"], strict=False
+        )  ## No encoder in this model so strict=False is compulsary. No other weight missing. Checked
+
+        # self.train_info = checkpoint["train_info"] ## No train info transfer too. ## Transferring
+
+        if self.config.model.doc_encoder.finetune:
+            # Load the document encoder params if encoder is finetuned
+            doc_encoder_dir = path.join(
+                path.dirname(location), self.config.paths.doc_encoder_dirname
+            )
+            logger.info(
+                "Loading document encoder from %s" % path.abspath(doc_encoder_dir)
+            )
+
+            # Load the encoder
+            self.model.mention_proposer.doc_encoder.lm_encoder = (
+                AutoModel.from_pretrained(pretrained_model_name_or_path=doc_encoder_dir)
+            )
+            self.model.mention_proposer.doc_encoder.tokenizer = (
+                AutoTokenizer.from_pretrained(
+                    pretrained_model_name_or_path=doc_encoder_dir,
+                    clean_up_tokenization_spaces=True,
+                )
+            )
+            if self.model.mention_proposer.doc_encoder.config.finetune:
+                self.model.mention_proposer.doc_encoder.lm_encoder.gradient_checkpointing_enable()
+
+            if torch.cuda.is_available():
+                self.model.cuda(device=self.config.device)
+
+        print("Loaded Model:", torch.cuda.memory_summary())
+        print(
+            "Gradient checkpointing enabled? ", torch.autograd.grad_checkpoint_enabled()
+        )
+        del checkpoint
+        torch.cuda.empty_cache()
+
+    def save_model(self, location: os.PathLike, last_checkpoint=True) -> None:
+        """Save model.
+
+        Args:
+                location: Location of checkpoint
+                last_checkpoint:
+                        Whether the checkpoint is the last one saved or not.
+                        If false, don't save optimizers and schedulers which take up a lot of space.
+        """
+
+        model_state_dict = OrderedDict(self.model.state_dict())
+        doc_encoder_state_dict = {}
+
+        # Separate the doc_encoder state dict
+        # We will save the model in two parts:
+        # (a) Doc encoder parameters - Useful for final upload to huggingface
+        # (b) Rest of the model parameters, optimizers, schedulers, and other bookkeeping variables
+        for key in self.model.state_dict():
+            if "lm_encoder." in key:
+                doc_encoder_state_dict[key] = model_state_dict[key]
+                del model_state_dict[key]
+
+        # Save the document encoder params
+        if self.config.model.doc_encoder.finetune:
+            doc_encoder_dir = path.join(
+                path.dirname(location), self.config.paths.doc_encoder_dirname
+            )
+            if not path.exists(doc_encoder_dir):
+                os.makedirs(doc_encoder_dir)
+
+            logger.info(f"Encoder saved at {path.abspath(doc_encoder_dir)}")
+            # Save the encoder
+            self.model.mention_proposer.doc_encoder.lm_encoder.save_pretrained(
+                save_directory=doc_encoder_dir, save_config=True
+            )
+            # Save the tokenizer
+            self.model.mention_proposer.doc_encoder.tokenizer.save_pretrained(
+                doc_encoder_dir
+            )
+
+        save_dict = {
+            "train_info": self.train_info,
+            "model": model_state_dict,
+            "rng_state": torch.get_rng_state(),
+            "np_rng_state": np.random.get_state(),
+            "config": self.config,
+        }
+
+        if self.scaler is not None:
+            save_dict["scaler"] = self.scaler.state_dict()
+
+        if last_checkpoint:
+            # For last checkpoint save the optimizer and scheduler states as well
+            save_dict["optimizer"] = {}
+            save_dict["scheduler"] = {}
+
+            param_groups: List[str] = (
+                ["mem", "doc"] if self.config.model.doc_encoder.finetune else ["mem"]
+            )
+            for param_group in param_groups:
+                save_dict["optimizer"][param_group] = self.optimizer[
+                    param_group
+                ].state_dict()
+                save_dict["scheduler"][param_group] = self.optim_scheduler[
+                    param_group
+                ].state_dict()
+
+        torch.save(save_dict, location)
+        logger.info(f"Model saved at: {path.abspath(location)}")