Repository: informagi/REL (language: python)

Issue:
Re-implement server using fastapi/pydantic
Hi all, this PR changes the server to use the `fastapi`/`pydantic` framework. See #146 for more info.

Closes #146
Re-implement server using fastapi/pydantic
Hi all, this PR changes the server to use the `fastapi`/`pydantic` framework. See #146 for more info.

Closes #146
Mention detection with Bert
New version of REL with the opportunity to use Bert instead of Flair for mention detection (named entity recognition)

## Installation

1. `python3 -m venv --prompt REL venv3`
1. `source venv3/bin/activate`
1. `git clone https://github.com/informagi/REL.git`
1. `cd REL`
1. `git checkout md-with-bert-2`
1. `pip install -e '.[develop]'`
1. `mkdir src/data`
1. `cd src/data`
1. `curl -O http://gem.cs.ru.nl/generic.tar.gz`
1. `curl -O http://gem.cs.ru.nl/ed-wiki-2019.tar.gz`
1. `curl -O http://gem.cs.ru.nl/wiki_2019.tar.gz`
1. `tar zxf wiki_2019.tar.gz`
1. `tar zxf ed-wiki-2019.tar.gz`
1. `tar zxf generic.tar.gz`
1. `cd ../..`

## Testing

Run 13 tests

* `pytest tests`

Some warnings might be reported but all tests should succeed

List test options:

* `python3 src/scripts/efficiency_test.py -h`

Process one test document with Flair, sentence-by-sentence, and report the performance:

* `python3 src/scripts/efficiency_test.py --tagger_ner_name flair --max_docs 1 --process_sentences`<br>...<br>Results: PMD RMD FMD PEL REL FEL: 94.9% 63.8% 76.3% | 69.2% 46.6% 55.7%

Process all 50 test documents with uncased Bert base, document-by-document, maximum of 500 tokens, and report the performance:

* `python3 src/scripts/efficiency_test.py --tagger_ner_name bert_base_uncased --split_docs_value 500`<br>...<br>Results: PMD RMD FMD PEL REL FEL: 93.5% 62.4% 74.8% | 62.9% 42.0% 50.3%

Use the server to process one test document with Flair, sentence-by-sentence, and report the performance:

* `python3 src/REL/server.py`
* `python3 src/scripts/efficiency_test.py --max_docs 1 --use_server`<br>...<br>Results: PMD RMD FMD PEL REL FEL: 95.0% 63.3% 76.0% | 67.5% 45.0% 54.0%

Mention detection with Bert
New version of REL with the opportunity to use Bert instead of Flair for mention detection (named entity recognition)

## Installation

1. `python3 -m venv --prompt REL venv3`
1. `source venv3/bin/activate`
1. `git clone https://github.com/informagi/REL.git`
1. `cd REL`
1. `git checkout md-with-bert-2`
1. `pip install -e '.[develop]'`
1. `mkdir src/data`
1. `cd src/data`
1. `curl -O http://gem.cs.ru.nl/generic.tar.gz`
1. `curl -O http://gem.cs.ru.nl/ed-wiki-2019.tar.gz`
1. `curl -O http://gem.cs.ru.nl/wiki_2019.tar.gz`
1. `tar zxf wiki_2019.tar.gz`
1. `tar zxf ed-wiki-2019.tar.gz`
1. `tar zxf generic.tar.gz`
1. `cd ../..`

## Testing

Run 13 tests

* `pytest tests`

Some warnings might be reported but all tests should succeed

List test options:

* `python3 src/scripts/efficiency_test.py -h`

Process one test document with Flair, sentence-by-sentence, and report the performance:

* `python3 src/scripts/efficiency_test.py --tagger_ner_name flair --max_docs 1 --process_sentences`<br>...<br>Results: PMD RMD FMD PEL REL FEL: 94.9% 63.8% 76.3% | 69.2% 46.6% 55.7%

Process all 50 test documents with uncased Bert base, document-by-document, maximum of 500 tokens, and report the performance:

* `python3 src/scripts/efficiency_test.py --tagger_ner_name bert_base_uncased --split_docs_value 500`<br>...<br>Results: PMD RMD FMD PEL REL FEL: 93.5% 62.4% 74.8% | 62.9% 42.0% 50.3%

Use the server to process one test document with Flair, sentence-by-sentence, and report the performance:

* `python3 src/REL/server.py`
* `python3 src/scripts/efficiency_test.py --max_docs 1 --use_server`<br>...<br>Results: PMD RMD FMD PEL REL FEL: 95.0% 63.3% 76.0% | 67.5% 45.0% 54.0%



Discussion / hints:

Fix patch:
diff --git a/docs/tutorials/conversations.md b/docs/tutorials/conversations.md
new file mode 100644
index 0000000..afaa276
--- /dev/null
+++ b/docs/tutorials/conversations.md
@@ -0,0 +1,71 @@
+# Conversational entity linking
+
+The `crel` submodule the conversational entity linking tool trained on the [ConEL-2 dataset](https://github.com/informagi/conversational-entity-linking-2022#conel-2-conversational-entity-linking-dataset).
+
+Unlike existing EL methods, `crel` is developed to identify both named entities and concepts.
+It also uses coreference resolution techniques to identify personal entities and references to the explicit entity mentions in the conversations.
+
+This tutorial describes how to start with conversational entity linking on a local machine.
+
+For more information, see the original [repository on conversational entity linking](https://github.com/informagi/conversational-entity-linking-2022).
+
+## Start with your local environment
+
+### Step 1: Download models
+
+First, download the models below:
+
+- **MD for concepts and NEs**: 
+	+ [Click here to download models](https://drive.google.com/file/d/1OoC2XZp4uBy0eB_EIuIhEHdcLEry2LtU/view?usp=sharing)
+	+ Extract `bert_conv-td` to your `base_url`
+- **Personal Entity Linking**:
+	+ [Click here to download models](https://drive.google.com/file/d/1-jW8xkxh5GV-OuUBfMeT2Tk7tEzvH181/view?usp=sharing)
+	+ Extract `s2e_ast_onto` to your `base_url`
+
+Additionally, conversational entity linking uses the wiki 2019 dataset. For more information on where to place the data and the `base_url`, check out [this page](../how_to_get_started). If setup correctly, your `base_url` should contain these directories:
+
+
+```bash
+.
+└── bert_conv-td
+└── s2e_ast_onto
+└── wiki_2019
+```
+
+
+### Step 2: Example code
+
+This example shows how to link a short conversation. Note that the speakers must be named "USER" and "SPEAKER".
+
+
+```python
+>>> from REL.crel.conv_el import ConvEL
+>>> 
+>>> cel = ConvEL(base_url="C:/path/to/base_url/")
+>>> 
+>>> conversation = [
+>>>     {"speaker": "USER", 
+>>>     "utterance": "I am allergic to tomatoes but we have a lot of famous Italian restaurants here in London.",}, 
+>>> 
+>>>     {"speaker": "SYSTEM", 
+>>>     "utterance": "Some people are allergic to histamine in tomatoes.",},
+>>> 
+>>>     {"speaker": "USER", 
+>>>     "utterance": "Talking of food, can you recommend me a restaurant in my city for our anniversary?",},
+>>> ]
+>>> 
+>>> annotated = cel.annotate(conversation)
+>>> [item for item in annotated if item['speaker'] == 'USER']
+[{'speaker': 'USER',
+  'utterance': 'I am allergic to tomatoes but we have a lot of famous Italian restaurants here in London.',
+  'annotations': [[17, 8, 'tomatoes', 'Tomato'],
+   [54, 19, 'Italian restaurants', 'Italian_cuisine'],
+   [82, 6, 'London', 'London']]},
+ {'speaker': 'USER',
+  'utterance': 'Talking of food, can you recommend me a restaurant in my city for our anniversary?',
+  'annotations': [[11, 4, 'food', 'Food'],
+   [40, 10, 'restaurant', 'Restaurant'],
+   [54, 7, 'my city', 'London']]}]
+
+```
+
diff --git a/docs/tutorials/index.md b/docs/tutorials/index.md
index c96d1bd..371db1a 100644
--- a/docs/tutorials/index.md
+++ b/docs/tutorials/index.md
@@ -14,3 +14,4 @@ The remainder of the tutorials are optional and for users who wish to e.g. train
 5. [Reproducing our results](reproducing_our_results/)
 6. [REL as systemd service](systemd_instructions/)
 7. [Notes on using custom models](custom_models/)
+7. [Conversational entity linking](conversations/)
diff --git a/mkdocs.yml b/mkdocs.yml
index 1ef6fd5..2a11e66 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -10,6 +10,7 @@ nav:
     - tutorials/reproducing_our_results.md
     - tutorials/systemd_instructions.md
     - tutorials/custom_models.md
+    - tutorials/conversations.md
   - Python API reference:
     - api/entity_disambiguation.md
     - api/generate_train_test.md
@@ -72,11 +73,10 @@ plugins:
         - https://numpy.org/doc/stable/objects.inv
         - https://docs.scipy.org/doc/scipy/objects.inv
         - https://pandas.pydata.org/docs/objects.inv
-        selection:
+        options:
           docstring_style: sphinx
           docstring_options:
             ignore_init_summary: yes
-        rendering:
           show_submodules: no
           show_source: true
           docstring_section_style: list
diff --git a/requirements.txt b/requirements.txt
index c5a5969..c84bf33 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,10 @@
+anyascii
 colorama
-konoha
+fastapi
 flair>=0.11
+konoha
+nltk
+pydantic
 segtok
 torch
-nltk
-anyascii
+uvicorn
diff --git a/setup.cfg b/setup.cfg
index 8fbd4af..92244af 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -43,13 +43,17 @@ package_dir =
     = src
 include_package_data = True
 install_requires =
+    anyascii
     colorama
-    konoha
+    fastapi
     flair>=0.11
+    konoha
+    nltk
+    pydantic
     segtok
+    spacy
     torch
-    nltk
-    anyascii
+    uvicorn
 
 [options.extras_require]
 develop =
diff --git a/src/REL/crel/__init__.py b/src/REL/crel/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/REL/crel/bert_md.py b/src/REL/crel/bert_md.py
new file mode 100644
index 0000000..23737bd
--- /dev/null
+++ b/src/REL/crel/bert_md.py
@@ -0,0 +1,94 @@
+import torch
+from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
+
+
+class BERT_MD:
+    def __init__(self, file_pretrained):
+        """
+
+        Args:
+            file_pretrained = "./tmp/ft-conel/"
+
+        Note:
+            The output of self.ner_model(s_input) is like
+              - s_input: e.g, 'Burger King franchise'
+              - return: e.g., [{'entity': 'B-ment', 'score': 0.99364895, 'index': 1, 'word': 'Burger', 'start': 0, 'end': 6}, ...]
+        """
+
+        model = AutoModelForTokenClassification.from_pretrained(file_pretrained)
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        model.to(device)
+        tokenizer = AutoTokenizer.from_pretrained(file_pretrained)
+        self.ner_model = pipeline(
+            "ner",
+            model=model,
+            tokenizer=tokenizer,
+            device=device.index if device.index != None else -1,
+            ignore_labels=[],
+        )
+
+    def md(self, s, flag_warning=False):
+        """Perform mention detection
+
+        Args:
+            s: input string
+            flag_warning: if True, print warning message
+
+        Returns:
+            REL style annotation results: [(start_position, length, mention), ...]
+            E.g., [[0, 15, 'The Netherlands'], ...]
+        """
+
+        ann = self.ner_model(s)  # Get ann results from BERT-NER model
+
+        ret = []
+        pos_start, pos_end = -1, -1  # Initialize variables
+
+        for i in range(len(ann)):
+            w, ner = ann[i]["word"], ann[i]["entity"]
+            assert ner in [
+                "B-ment",
+                "I-ment",
+                "O",
+            ], f"Unexpected ner tag: {ner}. If you use BERT-NER as it is, then you should flag_use_normal_bert_ner_tag=True."
+            if ner == "B-ment" and w[:2] != "##":
+                if (pos_start != -1) and (pos_end != -1):  # If B-ment is already found
+                    ret.append(
+                        [pos_start, pos_end - pos_start, s[pos_start:pos_end]]
+                    )  # save the previously identified mention
+                    pos_start, pos_end = -1, -1  # Initialize
+                pos_start, pos_end = ann[i]["start"], ann[i]["end"]
+
+            elif ner == "B-ment" and w[:2] == "##":
+                if (ann[i]["index"] == ann[i - 1]["index"] + 1) and (
+                    ann[i - 1]["entity"] != "B-ment"
+                ):  # If previous token has an entity (ner) label and it is NOT "B-ment" (i.e., ##xxx should not be the begin of the entity)
+                    if flag_warning:
+                        print(
+                            f"WARNING: ##xxx (in this case {w}) should not be the begin of the entity"
+                        )
+
+            elif (
+                i > 0
+                and (ner == "I-ment")
+                and (ann[i]["index"] == ann[i - 1]["index"] + 1)
+            ):  # If w is I-ment and previous word's index (i.e., ann[i-1]['index']) is also a mention
+                pos_end = ann[i]["end"]  # update pos_end
+
+            # This only happens when flag_ignore_o is False
+            elif (
+                ner == "O"
+                and w[:2] == "##"
+                and (
+                    ann[i - 1]["entity"] == "B-ment" or ann[i - 1]["entity"] == "I-ment"
+                )
+            ):  # If w is "O" and ##xxx, and previous token's index (i.e., ann[i-1]['index']) is B-ment or I-ment
+                pos_end = ann[i]["end"]  # update pos_end
+
+        # Append remaining ment
+        if (pos_start != -1) and (pos_end != -1):
+            ret.append(
+                [pos_start, pos_end - pos_start, s[pos_start:pos_end]]
+            )  # Save last mention
+
+        return ret
diff --git a/src/REL/crel/conv_el.py b/src/REL/crel/conv_el.py
new file mode 100644
index 0000000..868f962
--- /dev/null
+++ b/src/REL/crel/conv_el.py
@@ -0,0 +1,161 @@
+import importlib
+import sys
+from pathlib import Path
+
+from .bert_md import BERT_MD
+from .s2e_pe import pe_data
+from .s2e_pe.pe import EEMD, PEMD
+from REL.response_model import ResponseModel
+
+
+class ConvEL:
+    def __init__(
+        self, base_url=".", wiki_version="wiki_2019", ed_model=None, user_config=None, threshold=0
+    ):
+        self.threshold = threshold
+
+        self.wiki_version = wiki_version
+        self.base_url = base_url
+        self.file_pretrained = str(Path(base_url) / "bert_conv-td")
+
+        self.bert_md = BERT_MD(self.file_pretrained)
+
+        if not ed_model:
+            ed_model = self._default_ed_model()
+
+        self.response_model = ResponseModel(self.base_url, self.wiki_version, model=ed_model)
+        
+        self.eemd = EEMD(s2e_pe_model=str(Path(base_url) / "s2e_ast_onto"))
+        self.pemd = PEMD()
+
+        self.preprocess = pe_data.PreProcess()
+        self.postprocess = pe_data.PostProcess()
+
+        # These are always initialize when get_annotations() is called
+        self.conv_hist_for_pe = (
+            []
+        )  # initialize the history of conversation, which is used in PE Linking
+        self.ment2ent = {}  # This will be used for PE Linking
+
+    def _default_ed_model(self):
+        from REL.entity_disambiguation import EntityDisambiguation
+        return EntityDisambiguation(self.base_url, self.wiki_version, user_config={
+                "mode": "eval",
+                "model_path": f"{self.base_url}/{self.wiki_version}/generated/model",
+            })
+
+    def _error_check(self, conv):
+        assert type(conv) == list
+        for turn in conv:
+            assert type(turn) == dict
+            assert set(turn.keys()) == {"speaker", "utterance"}
+            assert turn["speaker"] in [
+                "USER",
+                "SYSTEM",
+            ], f'Speaker should be either "USER" or "SYSTEM", but got {turn["speaker"]}'
+
+    def _el(self, utt):
+        """Perform entity linking"""
+        # MD
+        md_results = self.bert_md.md(utt)
+
+        # ED
+        spans = [[r[0], r[1]] for r in md_results]  # r[0]: start, r[1]: length
+        el_results = self.ed(utt, spans)  # ED
+
+        self.conv_hist_for_pe[-1]["mentions"] = [r[2] for r in el_results]
+        self.ment2ent.update(
+            {r[2]: r[3] for r in el_results}
+        )  # If there is a mismatch of annotations for the same mentions, the last one (the most closest turn's one to the PEM) will be used.
+
+        return [r[:4] for r in el_results]  # [start_pos, length, mention, entity]
+
+    def _pe(self, utt):
+        """Perform PE Linking"""
+
+        ret = []
+
+        # Step 1: PE Mention Detection
+        pem_results = self.pemd.pem_detector(utt)
+        pem2result = {r[2]: r for r in pem_results}
+
+        # Step 2: Finding corresponding explicit entity mentions (EEMs)
+        # NOTE: Current implementation can handle only one target PEM at a time
+        outputs = []
+        for _, _, pem in pem_results:  # pems: [[start_pos, length, pem], ...]
+            self.conv_hist_for_pe[-1]["pems"] = [
+                pem
+            ]  # Create a conv for each target PEM that you want to link
+
+            # Preprocessing
+            token_with_info = self.preprocess.get_tokens_with_info(
+                self.conv_hist_for_pe
+            )
+            input_data = self.preprocess.get_input_of_pe_linking(token_with_info)
+
+            assert (
+                len(input_data) == 1
+            ), f"Current implementation can handle only one target PEM at a time"
+            input_data = input_data[0]
+
+            # Finding corresponding explicit entity mentions (EEMs)
+            scores = self.eemd.get_scores(input_data)
+
+            # Post processing
+            outputs += self.postprocess.get_results(
+                input_data, self.conv_hist_for_pe, self.threshold, scores
+            )
+
+        self.conv_hist_for_pe[-1]["pems"] = []  # Remove the target PEM
+
+        # Step 3: Get corresponding entity
+        for r in outputs:
+            pem = r["personal_entity_mention"]
+            pem_result = pem2result[pem]  # [start_pos, length, pem]
+            eem = r["mention"]  # Explicit entity mention
+            ent = self.ment2ent[eem]  # Corresponding entity
+            ret.append(
+                [pem_result[0], pem_result[1], pem_result[2], ent]
+            )  # [start_pos, length, PEM, entity]
+
+        return ret
+
+    def annotate(self, conv):
+        """Get conversational entity linking annotations
+
+        Args:
+            conv: A list of dicts, each dict contains "speaker" and "utterance" keys.
+
+        Returns:
+            A list of dicts, each dict contains conv's ones + "annotations" key.
+        """
+        self._error_check(conv)
+        ret = []
+        self.conv_hist_for_pe = []  # Initialize
+        self.ment2ent = {}  # Initialize
+
+        for turn in conv:
+            utt = turn["utterance"]
+            assert turn["speaker"] in [
+                "USER",
+                "SYSTEM",
+            ], f'Speaker should be either "USER" or "SYSTEM", but got {turn["speaker"]}'
+            ret.append({"speaker": turn["speaker"], "utterance": utt})
+
+            self.conv_hist_for_pe.append({})
+            self.conv_hist_for_pe[-1]["speaker"] = turn["speaker"]
+            self.conv_hist_for_pe[-1]["utterance"] = utt
+
+            if turn["speaker"] == "USER":
+                el_results = self._el(utt)
+                pe_results = self._pe(utt)
+                ret[-1]["annotations"] = el_results + pe_results
+
+        return ret
+
+    def ed(self, text, spans):
+        """Change tuple to list to match the output format of REL API."""
+        response = self.response_model.generate_response(text=text, spans=spans)
+        return [list(ent) for ent in response]
+
+
diff --git a/src/REL/crel/s2e_pe/__init__.py b/src/REL/crel/s2e_pe/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/REL/crel/s2e_pe/consts.py b/src/REL/crel/s2e_pe/consts.py
new file mode 100644
index 0000000..97bffe5
--- /dev/null
+++ b/src/REL/crel/s2e_pe/consts.py
@@ -0,0 +1,3 @@
+SPEAKER_START = 49518  # 'Ġ#####'
+SPEAKER_END = 22560  # 'Ġ###'
+NULL_ID_FOR_COREF = 0
diff --git a/src/REL/crel/s2e_pe/coref_bucket_batch_sampler.py b/src/REL/crel/s2e_pe/coref_bucket_batch_sampler.py
new file mode 100644
index 0000000..b252973
--- /dev/null
+++ b/src/REL/crel/s2e_pe/coref_bucket_batch_sampler.py
@@ -0,0 +1,68 @@
+import logging
+import math
+import random
+from typing import Iterable, List
+
+from torch.utils import data
+from torch.utils.data import DataLoader
+
+logger = logging.getLogger(__name__)
+
+
+class BucketBatchSampler(DataLoader):
+    def __init__(
+        self,
+        data_source: data.Dataset,
+        max_total_seq_len: int,
+        sorting_keys: List[str] = None,
+        padding_noise: float = 0.1,
+        drop_last: bool = False,
+        batch_size_1: bool = False,
+    ):
+        self.sorting_keys = sorting_keys
+        self.padding_noise = padding_noise
+        self.max_total_seq_len = max_total_seq_len
+        self.data_source = data_source
+        data_source.examples.sort(key=lambda x: len(x[1].token_ids), reverse=True)
+        self.drop_last = drop_last
+        self.batches = (
+            self.prepare_batches() if not batch_size_1 else self.prepare_eval_batches()
+        )
+
+    def prepare_batches(self):
+        batches = []
+        batch = []
+        per_example_batch_len = 0
+        for (
+            doc_key,
+            elem,
+        ) in self.data_source:  # NOTE: The `doc_key` also contains subtoken map
+            batch.append(elem)
+            batch = self.data_source.pad_batch(batch, len(batch[0].token_ids))
+            batches.append((doc_key, batch))
+            batch = []
+
+        if len(batch) == 0:
+            return batches
+        batch = self.data_source.pad_batch(batch, len(batch[0].token_ids))
+        batches.append(
+            (doc_key, batch)
+        )  # 220229: Change to return dockey as `prepare_eval_batches()` does.
+        return batches
+
+    def __iter__(self) -> Iterable[List[int]]:
+        random.shuffle(self.batches)
+        yield from self.batches
+
+    def __len__(self):
+        return len(self.batches)
+
+    def calc_effective_per_example_batch_len(self, example_len):
+        return math.ceil((example_len + 2) / 512) * 512
+
+    def prepare_eval_batches(self):
+        batches = []
+        for doc_key, elem in self.data_source:
+            batch = self.data_source.pad_batch([elem], len(elem.token_ids))
+            batches.append((doc_key, batch))
+        return batches
diff --git a/src/REL/crel/s2e_pe/data.py b/src/REL/crel/s2e_pe/data.py
new file mode 100644
index 0000000..7cb5488
--- /dev/null
+++ b/src/REL/crel/s2e_pe/data.py
@@ -0,0 +1,260 @@
+import json
+import logging
+import os
+import pickle
+from collections import namedtuple
+
+import torch
+from torch.utils.data import Dataset
+
+from .consts import NULL_ID_FOR_COREF, SPEAKER_END, SPEAKER_START
+from .utils import flatten_list_of_lists
+
+CorefExample = namedtuple("CorefExample", ["token_ids", "clusters"])
+
+logger = logging.getLogger(__name__)
+
+
+class CorefDataset(Dataset):
+    def __init__(self, input_data, tokenizer, model_name_or_path, max_seq_length=-1):
+        self.tokenizer = tokenizer
+        (
+            examples,
+            self.max_mention_num,
+            self.max_cluster_size,
+            self.max_num_clusters,
+            dockey2eems_tokenspan,
+            dockey2pems_tokenspan,
+        ) = self._parse_jsonlines(input_data)
+        self.max_seq_length = max_seq_length
+        (
+            self.examples,
+            self.lengths,
+            self.num_examples_filtered,
+            self.dockey2eems_subtokenspan,
+            self.dockey2pems_subtokenspan,
+        ) = self._tokenize(
+            examples, dockey2eems_tokenspan, dockey2pems_tokenspan, model_name_or_path
+        )
+        logger.info(
+            f"Finished preprocessing Coref dataset. {len(self.examples)} examples were extracted, {self.num_examples_filtered} were filtered due to sequence length."
+        )
+
+    def _parse_jsonlines(self, d):
+        examples = []
+        max_mention_num = -1
+        max_cluster_size = -1
+        max_num_clusters = -1
+        dockey2pems_tokenspan = {}
+        dockey2eems_tokenspan = {}
+        doc_key = d["doc_key"]
+        assert (
+            type(d["sentences"][0]) == list
+        ), "'sentences' should be 2d list, not just a 1d list of the tokens."
+        input_words = flatten_list_of_lists(d["sentences"])
+        clusters = d["clusters"]
+        max_mention_num = max(max_mention_num, len(flatten_list_of_lists(clusters)))
+        max_cluster_size = max(
+            max_cluster_size,
+            max(len(cluster) for cluster in clusters) if clusters else 0,
+        )
+        max_num_clusters = max(max_num_clusters, len(clusters) if clusters else 0)
+        speakers = flatten_list_of_lists(d["speakers"])
+        examples.append((doc_key, input_words, clusters, speakers))
+        dockey2eems_tokenspan[doc_key] = d["mentions"]
+        dockey2pems_tokenspan[doc_key] = d["pems"]
+        return (
+            examples,
+            max_mention_num,
+            max_cluster_size,
+            max_num_clusters,
+            dockey2eems_tokenspan,
+            dockey2pems_tokenspan,
+        )
+
+    def _tokenize(
+        self, examples, dockey2eems_tokenspan, dockey2pems_tokenspan, model_name_or_path
+    ):
+        coref_examples = []
+        lengths = []
+        num_examples_filtered = 0
+        dockey2eems_subtokenspan = {}
+        dockey2pems_subtokenspan = {}
+        for doc_key, words, clusters, speakers in examples:
+            word_idx_to_start_token_idx = dict()
+            word_idx_to_end_token_idx = dict()
+            end_token_idx_to_word_idx = [0]  # for 

+
+            token_ids = []
+            last_speaker = None
+            for idx, (word, speaker) in enumerate(zip(words, speakers)):
+                if last_speaker != speaker:
+                    speaker_prefix = (
+                        [SPEAKER_START]
+                        + self.tokenizer.encode(" " + speaker, add_special_tokens=False)
+                        + [SPEAKER_END]
+                    )
+                    last_speaker = speaker
+                else:
+                    speaker_prefix = []
+                for _ in range(len(speaker_prefix)):
+                    end_token_idx_to_word_idx.append(idx)
+                token_ids.extend(speaker_prefix)
+                word_idx_to_start_token_idx[idx] = len(token_ids) + 1  # +1 for 

+                tokenized = self.tokenizer.encode(" " + word, add_special_tokens=False)
+                for _ in range(len(tokenized)):
+                    end_token_idx_to_word_idx.append(idx)
+                token_ids.extend(tokenized)
+                word_idx_to_end_token_idx[idx] = len(
+                    token_ids
+                )  # old_seq_len + 1 (for 
) + len(tokenized_word) - 1 (we start counting from zero) = len(token_ids)
+
+            if 0 < self.max_seq_length < len(token_ids):
+                num_examples_filtered += 1
+                continue
+
+            new_clusters = [
+                [
+                    (word_idx_to_start_token_idx[start], word_idx_to_end_token_idx[end])
+                    for start, end in cluster
+                ]
+                for cluster in clusters
+            ]
+            lengths.append(len(token_ids))
+
+            coref_examples.append(
+                (
+                    (doc_key, end_token_idx_to_word_idx),
+                    CorefExample(token_ids=token_ids, clusters=new_clusters),
+                )
+            )
+
+            dockey2eems_subtokenspan[doc_key] = [
+                (word_idx_to_start_token_idx[start], word_idx_to_end_token_idx[end])
+                for start, end in dockey2eems_tokenspan[doc_key]
+            ]
+            dockey2pems_subtokenspan[doc_key] = [
+                (word_idx_to_start_token_idx[start], word_idx_to_end_token_idx[end])
+                for start, end in dockey2pems_tokenspan[doc_key]
+            ]
+
+        return (
+            coref_examples,
+            lengths,
+            num_examples_filtered,
+            dockey2eems_subtokenspan,
+            dockey2pems_subtokenspan,
+        )
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item):
+        return self.examples[item]
+
+    def pad_clusters_inside(self, clusters):
+        return [
+            cluster
+            + [(NULL_ID_FOR_COREF, NULL_ID_FOR_COREF)]
+            * (self.max_cluster_size - len(cluster))
+            for cluster in clusters
+        ]
+
+    def pad_clusters_outside(self, clusters):
+        return clusters + [[]] * (self.max_num_clusters - len(clusters))
+
+    def pad_clusters(self, clusters):
+        clusters = self.pad_clusters_outside(clusters)
+        clusters = self.pad_clusters_inside(clusters)
+        return clusters
+
+    def _pe_create_tensored_batch(self, padded_batch, len_example):
+        """Create tensored_batch avoiding errors
+        Original code was:
+        `tensored_batch = tuple(torch.stack([example[i].squeeze() for example in padded_batch], dim=0) for i in range(len(example)))`
+        However, this does not handle the single cluster case (E.g., "clusters": [[[135, 136], [273, 273]]] in the train.english.jsonlines)
+
+        The error caused by the above is like (220322):
+            gold_clusters = [tuple(tuple(m) for m in gc if NULL_ID_FOR_COREF not in m) for gc in gold_clusters.tolist()]
+            TypeError: argument of type 'int' is not iterable
+
+        - Updates:
+          - 220228: Created
+          - 220322: Write the error details
+        """
+        assert len_example == 3
+        tensored_batch = tuple()
+        for i in range(len_example):
+            to_stack = []
+            for example in padded_batch:
+                assert (
+                    len(example) == 3
+                ), f"example contains three components: input_ids, attention_mask, and clusters. Current len(examples): {len(example)}"
+                if i < 2:  # input_ids and attention_mask
+                    to_stack.append(example[i].squeeze())
+                elif i == 2:  # clusters
+                    to_stack.append(
+                        example[i]
+                    )  # squeeze() cause the error to single-cluster case
+            # add to_stack to tensored_batch (tuple)
+            tensored_batch += (torch.stack(to_stack, dim=0),)
+        return tensored_batch
+
+    def pad_batch(self, batch, max_length):
+        max_length += 2  # we have additional two special tokens 
, 

+        padded_batch = []
+        for example in batch:
+            # encoded_dict = self.tokenizer.encode_plus(example[0], # This does not work transformers v4.18.0 (works with v3.3.1)
+            # See: https://github.com/huggingface/transformers/issues/10297
+            encoded_dict = self.tokenizer.encode_plus(
+                example[0],
+                use_fast=False,
+                add_special_tokens=True,
+                pad_to_max_length=True,
+                max_length=max_length,
+                return_attention_mask=True,
+                return_tensors="pt",
+            )
+            clusters = self.pad_clusters(example.clusters)
+            example = (encoded_dict["input_ids"], encoded_dict["attention_mask"]) + (
+                torch.tensor(clusters),
+            )
+            padded_batch.append(example)
+        # tensored_batch = tuple(torch.stack([example[i].squeeze() for example in padded_batch], dim=0) for i in range(len(example)))
+        tensored_batch = self._pe_create_tensored_batch(
+            padded_batch, len(example)
+        )  # HJ: 220228
+        return tensored_batch
+
+
+def get_dataset(tokenizer, input_data, conf):
+    """Read input data
+
+    Args:
+        - tokenizer
+        - input_data (dict): Input dict containing the following keys:
+            dict_keys(['clusters', 'doc_key', 'mentions', 'pems', 'sentences', 'speakers'])
+            E.g.,
+                test_jsonl = {
+                    "clusters": [[[78, 83], [88, 89]]], # This can be blank when you want to perform prediction.
+                    "doc_key": "dialind:0_turn:3_pem:my-favorite-forms-of-science-fiction", # doc_key should be unique, no restrictions on the format
+                    "mentions": [[35, 35], [37, 38], [74, 74], [85, 85], [88, 89]], # mentions and spans should be token-level spans (i.e., different from REL). See original document of s2e-coref.
+                    "pems": [[78, 83]],
+                    "sentences": [["I", "think", "science", "fiction", "is", ...], ...],
+                    "speakers": [["SYSTEM", "SYSTEM", "SYSTEM", ..., "USER", "USER", "USER", ...], ...], }
+
+
+    Returns:
+        - dataset (CorefDataset):
+
+    Notes:
+        - Currently, parallel processing is not supported, i.e., you cannot input more than or equal to two sentences or PEMs at the same time.
+    """
+
+    coref_dataset = CorefDataset(
+        input_data,
+        tokenizer,
+        max_seq_length=conf.max_seq_length,
+        model_name_or_path=conf.model_name_or_path,
+    )
+    return coref_dataset
diff --git a/src/REL/crel/s2e_pe/modeling.py b/src/REL/crel/s2e_pe/modeling.py
new file mode 100644
index 0000000..7ed3150
--- /dev/null
+++ b/src/REL/crel/s2e_pe/modeling.py
@@ -0,0 +1,476 @@
+import torch
+from torch.nn import Dropout, LayerNorm, Linear, Module
+from transformers import BertModel, BertPreTrainedModel, LongformerModel
+
+try:  # If you use `211018_s2e_coref`
+    from transformers.modeling_bert import ACT2FN
+except:  # If you use `jupyterlab-debugger`
+    from transformers.models.bert.modeling_bert import ACT2FN
+
+import json
+import os
+
+from .utils import extract_clusters  # , ce_get_start_end_subtoken_num
+from .utils import extract_mentions_to_predicted_clusters_from_clusters, mask_tensor
+
+
+class FullyConnectedLayer(Module):
+    def __init__(self, config, input_dim, output_dim, dropout_prob):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.dropout_prob = dropout_prob
+
+        self.dense = Linear(self.input_dim, self.output_dim)
+        self.layer_norm = LayerNorm(self.output_dim, eps=config.layer_norm_eps)
+        self.activation_func = ACT2FN[config.hidden_act]
+        self.dropout = Dropout(self.dropout_prob)
+
+    def forward(self, inputs):
+        temp = inputs
+        temp = self.dense(temp)
+        temp = self.activation_func(temp)
+        temp = self.layer_norm(temp)
+        temp = self.dropout(temp)
+        return temp
+
+
+class S2E(BertPreTrainedModel):
+    def __init__(self, config, args):
+        super().__init__(config)
+        self.max_span_length = args.max_span_length
+        self.top_lambda = args.top_lambda
+        self.ffnn_size = args.ffnn_size
+        self.do_mlps = self.ffnn_size > 0
+        self.ffnn_size = self.ffnn_size if self.do_mlps else config.hidden_size
+        self.normalise_loss = args.normalise_loss
+
+        # self.longformer = LongformerModel(config)
+        self.longformer = BertModel(config)
+
+        self.start_mention_mlp = (
+            FullyConnectedLayer(
+                config, config.hidden_size, self.ffnn_size, args.dropout_prob
+            )
+            if self.do_mlps
+            else None
+        )
+        self.end_mention_mlp = (
+            FullyConnectedLayer(
+                config, config.hidden_size, self.ffnn_size, args.dropout_prob
+            )
+            if self.do_mlps
+            else None
+        )
+        self.start_coref_mlp = (
+            FullyConnectedLayer(
+                config, config.hidden_size, self.ffnn_size, args.dropout_prob
+            )
+            if self.do_mlps
+            else None
+        )
+        self.end_coref_mlp = (
+            FullyConnectedLayer(
+                config, config.hidden_size, self.ffnn_size, args.dropout_prob
+            )
+            if self.do_mlps
+            else None
+        )
+
+        self.start_coref_mlp = (
+            FullyConnectedLayer(
+                config, config.hidden_size, self.ffnn_size, args.dropout_prob
+            )
+            if self.do_mlps
+            else None
+        )
+        self.end_coref_mlp = (
+            FullyConnectedLayer(
+                config, config.hidden_size, self.ffnn_size, args.dropout_prob
+            )
+            if self.do_mlps
+            else None
+        )
+
+        self.mention_start_classifier = Linear(self.ffnn_size, 1)
+        self.mention_end_classifier = Linear(self.ffnn_size, 1)
+        self.mention_s2e_classifier = Linear(self.ffnn_size, self.ffnn_size)
+
+        self.antecedent_s2s_classifier = Linear(self.ffnn_size, self.ffnn_size)
+        self.antecedent_e2e_classifier = Linear(self.ffnn_size, self.ffnn_size)
+        self.antecedent_s2e_classifier = Linear(self.ffnn_size, self.ffnn_size)
+        self.antecedent_e2s_classifier = Linear(self.ffnn_size, self.ffnn_size)
+
+        self.init_weights()
+
+    def _get_span_mask(self, batch_size, k, max_k):
+        """
+        :param batch_size: int
+        :param k: tensor of size [batch_size], with the required k for each example
+        :param max_k: int
+        :return: [batch_size, max_k] of zero-ones, where 1 stands for a valid span and 0 for a padded span
+        """
+        size = (batch_size, max_k)
+        idx = torch.arange(max_k, device=self.device).unsqueeze(0).expand(size)
+        len_expanded = k.unsqueeze(1).expand(size).to(self.device)
+        ret = (idx < len_expanded).int()
+        return ret
+
+    def _prune_topk_mentions(self, mention_logits, attention_mask):
+        """
+        :param mention_logits: Shape [batch_size, seq_length, seq_length]
+        :param attention_mask: [batch_size, seq_length]
+        :param top_lambda:
+        :return:
+        """
+        batch_size, seq_length, _ = mention_logits.size()
+        actual_seq_lengths = torch.sum(attention_mask, dim=-1)  # [batch_size]
+
+        k = (
+            actual_seq_lengths * self.top_lambda
+        ).int()  # [batch_size] # top_lambda is in argument of the `run_coref.py`
+        max_k = int(
+            torch.max(k)
+        )  # This is the k for the largest input in the batch, we will need to pad
+
+        _, topk_1d_indices = torch.topk(
+            mention_logits.view(batch_size, -1), dim=-1, k=max_k
+        )  # [batch_size, max_k]
+        span_mask = self._get_span_mask(batch_size, k, max_k)  # [batch_size, max_k]
+        topk_1d_indices = (topk_1d_indices * span_mask) + (1 - span_mask) * (
+            (seq_length**2) - 1
+        )  # We take different k for each example
+        sorted_topk_1d_indices, _ = torch.sort(
+            topk_1d_indices, dim=-1
+        )  # [batch_size, max_k]
+
+        topk_mention_start_ids = (
+            sorted_topk_1d_indices // seq_length
+        )  # [batch_size, max_k]
+        topk_mention_end_ids = (
+            sorted_topk_1d_indices % seq_length
+        )  # [batch_size, max_k]
+
+        topk_mention_logits = mention_logits[
+            torch.arange(batch_size).unsqueeze(-1).expand(batch_size, max_k),
+            topk_mention_start_ids,
+            topk_mention_end_ids,
+        ]  # [batch_size, max_k]
+
+        topk_mention_logits = topk_mention_logits.unsqueeze(
+            -1
+        ) + topk_mention_logits.unsqueeze(
+            -2
+        )  # [batch_size, max_k, max_k]
+
+        return (
+            topk_mention_start_ids,
+            topk_mention_end_ids,
+            span_mask,
+            topk_mention_logits,
+        )
+
+    def _ce_prune_pem_eem(
+        self, mention_logits, pem_eem_subtokenspan
+    ):  # attention_mask, subtoken_map, pem_eem_subtokenspan):
+
+        batch_size, seq_length, _ = mention_logits.size()
+        assert batch_size == 1  # HJ: currently, only batch_size==1 is supported
+
+        k = torch.Tensor([len(pem_eem_subtokenspan)])
+
+        max_k = int(
+            torch.max(k)
+        )  # This is the k for the largest input in the batch, we will need to pad
+
+        span_mask = self._get_span_mask(batch_size, k, max_k)  # [batch_size, max_k]
+
+        pem_eem_start_ids, pem_eem_end_ids = torch.Tensor(
+            [[span[0] for span in pem_eem_subtokenspan]]
+        ).long().to(self.device), torch.Tensor(
+            [[span[1] for span in pem_eem_subtokenspan]]
+        ).long().to(
+            self.device
+        )  # HJ: 220302: [[...]] should be use because we have n batchs (here, we use n=1 though)
+
+        return (
+            pem_eem_start_ids,
+            pem_eem_end_ids,
+            span_mask,
+            None,
+        )  # topk_mention_logits
+
+    def _mask_antecedent_logits(self, antecedent_logits, span_mask):
+        # We now build the matrix for each pair of spans (i,j) - whether j is a candidate for being antecedent of i?
+        antecedents_mask = torch.ones_like(antecedent_logits, dtype=self.dtype).tril(
+            diagonal=-1
+        )  # [batch_size, k, k]
+        antecedents_mask = antecedents_mask * span_mask.unsqueeze(
+            -1
+        )  # [batch_size, k, k]
+        antecedent_logits = mask_tensor(antecedent_logits, antecedents_mask)
+        return antecedent_logits
+
+    def _get_cluster_labels_after_pruning(self, span_starts, span_ends, all_clusters):
+        """
+        :param span_starts: [batch_size, max_k]
+        :param span_ends: [batch_size, max_k]
+        :param all_clusters: [batch_size, max_cluster_size, max_clusters_num, 2]
+        :return: [batch_size, max_k, max_k + 1] - [b, i, j] == 1 if i is antecedent of j
+        """
+        batch_size, max_k = span_starts.size()
+        new_cluster_labels = torch.zeros((batch_size, max_k, max_k + 1), device="cpu")
+        all_clusters_cpu = all_clusters.cpu().numpy()
+        for b, (starts, ends, gold_clusters) in enumerate(
+            zip(span_starts.cpu().tolist(), span_ends.cpu().tolist(), all_clusters_cpu)
+        ):
+            gold_clusters = extract_clusters(gold_clusters)
+            mention_to_gold_clusters = (
+                extract_mentions_to_predicted_clusters_from_clusters(gold_clusters)
+            )
+            gold_mentions = set(mention_to_gold_clusters.keys())
+            for i, (start, end) in enumerate(zip(starts, ends)):
+                if (start, end) not in gold_mentions:
+                    continue
+                for j, (a_start, a_end) in enumerate(list(zip(starts, ends))[:i]):
+                    if (a_start, a_end) in mention_to_gold_clusters[(start, end)]:
+                        new_cluster_labels[b, i, j] = 1
+        new_cluster_labels = new_cluster_labels.to(self.device)
+        no_antecedents = 1 - torch.sum(new_cluster_labels, dim=-1).bool().float()
+        new_cluster_labels[:, :, -1] = no_antecedents
+        return new_cluster_labels
+
+    def _get_marginal_log_likelihood_loss(
+        self, coref_logits, cluster_labels_after_pruning, span_mask
+    ):
+        """
+        :param coref_logits: [batch_size, max_k, max_k]
+        :param cluster_labels_after_pruning: [batch_size, max_k, max_k]
+        :param span_mask: [batch_size, max_k]
+        :return:
+        """
+        gold_coref_logits = mask_tensor(coref_logits, cluster_labels_after_pruning)
+
+        gold_log_sum_exp = torch.logsumexp(
+            gold_coref_logits, dim=-1
+        )  # [batch_size, max_k]
+        all_log_sum_exp = torch.logsumexp(coref_logits, dim=-1)  # [batch_size, max_k]
+
+        gold_log_probs = gold_log_sum_exp - all_log_sum_exp
+        losses = -gold_log_probs
+        losses = losses * span_mask
+        per_example_loss = torch.sum(losses, dim=-1)  # [batch_size]
+        if self.normalise_loss:
+            per_example_loss = per_example_loss / losses.size(-1)
+        loss = per_example_loss.mean()
+        return loss
+
+    def _get_mention_mask(self, mention_logits_or_weights):
+        """
+        Returns a tensor of size [batch_size, seq_length, seq_length] where valid spans
+        (start <= end < start + max_span_length) are 1 and the rest are 0
+        :param mention_logits_or_weights: Either the span mention logits or weights, size [batch_size, seq_length, seq_length]
+        """
+        mention_mask = torch.ones_like(mention_logits_or_weights, dtype=self.dtype)
+        mention_mask = mention_mask.triu(diagonal=0)
+        mention_mask = mention_mask.tril(diagonal=self.max_span_length - 1)
+        return mention_mask
+
+    def _calc_mention_logits(self, start_mention_reps, end_mention_reps):
+        start_mention_logits = self.mention_start_classifier(
+            start_mention_reps
+        ).squeeze(
+            -1
+        )  # [batch_size, seq_length]
+        end_mention_logits = self.mention_end_classifier(end_mention_reps).squeeze(
+            -1
+        )  # [batch_size, seq_length]
+
+        temp = self.mention_s2e_classifier(
+            start_mention_reps
+        )  # [batch_size, seq_length]
+        joint_mention_logits = torch.matmul(
+            temp, end_mention_reps.permute([0, 2, 1])
+        )  # [batch_size, seq_length, seq_length]
+
+        mention_logits = (
+            joint_mention_logits
+            + start_mention_logits.unsqueeze(-1)
+            + end_mention_logits.unsqueeze(-2)
+        )
+        mention_mask = self._get_mention_mask(
+            mention_logits
+        )  # [batch_size, seq_length, seq_length]
+        mention_logits = mask_tensor(
+            mention_logits, mention_mask
+        )  # [batch_size, seq_length, seq_length]
+        return mention_logits
+
+    def _calc_coref_logits(self, top_k_start_coref_reps, top_k_end_coref_reps):
+        # s2s
+        temp = self.antecedent_s2s_classifier(
+            top_k_start_coref_reps
+        )  # [batch_size, max_k, dim]
+        top_k_s2s_coref_logits = torch.matmul(
+            temp, top_k_start_coref_reps.permute([0, 2, 1])
+        )  # [batch_size, max_k, max_k]
+
+        # e2e
+        temp = self.antecedent_e2e_classifier(
+            top_k_end_coref_reps
+        )  # [batch_size, max_k, dim]
+        top_k_e2e_coref_logits = torch.matmul(
+            temp, top_k_end_coref_reps.permute([0, 2, 1])
+        )  # [batch_size, max_k, max_k]
+
+        # s2e
+        temp = self.antecedent_s2e_classifier(
+            top_k_start_coref_reps
+        )  # [batch_size, max_k, dim]
+        top_k_s2e_coref_logits = torch.matmul(
+            temp, top_k_end_coref_reps.permute([0, 2, 1])
+        )  # [batch_size, max_k, max_k]
+
+        # e2s
+        temp = self.antecedent_e2s_classifier(
+            top_k_end_coref_reps
+        )  # [batch_size, max_k, dim]
+        top_k_e2s_coref_logits = torch.matmul(
+            temp, top_k_start_coref_reps.permute([0, 2, 1])
+        )  # [batch_size, max_k, max_k]
+
+        # sum all terms
+        coref_logits = (
+            top_k_s2e_coref_logits
+            + top_k_e2s_coref_logits
+            + top_k_s2s_coref_logits
+            + top_k_e2e_coref_logits
+        )  # [batch_size, max_k, max_k]
+        return coref_logits
+
+    def _ce_get_scores(
+        self, mention_start_ids, mention_end_ids, subtoken_map, final_logits, doc_key
+    ):
+        """Get scores"""
+        scores = (
+            []
+        )  # score_dc = [{'span_subtoken':(start_id, end_id), 'coref_logits':coref_logits}, ...]
+        _mention_start_ids_flatten = mention_start_ids[0]  # [N_mention]
+        _mention_end_ids_flatten = mention_end_ids[0]  # [N_mention]
+        _final_logits_2d = final_logits[
+            0
+        ]  # [N_mention, N_mention+1] (+1 is for threshold)
+        # Get the length of the _mention_start_ids_flatten
+        N = len(_mention_start_ids_flatten)
+
+        for i in range(N):  # loop for anaphoras
+            for j in range(N):  # loop for antecedents
+                if i <= j:
+                    continue  # anaphoras should not come before antecedents
+                scores.append(
+                    {
+                        "doc_key": doc_key,
+                        "span_token_anaphora": (
+                            int(subtoken_map[_mention_start_ids_flatten[i]]),
+                            int(subtoken_map[_mention_end_ids_flatten[i]]),
+                        ),
+                        "span_token_antecedent": (
+                            int(subtoken_map[_mention_start_ids_flatten[j]]),
+                            int(subtoken_map[_mention_end_ids_flatten[j]]),
+                        ),
+                        # 'span_subtoken': (int(_mention_start_ids_flatten[i]), int(_mention_end_ids_flatten[i])),
+                        # 'subtoken_map': subtoken_map,
+                        "coref_logits": float(_final_logits_2d[i][j]),
+                    }
+                )
+
+        return scores
+
+    # def forward(self, input_ids, attention_mask=None, gold_clusters=None, return_all_outputs=False):
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        gold_clusters=None,
+        return_all_outputs=False,
+        subtoken_map=None,
+        pem_eem_subtokenspan=None,
+        doc_key=None,
+    ):
+        # TODO: Change the argument of this forward func, from `pem_eem` to `pem` and `eem`
+        # And do pem_eem = pem+eem
+
+        outputs = self.longformer(input_ids, attention_mask=attention_mask)
+        sequence_output = outputs[0]  # [batch_size, seq_len, dim]
+        # MEMO: `sequence_output` should be a hidden vector of the model.
+        #       Originally, this seems to be acquired by `outputs.last_hidden_state` (https://huggingface.co/transformers/master/model_doc/longformer.html)
+
+        # Compute representations
+        start_mention_reps = (
+            self.start_mention_mlp(sequence_output) if self.do_mlps else sequence_output
+        )
+        end_mention_reps = (
+            self.end_mention_mlp(sequence_output) if self.do_mlps else sequence_output
+        )
+
+        start_coref_reps = (
+            self.start_coref_mlp(sequence_output) if self.do_mlps else sequence_output
+        )
+        end_coref_reps = (
+            self.end_coref_mlp(sequence_output) if self.do_mlps else sequence_output
+        )
+
+        # mention scores
+        mention_logits = self._calc_mention_logits(start_mention_reps, end_mention_reps)
+
+        # prune mentions
+        # (span_mask: [batch_size, max_k] of zero-ones, where 1 stands for a valid span and 0 for a padded span)
+        # mention_start_ids, mention_end_ids, span_mask, topk_mention_logits = self._prune_topk_mentions(mention_logits, attention_mask)
+        mention_start_ids, mention_end_ids, span_mask, _ = self._ce_prune_pem_eem(
+            mention_logits, pem_eem_subtokenspan
+        )
+
+        batch_size, _, dim = start_coref_reps.size()
+        max_k = mention_start_ids.size(-1)
+        size = (batch_size, max_k, dim)
+
+        # Antecedent scores
+        # gather reps
+        topk_start_coref_reps = torch.gather(
+            start_coref_reps, dim=1, index=mention_start_ids.unsqueeze(-1).expand(size)
+        )
+        topk_end_coref_reps = torch.gather(
+            end_coref_reps, dim=1, index=mention_end_ids.unsqueeze(-1).expand(size)
+        )
+        coref_logits = self._calc_coref_logits(
+            topk_start_coref_reps, topk_end_coref_reps
+        )
+        final_logits = coref_logits  # topk_mention_logits + coref_logits
+        final_logits = self._mask_antecedent_logits(final_logits, span_mask)
+        # adding zero logits for null span
+        final_logits = torch.cat(
+            (final_logits, torch.zeros((batch_size, max_k, 1), device=self.device)),
+            dim=-1,
+        )  # [batch_size, max_k, max_k + 1]
+        scores = self._ce_get_scores(
+            mention_start_ids, mention_end_ids, subtoken_map, final_logits, doc_key
+        )
+
+        if return_all_outputs:
+            outputs = (mention_start_ids, mention_end_ids, final_logits, mention_logits)
+        else:
+            outputs = tuple()
+
+        if gold_clusters is not None:
+            losses = {}
+            labels_after_pruning = self._get_cluster_labels_after_pruning(
+                mention_start_ids, mention_end_ids, gold_clusters
+            )
+            loss = self._get_marginal_log_likelihood_loss(
+                final_logits, labels_after_pruning, span_mask
+            )  # HJ: 220303: `labels_after_pruning` is strange...
+            losses.update({"loss": loss})
+            outputs = (loss,) + outputs + (losses,)
+
+        return outputs, scores
diff --git a/src/REL/crel/s2e_pe/pe.py b/src/REL/crel/s2e_pe/pe.py
new file mode 100644
index 0000000..e6ca80c
--- /dev/null
+++ b/src/REL/crel/s2e_pe/pe.py
@@ -0,0 +1,259 @@
+# PEMD
+from tokenizers.pre_tokenizers import Whitespace
+
+pre_tokenizer = Whitespace()
+try:
+    import spacy
+
+    nlp = spacy.load("en_core_web_md")
+except:  # From Google Colab (see https://stackoverflow.com/a/59197634)
+    import spacy.cli
+
+    spacy.cli.download("en_core_web_md")
+    import en_core_web_md
+
+    nlp = en_core_web_md.load()
+import torch
+from transformers import AutoConfig, AutoTokenizer, LongformerConfig
+
+# EEMD
+from . import data
+from .coref_bucket_batch_sampler import BucketBatchSampler
+from .modeling import S2E
+from .pe_data import PreProcess  # to use get_span()
+
+
+class PEMD:
+    """Responsible for personal entity mention detection"""
+
+    def __init__(self):
+        self.pronouns = ["my", "our"]  # These should be lowercase
+        self.preprocess = PreProcess()  # to use get_span()
+
+    def _extract_text_with_pronoun(self, utt: str, max_candidate_num=10):
+        """
+
+        Args:
+            max_candidate_num (int): Max following words num (which equals to candidate num). Does not contain "my/our" in this count.
+
+        Example:
+            Input: 'Our town is big into high school football - our quarterback just left to go play for Clemson. Oh, that is my town.'
+            Output:
+                [{'extracted_text': 'Our town is big into high school football - our quarterback',
+                'pronoun': ('Our', (0, 3))}, ...]
+        """
+        if any(
+            [True for p in self.pronouns if p in utt.lower()]
+        ):  # If at least one pronoun is in utt.lower()
+            ret = []
+        else:  # If no pronouns are in utt.lower()
+            return []
+
+        try:  # if tokenizer version is 0.10.3 etc where pre_tokenize_str is available
+            ws = pre_tokenizer.pre_tokenize_str(
+                utt
+            )  # E.g., [('Our', (0, 3)), ('town', (4, 8)), ...]
+        except:  # if 0.8.1.rc2 etc where pre_tokenizer_str is NOT available
+            ws = pre_tokenizer.pre_tokenize(utt)
+        for i, (w, _) in enumerate(ws):
+            if w.lower() in self.pronouns:
+                n_options = min(
+                    max_candidate_num, len(ws[i:]) - 1
+                )  # `len(ws[i:])` contains "my/our" so have to operate -1
+                text_w_pronoun = utt[
+                    ws[i][1][0] : ws[i + n_options][1][1]
+                ]  # E.g., 'our quarterback just ...' # `ws[i][1][0]`: start position. `ws[i+n+2][1][1]`: end position.
+                ret.append({"pronoun": ws[i], "extracted_text": text_w_pronoun})
+        return ret
+
+    def pem_detector(self, utt):
+        """Mention detection for personal entities
+
+        Args:
+            utt (str): Input utterance.
+                E.g., 'I agree. One of my favorite forms of science fiction is anything related to time travel! I find it fascinating.'
+
+        Returns:
+            list of detected personal entity mentions.
+                E.g., ['my favorite forms of science fiction']
+
+        """
+        _dc_list = self._extract_text_with_pronoun(
+            utt
+        )  # E.g., [{'extracted_text': 'Our town is big into high ...', 'pronoun': ('Our', (0, 3))}, ...]
+        if len(_dc_list) == 0:
+            return []
+        else:
+            texts_w_pronoun = [
+                _dc["extracted_text"] for _dc in _dc_list
+            ]  # E.g., ['Our town is big into ...', 'My dog loves human food!']
+
+        ret = []
+        for text in texts_w_pronoun:  # For each extracted text
+            doc = nlp(text)
+            ment = ""
+            end_pos = 0  # start_pos is always 0
+            for i, token in enumerate(doc):
+                # print(token.pos_, token.text)
+                if i == 0:
+                    assert (
+                        token.text.lower() in self.pronouns
+                    ), f"{token.text} does not start with 'my' or 'our'"
+                    end_pos = token.idx + len(token.text)  # update end_pos
+                else:  # i > 0
+                    if token.pos_ in [
+                        "ADJ",
+                        "NOUN",
+                        "PROPN",
+                        "NUM",
+                        "PART",
+                    ] or token.text.lower() in [
+                        "of",
+                        "in",
+                        "the",
+                        "a",
+                        "an",
+                    ]:
+                        end_pos = token.idx + len(token.text)  # update end_pos
+                    else:
+                        break
+            ment = text[:end_pos]
+
+            ###### Post process #######
+            # if end with " of " then remove it
+            for drop_term in [" of", " in", " to"]:
+                ment = (
+                    ment[: -(len(drop_term) - 1)] if ment.endswith(drop_term) else ment
+                )
+
+            if (
+                len(ment) > min([len(pron) for pron in self.pronouns]) + 1
+            ):  # Want to ignore the case: "My "
+                ret.append(ment.strip())
+
+            # 220406: TMP error check
+            # TODO: Check this part whether it is needed or not
+            assert len(ment) != "our ", f'Should fix "if len(ment)>len(CLUE)+1" part.'
+
+            # Sanity check
+            for ment in ret:
+                assert ment in utt, f"{ment} is not in {utt}"
+
+        # Change to REL format [start_pos, length, mention]
+        ret = [
+            [start_pos, len(m), m]
+            for m in ret
+            for start_pos, _ in self.preprocess.get_span(
+                m, utt, flag_start_end_span_representation=False
+            )
+        ]
+
+        return ret
+
+
+class EEMD:
+    """Find corresponding explicit entity mention using s2e-coref-based method"""
+
+    def __init__(self, s2e_pe_model):
+        self.conf = self.Config(s2e_pe_model=s2e_pe_model)
+        self.model = self._read_model()
+
+    class Config:
+        """Inner class for config"""
+
+        def __init__(self, s2e_pe_model):
+            self.max_seq_length = 4096
+            self.model_name_or_path = s2e_pe_model
+            self.max_total_seq_len = 4096
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            # self.device = torch.device("cpu") # TMP
+
+            # Config for NN model
+            # Params are from: https://github.com/yuvalkirstain/s2e-coref/blob/main/cli.py
+            self.max_span_length = 30
+            self.top_lambda = 0.4
+            self.ffnn_size = 3072
+            self.normalise_loss = False
+            self.dropout_prob = 0.3
+
+    def _read_model(self):
+        config_class = LongformerConfig
+        base_model_prefix = "longformer"
+
+        transformer_config = AutoConfig.from_pretrained(
+            self.conf.model_name_or_path
+        )  # , cache_dir=args.cache_dir)
+
+        S2E.config_class = config_class
+        S2E.base_model_prefix = base_model_prefix
+
+        model = S2E.from_pretrained(
+            self.conf.model_name_or_path, config=transformer_config, args=self.conf
+        )
+
+        model.to(self.conf.device)
+
+        return model
+
+    def get_scores(self, input_data):
+        """Calculate the score of each mention pair
+        Args:
+            input_data (dict): Input data.
+                E.g., {'clusters': [], # Not used for inference
+                       'doc_key': 'tmp', # Not used for inference
+                       'mentions': [[2, 3], [77, 78]], # Detected concept and NE mentions
+                       'pems': [[67, 72]], # Detected personal entity mention. Only one PEM is allowed now.
+                       'sentences': [['I', 'think', 'science', 'fiction', 'is', ...]], # tokenized sentences using tokenizers.pre_tokenizers
+                       'speakers': [['USER', 'USER', 'USER', 'USER', 'USER', ...]], # speaker information
+                       'text': None
+                       }
+
+        Returns:
+            The scores for each mention pair. The pairs which does not have any PEM are removed in the later post-processing.
+                E.g.,
+                    [{'doc_key': 'tmp',
+                    'span_token_anaphora': (67, 72),
+                    'span_token_antecedent': (2, 3), ...]
+        """
+        assert (
+            type(input_data) == dict
+        ), f"input_data should be a dict, but got {type(input_data)}"
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.conf.model_name_or_path, use_fast=False
+        )
+        # `use_fast=False` should be supecified for v4.18.0 (do not need to do this for v3.3.1)
+        # See: https://github.com/huggingface/transformers/issues/10297#issuecomment-812548803
+
+        eval_dataset = data.get_dataset(tokenizer, input_data, self.conf)
+
+        eval_dataloader = BucketBatchSampler(
+            eval_dataset,
+            max_total_seq_len=self.conf.max_total_seq_len,
+            batch_size_1=True,
+        )
+
+        assert len(eval_dataloader) == 1, f"Currently, only 1 batch is supported"
+        for i, ((doc_key, subtoken_maps), batch) in enumerate(eval_dataloader):
+            # NOTE: subtoken_maps should NOT be used to map word -> subtoken!!!
+            #       The original name of subtoken_maps is `end_token_idx_to_word_idx`, meaning this is intended to map subtoken (end token) -> word.
+            # NOTE: Currently, this code only supports only one example at a time, however, for the futurework, we keep this for loop here.
+            batch = tuple(tensor.to(self.conf.device) for tensor in batch)
+            input_ids, attention_mask, gold_clusters = batch
+
+            with torch.no_grad():
+                _, scores = self.model(
+                    input_ids=input_ids,  # This calls __call__ in module.py in PyTorch, and it calls S2E.forward().
+                    attention_mask=attention_mask,
+                    gold_clusters=gold_clusters,
+                    return_all_outputs=True,
+                    subtoken_map=subtoken_maps,  # HJ
+                    pem_eem_subtokenspan=(
+                        sorted(
+                            eval_dataset.dockey2eems_subtokenspan[doc_key]
+                            + eval_dataset.dockey2pems_subtokenspan[doc_key]
+                        )
+                    ),  # HJ
+                    doc_key=doc_key,
+                )  # HJ: 220221
+
+        return scores
diff --git a/src/REL/crel/s2e_pe/pe_data.py b/src/REL/crel/s2e_pe/pe_data.py
new file mode 100644
index 0000000..250243d
--- /dev/null
+++ b/src/REL/crel/s2e_pe/pe_data.py
@@ -0,0 +1,468 @@
+from tokenizers.pre_tokenizers import Whitespace
+
+pre_tokenizer = Whitespace()
+
+TMP_DOC_ID = "tmp"  # temporary doc id
+
+
+class PreProcess:
+    """Create input for PE Linking module"""
+
+    def _error_check(self, conv):
+        assert type(conv) == list
+        assert len(conv) > 0, f"conv should be a list of dicts, but got {conv}"
+        for turn in conv:
+            assert type(turn) == dict, f"conv should be a list of dicts, but got {turn}"
+            assert set(turn.keys()) == {
+                "speaker",
+                "utterance",
+                "mentions",
+                "pems",
+            } or set(turn.keys()) == {
+                "speaker",
+                "utterance",
+            }, f"Each turn should have either [speaker, utterance, mentions, pems] keys for USER or [speaker, utterance] keys for SYSTEM. If there is no pems or mentions, then set them to empty list."
+            assert turn["speaker"] in [
+                "USER",
+                "SYSTEM",
+            ], f'The speaker should be either USER or SYSTEM, but got {turn["speaker"]}'
+            assert (
+                type(turn["utterance"]) == str
+            ), f'The utterance should be a string, but got {turn["utterance"]}'
+            if turn["speaker"] == "USER":
+                assert (
+                    type(turn["mentions"]) == list
+                ), f'The mentions should be a list, but got {turn["mentions"]}'
+                assert (
+                    type(turn["pems"]) == list
+                ), f'The pems should be a list, but got {turn["pems"]}'
+
+        # Check there are only one pem per conv
+        pems = [pem for turn in conv if "pems" in turn for pem in turn["pems"]]
+        assert (
+            len(pems) == 1
+        ), f"Current implementation only supports one pem per input conv. If there are multiple PEM, then split them into multiple conv."  # This is also a TODO for the future
+
+    def get_span(self, ment, text, flag_start_end_span_representation=True):
+        """Get (start, end) span of a mention (inc. PEM) in a text
+
+        Args:
+            ment (str): E.g., 'Paris'
+            text (str): E.g., 'Paris. The football club Paris Saint-Germain and the rugby union club Stade Français are based in Paris.'
+
+        Returns: mention spans
+            if flag_start_end_span_representation==True:
+                E.g.,  [(0, 5), (25, 30), (98, 103)]
+            if flag_start_end_span_representation==False:
+                E.g.,  [(0, 5), (25, 5), (98, 5)]
+
+        Note:
+            - re.finditer is NOT used since it takes regex pattern (not string) as input and fails for the patterns such as:
+                text = 'You you dance? I love cuban Salsa but also like other types as well. dance-dance.'
+                ment = 'dance? '
+        """
+        assert ment in text, f"The mention {ment} is not in the text {text}"
+        spans = []  # [(start_pos, length), ...]
+        offset = 0
+        while True:
+            try:
+                start_pos = text.index(ment, offset)
+                spans.append((start_pos, len(ment)))
+                offset = start_pos + len(ment)
+            except:
+                break
+
+        if (
+            flag_start_end_span_representation
+        ):  # (start_pos, length) --> (start_pos, end_pos)
+            spans = [(s, l + s) for s, l in spans]  # pos_end = pos_start + length
+
+        return spans
+
+    def _token_belongs_to_mention(
+        self, m_spans: list, t_span: tuple, utt: str, print_warning=False
+    ) -> bool:
+        """Check whether token span is in ment span(s) or not
+
+        Args:
+            m_spans: e.g., [(0, 4), (10, 14), (2,4)]
+            t_span: e.g., (1, 3)
+        """
+
+        def _error_check(m_spans, t_span):
+            assert len(t_span) == 2
+            assert (
+                t_span[1] > t_span[0]
+            )  # Note that span must be (start_ind, end_ind), NOT the REL style output of (start_ind, length)
+            assert any(
+                [True if m_span[1] > m_span[0] else False for m_span in m_spans]
+            )  # The same as above
+
+        _error_check(m_spans, t_span)
+
+        # Main
+        for m_span in m_spans:
+
+            # if token span is out of mention span (i.e., does not have any overlaps), then go to next
+            t_out_m = (t_span[1] <= m_span[0]) or (m_span[1] <= t_span[0])
+            if t_out_m:
+                continue
+
+            # Check whether (1) token is in mention, (2) mention is in token, or (3) token partially overlaps with mention
+            t_in_m = (m_span[0] <= t_span[0]) and (t_span[1] <= m_span[1])
+            m_in_t = (t_span[0] <= m_span[0]) and (
+                m_span[1] <= t_span[1]
+            )  # To deal with the case of ment:"physicians" ent:"Physician"
+            t_ol_m = not (t_in_m or t_out_m)  # token overlaps with mention
+
+            if t_in_m or m_in_t:  # Case 1 or 2
+                return True
+            elif t_ol_m:  # Case 3
+                if print_warning:
+                    print(
+                        f"WARNING: There is overlaps between mention and word spans. t_span:{t_span} ({utt[t_span[0]:t_span[1]]}), m_span:{m_span} ({utt[m_span[0]:m_span[1]]})"
+                    )
+                # NOTE: Treat this token as not belonging to the mention
+
+        return False
+
+    def _tokens_info(
+        self,
+        text: str,
+        speaker: str,
+        ments: list,
+        pems: list,
+    ):
+        """Append information for each token
+        The information includes:
+            - span: (start_pos, end_pos) of the token, acquired from the pre_tokenizers
+            - mention: what mention the token belongs to (or if not in any mention, None)
+            - pem: same as mention, but for pems
+            - speaker: the speaker of the utterance, either "USER" or "SYSTEM"
+
+        Args:
+            text: the text of the utterance
+            speaker: the speaker of the utterance, either "USER" or "SYSTEM"
+            ments: list of mentions
+            pems: list of pems
+
+        Returns: list of tokens with information
+            E.g., [{'token': 'Blue',
+                    'span': (0, 4),
+                    'mention': 'Blue',
+                    'pem': None,
+                    'speaker': 'USER'}, ...]
+        """
+        ments = list(sorted(ments, key=len))  # to perform longest match
+        tokens_conv = []
+
+        try:  # if tokenizer version is 0.10.3 etc where pre_tokenize_str is available
+            tokens_per_utt = pre_tokenizer.pre_tokenize(text)
+        except:  # if 0.8.1.rc2 etc where pre_tokenizer_str is NOT available
+            tokens_per_utt = pre_tokenizer.pre_tokenize_str(text)
+
+        ment2span = {ment: self.get_span(ment, text) for ment in ments}  # mention spans
+        pem2span = {pem: self.get_span(pem, text) for pem in pems}  # pem spans
+        # NOTE: get_span() cannot consider the case of the same surface mention occurring multiple times in the same utterance, only one of which is detected by the MD module.
+        # This is because the function detects ALL of the spans for the given surface mention.
+        # However, this case is extremely rare thus here the code ignores this case.
+
+        for tkn_span in tokens_per_utt:
+            tkn = tkn_span[0]
+            span = tkn_span[1]
+            ment_out, pem_out = None, None  # Initialize
+
+            # First check if token is in any PEMs
+            for pem in pems:
+                if self._token_belongs_to_mention(pem2span[pem], span, text):
+                    pem_out = pem
+                    break
+
+            # If token is not in any pem, then check if it is in any mention
+            if pem_out is None:
+                for ment in ments:
+                    if self._token_belongs_to_mention(ment2span[ment], span, text):
+                        ment_out = ment
+                        break
+
+            tokens_conv.append(
+                {
+                    "token": tkn,
+                    "span": span,
+                    "mention": ment_out,
+                    "pem": pem_out,
+                    "speaker": speaker,
+                }
+            )
+        return tokens_conv
+
+    def get_tokens_with_info(self, conv):
+        """Get tokens with information of:
+            - span: (start_pos, end_pos) of the token, acquired from the pre_tokenizers
+            - mention: what mention the token belongs to (or if not in any mention, None)
+            - pem: same as mention, but for pems
+            - speaker: the speaker of the utterance, either "USER" or "SYSTEM"
+
+        Args:
+            conv: a conversation
+                E.g.,
+                    {"speaker": "USER",
+                    "utterance": "I agree. One of my favorite forms of science fiction is anything related to time travel! I find it fascinating.",
+                    "mentions": ["science fiction", "time travel", ],
+                    "pems": ["my favorite forms of science fiction", ],},
+        """
+        self._error_check(conv)
+        ret = []
+        for turn_num, turn in enumerate(conv):
+            speaker = turn["speaker"]
+            if speaker == "USER":
+                ments = turn["mentions"] if "mentions" in turn else []
+                pems = turn["pems"] if "pems" in turn else []
+            elif speaker == "SYSTEM":
+                ments = []
+                pems = []
+            else:
+                raise ValueError(
+                    f'Unknown speaker: {speaker}. Speaker must be either "USER" or "SYSTEM".'
+                )
+            tkn_info = self._tokens_info(turn["utterance"], speaker, ments, pems)
+            for elem in tkn_info:
+                elem["turn_number"] = turn_num
+            ret += tkn_info
+        return ret
+
+    def _get_token_level_span(self, token_info: list, key_for_ment: str):
+        """Get token-level spans for mentions and pems
+        Token-level span is the input for s2e
+
+        Args:
+            token_info (list): E.g., [{'token': 'Blue',  'span': (0, 4),  'corresponding_ment': 'Blue',  'speaker': 'USER',  'turn_number': 0}, ...]
+            key_ment_or_pem (str): 'mention' or 'pem'
+
+        Returns:
+            E.g., [[2, 5], [8, 9], [0, 3]]
+        """
+        # Error check
+        assert key_for_ment in [
+            "mention",
+            "pem",
+        ]  # key_for_ment should be mention or pem
+
+        pem_and_eem = []
+        start_pos, end_pos = None, None
+        for i in range(len(token_info)):
+            prev_ment = token_info[i - 1][key_for_ment] if i > 0 else None
+            curr_ment = token_info[i][key_for_ment]
+            futr_ment = (
+                token_info[i + 1][key_for_ment] if (i + 1) < len(token_info) else None
+            )
+
+            if (prev_ment != curr_ment) and (curr_ment != None):  # mention start
+                if start_pos == None:  # Error check
+                    start_pos = i
+                else:
+                    raise ValueError("pos should be None to assign the value")
+            if (futr_ment != curr_ment) and (curr_ment != None):
+                # print(curr_ment,start_pos,end_pos)
+                if end_pos == None:  # Error check
+                    end_pos = i
+                else:
+                    raise ValueError("pos should be None to assign the value")
+
+            # print(prev_ment,curr_ment,futr_ment,'\tSTART_END_POS:',start_pos,end_pos)
+
+            if (start_pos != None) and (end_pos != None):
+                # print('curr_ment:',curr_ment)
+                pem_and_eem.append([start_pos, end_pos])
+                start_pos, end_pos = None, None  # Initialize
+
+        return pem_and_eem
+
+    def get_input_of_pe_linking(self, token_info):
+        """Get the input of PE Linking module
+
+        Args:
+            token_info (list): list of dict where tokens and their information is stored. Output of get_tokens_with_info()
+                E.g.,
+                    [{'token': 'I', # token
+                    'span': (0, 1), # (start_pos, end_pos) of the token
+                    'mention': None, # what mention the token belongs to (or if not in any mention, None)
+                    'pem': None, # same as mention, but for pems
+                    'speaker': 'USER', # the speaker of the utterance, either "USER" or "SYSTEM"
+                    'turn_number': 0},  # turn number of the utterance which the token belongs to (0-based)
+                    ...]
+
+        Returns:
+            Input of PE Linking module (the same input as s2e-coref)
+        """
+        ret = []
+        pem_spans = self._get_token_level_span(token_info, "pem")
+        # TODO: This is redundant and not efficient.
+        # Instead `pem_spans` can be acquired by pos_to_token mapping, which maps char_position --> token_position
+        # But I am too lazy to implement this code now. (Especially since it affects all other functions and data flows)
+
+        for (
+            pem_span
+        ) in (
+            pem_spans
+        ):  # [[142, 143], [256, 258], [309, 310]]. Note that this is token-level, not char-level
+            (start, end) = pem_span
+            turn_num = token_info[start]["turn_number"]
+            assert (
+                turn_num == token_info[end]["turn_number"]
+            ), f"Start token and end token should have the same turn_number. start: {start}, end: {end}"
+            tokens_until_current_turn = [
+                e for e in token_info if e["turn_number"] <= turn_num
+            ]  # Extract tokens until the current turn
+
+            ret.append(
+                {
+                    "clusters": [],  # Not used for inference
+                    "doc_key": "tmp",  # Not used for inference
+                    "mentions": self._get_token_level_span(
+                        tokens_until_current_turn, "mention"
+                    ),  # Detected mention spans, where format is (start_token_ind, end_token_ind) # E.g., [[28, 43], [67, 78]]. TODO: The same as above todo
+                    "pems": [
+                        pem_span
+                    ],  # Detected personal entity mention span. The format is the same as mention spans # E.g., [[7, 43]]. NOTE: Currently our tool support only one mention at a time.
+                    "sentences": [
+                        [e["token"] for e in tokens_until_current_turn]
+                    ],  # Tokenized sentences. E.g., ['I', 'think', 'science', 'fiction', 'is', ...]
+                    "speakers": [
+                        [e["speaker"] for e in tokens_until_current_turn]
+                    ],  # Speaker information. E.g., ['SYSTEM', 'SYSTEM', ..., 'USER', 'USER', ...]
+                    "text": None,
+                }
+            )
+        return ret
+
+
+class PostProcess:
+    """Handle output of PE Linking module"""
+
+    def _get_ment2score(
+        self, doc_key: str, mentions: list, pems: list, scores: list, flag_print=False
+    ) -> dict:
+        """Get mention to score map
+
+        Args:
+            doc_key (str): E.g., 'dialind:1_turn:0_pem:My-favourite-type-of-cake'
+            mentions: E.g., [[6, 7], [12, 12], [14, 15]]
+            pems: E.g., [[0, 4]]
+            scores: The scores for all mention (inc. PE) pair
+                E.g., [{'doc_key': 'tmp', 'span_token_anaphora': [8, 8], 'span_token_antecedent': [0, 0], 'coref_logits': -66.80387115478516}, ...]
+
+        Returns:
+            {(6, 7): -42.52804183959961,
+            (12, 12): -83.429443359375,
+            (14, 15): -47.706520080566406}
+
+        """
+        assert all(
+            [isinstance(m, list) for m in mentions]
+        )  # Check both mentions and pems are 2d lists
+        assert all([isinstance(m, list) for m in pems])  # The same for pems
+        assert len(pems) == 1  # Check we only have one PEM
+        if doc_key not in [sj["doc_key"] for sj in scores]:  # 220403
+            if flag_print:
+                print(
+                    f"{doc_key} not in scores. It might be that EL tool could not detect any candidate EEMs for this PEM. Return empty dict."
+                )
+            return {}  # ment2score = {}
+
+        # Change all ments and pems to tuple to be able to compare
+        ment_tpl_list = [
+            tuple(m) for m in mentions
+        ]  # E.g., [(6, 7), (12, 12), (14, 15)]
+        pem_tpl = tuple(pems[0])  # E.g., (0, 4)
+
+        ment2score = {}
+        span_hist = set()  # This is used to error check
+        for sj in scores:
+            if sj["doc_key"] == doc_key:
+                # print(sj)
+                span_ano = tuple(sj["span_token_anaphora"])  # E.g., (6, 7)
+                span_ant = tuple(sj["span_token_antecedent"])  # E.g., (0, 4)
+                span_hist.add(span_ano)
+                span_hist.add(span_ant)
+
+                if (
+                    span_ano == pem_tpl and span_ant in ment_tpl_list
+                ):  # anaphora is the PEM case
+                    ment2score[span_ant] = sj["coref_logits"]
+                elif (
+                    span_ant == pem_tpl and span_ano in ment_tpl_list
+                ):  # antecedent is the PEM case
+                    ment2score[span_ano] = sj["coref_logits"]
+
+        # Check all ment_tpl_list and pem_tpl are in span_hist
+        assert sorted(ment_tpl_list + [pem_tpl]) == sorted(
+            list(span_hist)
+        ), f"mentions in score.json and pred.jsonl should be the same. {sorted(ment_tpl_list + [pem_tpl])} != {sorted(list(span_hist))}. doc_key: {doc_key}"
+        return ment2score
+
+    def _convert_to_mention_from_token(self, mention, comb_text):
+        """
+        Args:
+            mention (list): [start, end] (this is token-level (NOT subtoken-level))
+        """
+        start = mention[0]  # output['subtoken_map'][mention[0]]
+        end = mention[1] + 1  # output['subtoken_map'][mention[1]] + 1
+        mtext = "".join(" ".join(comb_text[start:end]).split(" ##"))
+        return mtext
+
+    def get_results(self, pel_input, conv, threshold, scores):
+        """Get PE Linking post-processed results
+
+        Args:
+            pel_input (dict): input for PE Linking module
+                E.g., {'clusters': [], # Not used for inference
+                       'doc_key': 'tmp', # Not used for inference
+                       'mentions': [[2, 3], [77, 78]], # Detected concept and NE mentions
+                       'pems': [[67, 72]], # Detected personal entity mention. Only one PEM is allowed now.
+                       'sentences': [['I', 'think', 'science', 'fiction', 'is', ...]], # tokenized sentences using tokenizers.pre_tokenizers
+                       'speakers': [['USER', 'USER', 'USER', 'USER', 'USER', ...]], # speaker information
+                       'text': None
+                       }
+            threshold: default 0
+            conv: The conversation input to preprocessing module (conversation before preprocessing)
+            scores: The scores for all mention (inc. PE) pair
+                E.g.,
+                    [{'doc_key': 'tmp',
+                    'span_token_anaphora': (67, 72), # This could be either a mention or a PEM
+                    'span_token_antecedent': (2, 3), # The same as above
+                    'coref_logits': -4.528693675994873}, # Output score from PE Linking module
+        Returns:
+            E.g.,
+                [{'personal_entity_mention': 'my favorite forms of science fiction',
+                'mention': 'time travel',
+                'score': 4.445976734161377}]
+
+        """
+        assert type(pel_input) == dict, f"pel_input should be a dict. {type(pel_input)}"
+        ments_span_tokenlevel = [m for m in pel_input["mentions"]]
+        pems_span_tokenlevel = [p for p in pel_input["pems"]]  # len(pems) == 1
+        assert (
+            len(pems_span_tokenlevel) == 1
+        ), f"len(pems_span_tokenlevel) should be 1. {len(pems_span_tokenlevel)}"
+
+        mspan2score = self._get_ment2score(
+            TMP_DOC_ID, ments_span_tokenlevel, pems_span_tokenlevel, scores
+        )  # Mention span to score
+        comb_text = pel_input["sentences"][
+            0
+        ]  # pel_input['sentences'] should have only one sentence
+
+        pred_peas = []
+
+        pem = [m for turn in conv if turn["speaker"] == "USER" for m in turn["pems"]][
+            0
+        ]  # Each conv has only one pem for current implementation
+        for ment_span_tokenlevel in pel_input["mentions"]:
+            score = mspan2score[tuple(ment_span_tokenlevel)]
+            ment = self._convert_to_mention_from_token(ment_span_tokenlevel, comb_text)
+            if score > threshold:
+                pred_peas.append(
+                    {"personal_entity_mention": pem, "mention": ment, "score": score}
+                )
+        return pred_peas
diff --git a/src/REL/crel/s2e_pe/utils.py b/src/REL/crel/s2e_pe/utils.py
new file mode 100644
index 0000000..7207b59
--- /dev/null
+++ b/src/REL/crel/s2e_pe/utils.py
@@ -0,0 +1,134 @@
+import json
+import os
+from datetime import datetime
+from time import time
+
+import numpy as np
+
+import torch
+
+from .consts import NULL_ID_FOR_COREF
+
+
+def flatten_list_of_lists(lst):
+    return [elem for sublst in lst for elem in sublst]
+
+
+def extract_clusters(gold_clusters):
+    gold_clusters = [
+        tuple(tuple(m) for m in gc if NULL_ID_FOR_COREF not in m)
+        for gc in gold_clusters.tolist()
+    ]
+    return gold_clusters
+
+
+def extract_mentions_to_predicted_clusters_from_clusters(gold_clusters):
+    mention_to_gold = {}
+    for gc in gold_clusters:
+        for mention in gc:
+            mention_to_gold[tuple(mention)] = gc
+    return mention_to_gold
+
+
+def extract_clusters_for_decode(mention_to_antecedent, pems_subtoken):
+    """
+    Args:
+        pems (list): E.g., [(2,3), (8,11), ...]
+    """
+
+    mention_to_antecedent = sorted(mention_to_antecedent)
+    mention_to_cluster = {}
+    clusters = []
+    for mention, antecedent in mention_to_antecedent:
+        if (mention in pems_subtoken) or (antecedent in pems_subtoken):
+            if antecedent in mention_to_cluster:
+                cluster_idx = mention_to_cluster[antecedent]
+                clusters[cluster_idx].append(mention)
+                mention_to_cluster[mention] = cluster_idx
+
+            else:
+                cluster_idx = len(clusters)
+                mention_to_cluster[mention] = cluster_idx
+                mention_to_cluster[antecedent] = cluster_idx
+                clusters.append([antecedent, mention])
+    clusters = [tuple(cluster) for cluster in clusters]
+    return clusters, mention_to_cluster
+
+
+def ce_extract_clusters_for_decode_with_one_mention_per_pem(
+    starts, end_offsets, coref_logits, pems_subtoken, flag_use_threshold
+):
+    """
+
+    Args:
+        - flag_use_threshold:
+            True: Default. If PEM does not meet a threshold (default: 0), then all mentions are ignored. The threshold is stored in final element of each row of coref_logits.
+            False: Ignore threshold, pick the highest logit EEM for each PEM.
+    Updates:
+      - 220302: Created
+    """
+    if flag_use_threshold:
+        max_antecedents = np.argmax(
+            coref_logits, axis=1
+        ).tolist()  # HJ: 220225: mention_to_antecedents takes max score. We have at most two predicted EEMs (one is coreference is PEM case, and the other is antecedent is PEM case).
+    else:
+        max_antecedents = np.argmax(
+            coref_logits[:, :-1], axis=1
+        ).tolist()  # HJ: 220225: mention_to_antecedents takes max score. We have at most two predicted EEMs (one is coreference is PEM case, and the other is antecedent is PEM case).
+
+    # Create {(ment, antecedent): logits} dict
+    mention_antecedent_to_coreflogit_dict = {
+        (
+            (int(start), int(end)),
+            (int(starts[max_antecedent]), int(end_offsets[max_antecedent])),
+        ): logit[max_antecedent]
+        for start, end, max_antecedent, logit in zip(
+            starts, end_offsets, max_antecedents, coref_logits
+        )
+        if max_antecedent < len(starts)
+    }
+    # 220403: Drop if key has the same start and end pos for anaphora and antecedent
+    mention_antecedent_to_coreflogit_dict = {
+        k: v for k, v in mention_antecedent_to_coreflogit_dict.items() if k[0] != k[1]
+    }
+    if len(mention_antecedent_to_coreflogit_dict) == 0:
+        return []
+
+    # Select the ment-ant pair containing the PEM
+
+    mention_antecedent_to_coreflogit_dict_with_pem = {
+        (m, a): logit
+        for (m, a), logit in mention_antecedent_to_coreflogit_dict.items()
+        if (m in pems_subtoken) or (a in pems_subtoken)
+    }
+    if len(mention_antecedent_to_coreflogit_dict_with_pem) == 0:
+        return []
+
+    # Select the max score
+    _max_logit = max(mention_antecedent_to_coreflogit_dict_with_pem.values())
+    if flag_use_threshold and (_max_logit <= 0):
+        print(f"WARNING: _max_logit = {_max_logit}")
+    # _max_logit = _max_logit if _max_logit > 0 else 0 # HJ: 220302: If we set a threshold, then this does not work.
+    assert (
+        coref_logits[-1][-1] == 0
+    ), f"The threshold should be 0. If you set your threshold, then the code above should be fixed."
+    # Select the pair with max score
+    mention_to_antecedent_max_pem = {
+        ((m[0], m[1]), (a[0], a[1]))
+        for (m, a), logit in mention_antecedent_to_coreflogit_dict_with_pem.items()
+        if logit == _max_logit
+    }
+    assert (
+        len(mention_to_antecedent_max_pem) <= 1
+    ), f"Two or more mentions have the same max score: {mention_to_antecedent_max_pem}"
+
+    predicted_clusters, _ = extract_clusters_for_decode(
+        mention_to_antecedent_max_pem, pems_subtoken
+    )  # TODO: 220302: Using `extract_clusters_for_decode` here is redundant.
+    return predicted_clusters
+
+
+def mask_tensor(t, mask):
+    t = t + ((1.0 - mask.float()) * -10000.0)
+    t = torch.clamp(t, min=-10000.0, max=10000.0)
+    return t
diff --git a/src/REL/db/base.py b/src/REL/db/base.py
index 8eec44d..2526946 100644
--- a/src/REL/db/base.py
+++ b/src/REL/db/base.py
@@ -40,7 +40,7 @@ def initialize_db(self, fname, table_name, columns):
             db (sqlite3.Connection): a SQLite3 database with an embeddings table.
         """
         # open database in autocommit mode by setting isolation_level to None.
-        db = sqlite3.connect(fname, isolation_level=None)
+        db = sqlite3.connect(fname, isolation_level=None, check_same_thread=False)
 
         q = "create table if not exists {}(word text primary key, {})".format(
             table_name, ", ".join(["{} {}".format(k, v) for k, v in columns.items()])
diff --git a/src/REL/response_model.py b/src/REL/response_model.py
new file mode 100644
index 0000000..b41da49
--- /dev/null
+++ b/src/REL/response_model.py
@@ -0,0 +1,65 @@
+from REL.entity_disambiguation import EntityDisambiguation
+from REL.ner import load_flair_ner
+from flair.models import SequenceTagger
+from REL.mention_detection import MentionDetection
+from REL.utils import process_results
+
+
+class ResponseModel:
+    API_DOC = "API_DOC"
+
+    def __init__(self, base_url, wiki_version, model, tagger_ner=None):
+        self.model = model
+        self.tagger_ner = tagger_ner
+
+        self.base_url = base_url
+        self.wiki_version = wiki_version
+
+        self.custom_ner = not isinstance(tagger_ner, SequenceTagger)
+        self.mention_detection = MentionDetection(base_url, wiki_version)
+
+    def generate_response(self,
+                          *,
+                          text: list,
+                          spans: list,
+                          ):
+        """
+        Generates response for API. Can be either ED only or EL, meaning end-to-end.
+
+        :return: list of tuples for each entity found.
+        """
+
+        if len(text) == 0:
+            return []
+        
+        processed = {self.API_DOC: [text, spans]}
+
+        if len(spans) > 0:
+            # ED.
+            mentions_dataset, total_ment = self.mention_detection.format_spans(
+                processed
+            )
+        else:
+            # EL
+            mentions_dataset, total_ment = self.mention_detection.find_mentions(
+                processed, self.tagger_ner
+            )
+
+        # Disambiguation
+        predictions, timing = self.model.predict(mentions_dataset)
+
+        include_offset = (len(spans) == 0) and not self.custom_ner
+
+        # Process result.
+        result = process_results(
+            mentions_dataset,
+            predictions,
+            processed,
+            include_offset=include_offset,
+        )
+
+        # Singular document.
+        if len(result) > 0:
+            return [*result.values()][0]
+
+        return []
\ No newline at end of file
diff --git a/src/REL/server.py b/src/REL/server.py
index d26d6a9..b321ad2 100644
--- a/src/REL/server.py
+++ b/src/REL/server.py
@@ -1,152 +1,54 @@
-import json
-from http.server import BaseHTTPRequestHandler
-
-from flair.models import SequenceTagger
-
-from REL.mention_detection import MentionDetection
-from REL.utils import process_results
-
-API_DOC = "API_DOC"
-
-
-
-def make_handler(base_url, wiki_version, model, tagger_ner):
-    """
-    Class/function combination that is used to setup an API that can be used for e.g. GERBIL evaluation.
-    """
-    class GetHandler(BaseHTTPRequestHandler):
-        def __init__(self, *args, **kwargs):
-            self.model = model
-            self.tagger_ner = tagger_ner
-
-            self.base_url = base_url
-            self.wiki_version = wiki_version
-
-            self.custom_ner = not isinstance(tagger_ner, SequenceTagger)
-            self.mention_detection = MentionDetection(base_url, wiki_version)
-
-            super().__init__(*args, **kwargs)
-
-        def do_GET(self):
-            self.send_response(200)
-            self.end_headers()
-            self.wfile.write(
-                bytes(
-                    json.dumps(
-                        {
-                            "schemaVersion": 1,
-                            "label": "status",
-                            "message": "up",
-                            "color": "green",
-                        }
-                    ),
-                    "utf-8",
-                )
-            )
-            return
-
-        def do_HEAD(self):
-            # send bad request response code
-            self.send_response(400)
-            self.end_headers()
-            self.wfile.write(bytes(json.dumps([]), "utf-8"))
-            return
-
-        def do_POST(self):
-            """
-            Returns response.
-
-            :return:
-            """
-            try:
-                content_length = int(self.headers["Content-Length"])
-                post_data = self.rfile.read(content_length)
-                self.send_response(200)
-                self.end_headers()
-
-                text, spans = self.read_json(post_data)
-                response = self.generate_response(text, spans)
-
-                self.wfile.write(bytes(json.dumps(response), "utf-8"))
-            except Exception as e:
-                print(f"Encountered exception: {repr(e)}")
-                self.send_response(400)
-                self.end_headers()
-                self.wfile.write(bytes(json.dumps([]), "utf-8"))
-            return
-
-        def read_json(self, post_data):
-            """
-            Reads input JSON message.
-
-            :return: document text and spans.
-            """
-
-            data = json.loads(post_data.decode("utf-8"))
-            text = data["text"]
-            text = text.replace("&amp;", "&")
-
-            # GERBIL sends dictionary, users send list of lists.
-            if "spans" in data:
-                try:
-                    spans = [list(d.values()) for d in data["spans"]]
-                except Exception:
-                    spans = data["spans"]
-                    pass
-            else:
-                spans = []
-
-            return text, spans
-
-        def generate_response(self, text, spans):
-            """
-            Generates response for API. Can be either ED only or EL, meaning end-to-end.
-
-            :return: list of tuples for each entity found.
-            """
-
-            if len(text) == 0:
-                return []
-
-            if len(spans) > 0:
-                # ED.
-                processed = {API_DOC: [text, spans]}
-                mentions_dataset, total_ment = self.mention_detection.format_spans(
-                    processed
-                )
-            else:
-                # EL
-                processed = {API_DOC: [text, spans]}
-                mentions_dataset, total_ment = self.mention_detection.find_mentions(
-                    processed, self.tagger_ner
-                )
-
-            # Disambiguation
-            predictions, timing = self.model.predict(mentions_dataset)
-
-            # Process result.
-            result = process_results(
-                mentions_dataset,
-                predictions,
-                processed,
-                include_offset=False if ((len(spans) > 0) or self.custom_ner) else True,
-            )
-
-            # Singular document.
-            if len(result) > 0:
-                return [*result.values()][0]
-
-            return []
-
-    return GetHandler
+from REL.response_model import ResponseModel
+
+from fastapi import FastAPI
+from pydantic import BaseModel, Field
+from typing import List, Optional, Literal
+
+app = FastAPI()
+
+@app.get("/")
+def root():
+    """Returns server status."""
+    return {
+        "schemaVersion": 1,
+        "label": "status",
+        "message": "up",
+        "color": "green",
+    }
+
+
+class EntityConfig(BaseModel):
+    text: str = Field(..., description="Text for entity linking or disambiguation.")
+    spans: List[str] = Field(..., description="Spans for entity disambiguation.")
+
+
+@app.post("/")
+def root(config: EntityConfig):
+    """Submit your text here for entity disambiguation or linking."""
+    response = handler.generate_response(text=config.text, spans=config.spans)
+    return response
+
+
+class ConversationTurn(BaseModel):
+    speaker: Literal["USER", "SYSTEM"] = Field(..., description="Speaker for this turn.")
+    utterance: str = Field(..., description="Input utterance.")
+
+
+class ConversationConfig(BaseModel):
+    text: List[ConversationTurn] = Field(..., description="Conversation as list of turns between two speakers.")
+
+
+@app.post("/conversation/")
+def conversation(config: ConversationConfig):
+    """Submit your text here for conversational entity linking."""
+    text = config.dict()['text']
+    response = conv_handler.annotate(text)
+    return response
 
 
 if __name__ == "__main__":
     import argparse
-    from http.server import HTTPServer
-
-    from REL.entity_disambiguation import EntityDisambiguation
-    from REL.ner import load_flair_ner
+    import uvicorn
 
     p = argparse.ArgumentParser()
     p.add_argument("base_url")
@@ -157,18 +59,17 @@ def generate_response(self, text, spans):
     p.add_argument("--port", "-p", default=5555, type=int)
     args = p.parse_args()
 
+    from REL.crel.conv_el import ConvEL
+    from REL.entity_disambiguation import EntityDisambiguation
+    from REL.ner import load_flair_ner
+
     ner_model = load_flair_ner(args.ner_model)
     ed_model = EntityDisambiguation(
         args.base_url, args.wiki_version, {"mode": "eval", "model_path": args.ed_model}
     )
-    server_address = (args.bind, args.port)
-    server = HTTPServer(
-        server_address,
-        make_handler(args.base_url, args.wiki_version, ed_model, ner_model),
-    )
 
-    try:
-        print("Ready for listening.")
-        server.serve_forever()
-    except KeyboardInterrupt:
-        exit(0)
+    handler = ResponseModel(args.base_url, args.wiki_version, ed_model, ner_model)
+
+    conv_handler = ConvEL(args.base_url, args.wiki_version, ed_model=ed_model)
+
+    uvicorn.run(app, port=args.port, host=args.bind)


Test patch:
diff --git a/scripts/efficiency_test.py b/scripts/efficiency_test.py
index fc66f54..b903712 100644
--- a/scripts/efficiency_test.py
+++ b/scripts/efficiency_test.py
@@ -1,12 +1,15 @@
 import numpy as np
 import requests
+import os
 
 from REL.training_datasets import TrainingEvaluationDatasets
 
 np.random.seed(seed=42)
 
-base_url = "/Users/vanhulsm/Desktop/projects/data/"
-wiki_version = "wiki_2014"
+base_url = os.environ.get("REL_BASE_URL")
+wiki_version = "wiki_2019"
+host = 'localhost'
+port = '5555'
 datasets = TrainingEvaluationDatasets(base_url, wiki_version).load()["aida_testB"]
 
 # random_docs = np.random.choice(list(datasets.keys()), 50)
@@ -40,7 +43,7 @@
             print(myjson)
 
             print("Output API:")
-            print(requests.post("http://192.168.178.11:1235", json=myjson).json())
+            print(requests.post(f"http://{host}:{port}", json=myjson).json())
             print("----------------------------")
 
 
diff --git a/scripts/test_server.py b/scripts/test_server.py
new file mode 100644
index 0000000..2953248
--- /dev/null
+++ b/scripts/test_server.py
@@ -0,0 +1,62 @@
+import os
+import requests
+
+# Script for testing the implementation of the conversational entity linking API
+#
+# To run:
+#
+#    python .\src\REL\server.py $REL_BASE_URL wiki_2019
+# or
+#    python .\src\REL\server.py $env:REL_BASE_URL wiki_2019
+#
+# Set $REL_BASE_URL to where your data are stored (`base_url`)
+#
+# These paths must exist:
+# - `$REL_BASE_URL/bert_conv`
+# - `$REL_BASE_URL/s2e_ast_onto `
+#
+# (see https://github.com/informagi/conversational-entity-linking-2022/tree/main/tool#step-1-download-models)
+#
+
+
+host = 'localhost'
+port = '5555'
+
+text1 = {
+    "text": "REL is a modular Entity Linking package that can both be integrated in existing pipelines or be used as an API.",
+    "spans": []
+}
+
+conv1 = {
+    "text" : [
+        {
+            "speaker":
+            "USER",
+            "utterance":
+            "I am allergic to tomatoes but we have a lot of famous Italian restaurants here in London.",
+        },
+        {
+            "speaker": "SYSTEM",
+            "utterance": "Some people are allergic to histamine in tomatoes.",
+        },
+        {
+            "speaker":
+            "USER",
+            "utterance":
+            "Talking of food, can you recommend me a restaurant in my city for our anniversary?",
+        },
+    ]
+}
+
+
+for endpoint, myjson in (
+        ('', text1), 
+        ('conversation/', conv1)
+    ):
+    print('Input API:')
+    print(myjson)
+    print()
+    print('Output API:')
+    print(requests.post(f"http://{host}:{port}/{endpoint}", json=myjson).json())
+    print('----------------------------')
+
diff --git a/tests/test_crel.py b/tests/test_crel.py
new file mode 100644
index 0000000..36fc551
--- /dev/null
+++ b/tests/test_crel.py
@@ -0,0 +1,94 @@
+import os
+import pytest
+from REL.crel.conv_el import ConvEL
+import yaml
+from pathlib import Path
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+
+@pytest.fixture()
+def cel():
+    return ConvEL(base_url=os.environ.get("REL_BASE_URL", "."))
+
+
+@pytest.mark.skipif(
+    os.getenv("GITHUB_ACTIONS")=='true', reason="No way of testing this on Github actions."
+)
+def test_conv1(cel):
+    example = [
+        {
+            "speaker": "USER",
+            "utterance": "I think science fiction is an amazing genre for anything. Future science, technology, time travel, FTL travel, they're all such interesting concepts.",
+        },
+        {
+            "speaker": "SYSTEM",
+            "utterance": "Awesome! I really love how sci-fi storytellers focus on political/social/philosophical issues that would still be around even in the future. Makes them relatable.",
+        },
+        {
+            "speaker": "USER",
+            "utterance": "I agree. One of my favorite forms of science fiction is anything related to time travel! I find it fascinating.",
+        },
+    ]
+
+    result = cel.annotate(example)
+    assert isinstance(result, list)
+
+    expected_annotations = [
+        [
+            [8, 15, "science fiction", "Science_fiction"],
+            [38, 5, "genre", "Genre_fiction"],
+            [74, 10, "technology", "Technology"],
+            [86, 11, "time travel", "Time_travel"],
+            [99, 10, "FTL travel", "Faster-than-light"],
+        ],
+        [
+            [37, 15, "science fiction", "Science_fiction"],
+            [76, 11, "time travel", "Time_travel"],
+            [16, 36, "my favorite forms of science fiction", "Time_travel"],
+        ],
+    ]
+
+    annotations = [res["annotations"] for res in result if res["speaker"] == "USER"]
+
+    assert annotations == expected_annotations
+
+
+@pytest.mark.skipif(
+    os.getenv("GITHUB_ACTIONS")=='true', reason="No way of testing this on Github actions."
+)
+def test_conv2(cel):
+    example = [
+        {
+            "speaker": "USER",
+            "utterance": "I am allergic to tomatoes but we have a lot of famous Italian restaurants here in London.",
+        },
+        {
+            "speaker": "SYSTEM",
+            "utterance": "Some people are allergic to histamine in tomatoes.",
+        },
+        {
+            "speaker": "USER",
+            "utterance": "Talking of food, can you recommend me a restaurant in my city for our anniversary?",
+        },
+    ]
+
+    result = cel.annotate(example)
+    assert isinstance(result, list)
+
+    annotations = [res["annotations"] for res in result if res["speaker"] == "USER"]
+
+    expected_annotations = [
+        [
+            [17, 8, "tomatoes", "Tomato"],
+            [54, 19, "Italian restaurants", "Italian_cuisine"],
+            [82, 6, "London", "London"],
+        ],
+        [
+            [11, 4, "food", "Food"],
+            [40, 10, "restaurant", "Restaurant"],
+            [54, 7, "my city", "London"],
+        ],
+    ]
+
+    assert annotations == expected_annotations