--- a
+++ b/src/training/bilstm.py
@@ -0,0 +1,646 @@
+# Base Dependencies
+# -----------------
+import numpy as np
+import time
+from copy import deepcopy
+from functools import partial
+from tqdm import tqdm
+from typing import Dict, Optional
+from pathlib import Path
+from os.path import join
+
+# Package Dependencies
+# --------------------
+from .base import BaseTrainer
+from .config import PLExperimentConfig, BaalExperimentConfig
+from .early_stopping import EarlyStopping
+from .utils import get_baal_query_strategy
+
+# Local Dependencies
+# -------------------
+from extensions.baal import (
+    MyModelWrapperBilstm,
+    MyActiveLearningDatasetBilstm,
+    MyActiveLearningLoop,
+)
+from extensions.torchmetrics import (
+    DetectionF1Score,
+    DetectionPrecision,
+    DetectionRecall,
+)
+from ml_models.bilstm import (
+    HasanModel,
+    EmbeddingConfig,
+    LSTMConfig,
+    RDEmbeddingConfig,
+)
+from re_datasets.bilstm_utils import pad_and_sort_batch, custom_collate
+from vocabulary import Vocabulary, read_list_from_file
+
+# 3rd-Party Dependencies
+# ----------------------
+import neptune
+import torch
+
+from baal.bayesian.dropout import patch_module
+from datasets import Dataset
+from torch.optim import Adam
+from torch.nn import CrossEntropyLoss, Module
+from torch.utils.data import DataLoader
+from torch.utils.data.sampler import BatchSampler, RandomSampler
+from torchmetrics import Accuracy
+from torchmetrics.classification import F1Score, Precision, Recall
+
+# Constants
+# ---------
+from constants import (
+    N2C2_VOCAB_PATH,
+    DDI_VOCAB_PATH,
+    N2C2_IOB_TAGS,
+    DDI_IOB_TAGS,
+    N2C2_RD_MAX,
+    DDI_RD_MAX,
+    RD_EMB_DIM,
+    IOB_EMB_DIM,
+    BIOWV_EMB_DIM,
+    POS_EMB_DIM,
+    DEP_EMB_DIM,
+    BIOWORD2VEC_PATH,
+    U_POS_TAGS,
+    DEP_TAGS,
+    BaalQueryStrategy,
+)
+from config import NEPTUNE_API_TOKEN, NEPTUNE_PROJECT
+
+
+class BilstmTrainer(BaseTrainer):
+    """Trainer for BiLSTM method."""
+
+    def __init__(
+        self,
+        dataset: str,
+        train_dataset: Dataset,
+        test_dataset: Dataset,
+        relation_type: Optional[str] = None,
+    ):
+        """
+        Args:
+            dataset (str): name of the dataset, e.g., "n2c2".
+            train_dataset (Dataset): train split of the dataset.
+            test_dataset (Dataset): test split of the dataset.
+            relation_type (str, optional): relation type.
+
+        Raises:
+            ValueError: if the name dataset provided is not supported
+        """
+        super().__init__(dataset, train_dataset, test_dataset, relation_type)
+
+        # vocabulary
+        self.vocab = self._init_vocab()
+
+        # transform datasets
+        self.transform = partial(
+            pad_and_sort_batch, padding_idx=self.vocab.pad_index, rd_max=self.RD_MAX
+        )
+
+    @property
+    def method_name(self) -> str:
+        return "bilstm"
+
+    @property
+    def method_name_pretty(self) -> str:
+        return "BiLSTM"
+
+    @property
+    def task(self) -> str:
+        if self.dataset == "n2c2":
+            task = "binary"
+        else:
+            task = "multiclass"
+        return task
+
+    @property
+    def model_class(self) -> str:
+        return HasanModel
+
+    @property
+    def RD_MAX(self) -> str:
+        if self.dataset == "n2c2":
+            rd_max = N2C2_RD_MAX
+        else:
+            rd_max = DDI_RD_MAX
+        return rd_max
+
+    @property
+    def IOB_TAGS(self) -> str:
+        if self.dataset == "n2c2":
+            iob_tags = N2C2_IOB_TAGS
+        else:
+            iob_tags = DDI_IOB_TAGS
+        return iob_tags
+
+    def _init_optimizer(self, model: Module):
+        return Adam(model.parameters(), lr=0.0001)
+
+    def _init_vocab(self):
+        """Loads the vocabulary of the dataset"""
+        if self.dataset == "n2c2":
+            vocab_path = N2C2_VOCAB_PATH
+        else:
+            vocab_path = DDI_VOCAB_PATH
+
+        return Vocabulary(read_list_from_file(vocab_path))
+    
+    def _init_model(self, patch: bool = False) -> HasanModel:
+        """Builds the BiLSTM model setting the right configuration for the chosen dataset"""
+        # word embedding configuration
+        biowv_config = EmbeddingConfig(
+            embedding_dim=BIOWV_EMB_DIM,
+            vocab_size=len(self.vocab),
+            emb_path=BIOWORD2VEC_PATH,
+            freeze=True,
+            padding_idx=self.vocab.pad_index,
+        )
+
+        # relative-distance embedding configuration
+        rd_config = RDEmbeddingConfig(
+            input_dim=self.RD_MAX, embedding_dim=RD_EMB_DIM, freeze=False
+        )
+
+        # IOB embedding configuration
+        iob_config = EmbeddingConfig(
+            embedding_dim=IOB_EMB_DIM, vocab_size=(len(self.IOB_TAGS) + 1), freeze=False
+        )
+
+        # Part-of-Speach tag embedding configuration
+        pos_config = EmbeddingConfig(
+            embedding_dim=POS_EMB_DIM, vocab_size=(len(U_POS_TAGS) + 1), freeze=False
+        )
+
+        dep_config = EmbeddingConfig(
+            embedding_dim=DEP_EMB_DIM, vocab_size=(len(DEP_TAGS) + 1), freeze=False
+        )
+
+        # BiLSTM configuration
+        lstm_config = LSTMConfig(
+            emb_size=(
+                BIOWV_EMB_DIM + 2 * RD_EMB_DIM + POS_EMB_DIM + DEP_EMB_DIM + IOB_EMB_DIM
+            )
+        )
+
+        model = self.model_class(
+            vocab=self.vocab,
+            lstm_config=lstm_config,
+            bioword2vec_config=biowv_config,
+            rd_config=rd_config,
+            pos_config=pos_config,
+            dep_config=dep_config,
+            iob_config=iob_config,
+            num_classes=self.num_classes,
+        )
+
+        if patch:
+            model = patch_module(model)
+
+        return model
+
+    def _reset_trainer(self):
+        self.train_dataset.reset_format()
+        self.test_dataset.reset_format()
+
+    def create_dataloader(self, dataset: Dataset, batch_size: int = 6) -> DataLoader:
+        """Creates a dataloader from a dataset with the adequate configuration
+
+        Args:
+            dataset (Dataset): dataset to load
+
+        Returns:
+            DataLoader: dataloader for the given dataset
+        """
+        dataset.set_transform(self.transform)
+
+        # create dataloader
+        sampler = BatchSampler(
+            RandomSampler(dataset), batch_size=batch_size, drop_last=False
+        )
+        dataloader = DataLoader(dataset, sampler=sampler, collate_fn=custom_collate)
+
+        return dataloader
+
+    def eval_model(
+        self,
+        model: Module,
+        dataloader: DataLoader,
+        criterion: Module,
+    ) -> Dict[str, float]:
+        """Evaluates the current model on the dev or test set
+
+        Args:
+            model (Module): model to use for evaluation.
+            dataloader (DataLoader): dataloader of evaluation dataset
+        Returns:
+            Dict: metrics including loss (`loss`), precision (`p`), recall (`r`) and F1-score (`f1`)
+        """
+
+        y_true = np.array([], dtype=np.int8)
+        y_pred = np.array([], dtype=np.int8)
+
+        val_loss = 0.0
+
+        with torch.no_grad():
+            for inputs, labels in dataloader:
+                # send (inputs, labels) to device
+                labels = labels.to(self.device)
+                for key, value in inputs.items():
+                    inputs[key] = value.to(self.device)
+
+                # calculate outputs
+                outputs = model(inputs)
+                loss = criterion(outputs, labels)
+                val_loss += len(inputs) * loss.item()
+
+                # calculate predictions
+                _, predicted = torch.max(outputs.data, 1)
+
+                # store labels and predictions
+                y_true = np.append(y_true, labels.cpu().detach().numpy())
+                y_pred = np.append(y_pred, predicted.cpu().detach().numpy())
+
+        metrics = self.compute_metrics(y_true, y_pred)
+        metrics["loss"] = val_loss / len(dataloader)
+
+        return metrics
+
+    def train_passive_learning(
+        self, config: PLExperimentConfig, verbose: bool = True, logging: bool = True
+    ):
+        """Trains the BiLSTM model using passive learning and early stopping
+
+        Args:
+            config (PLExperimentConfig): cofiguration
+            verbose (bool): determines if information is printed during training. Daults to True.
+            logging (bool): log the test metrics on Neptune. Defaults to True.
+        """
+        self._reset_trainer()
+
+        # setup
+        train_val_split = self.train_dataset.train_test_split(
+            test_size=config.val_size, stratify_by_column="label"
+        )
+        labels = np.array(train_val_split["train"]["label"])
+
+        train_dataloader = self.create_dataloader(
+            train_val_split["train"], batch_size=config.batch_size
+        )
+
+        val_dataloader = self.create_dataloader(
+            train_val_split["test"], batch_size=config.batch_size
+        )
+        test_dataloader = self.create_dataloader(
+            self.test_dataset, batch_size=config.batch_size
+        )
+
+        if logging:
+            run = neptune.init_run(project=NEPTUNE_PROJECT, api_token=NEPTUNE_API_TOKEN)
+
+        model = self._init_model()
+        model = model.to(self.device)
+        criterion = CrossEntropyLoss(weight=self.compute_class_weights(labels))
+        optimizer = self._init_optimizer(model)
+
+        # print info
+        if verbose:
+            self.print_info_passive_learning()
+
+        # early stopper
+        ES = EarlyStopping(
+            patience=config.es_patience,
+            verbose=True,
+            path=Path(join(self.pl_checkpoint_path, "best_model.pt")),
+        )
+
+        # training loop
+        for epoch in range(config.max_epoch):
+            running_loss = 0.0
+            for i, (inputs, labels) in tqdm(enumerate(train_dataloader, 0)):
+                # get the inputs; data is a list of [inputs, labels]
+                labels = labels.to(self.device)
+                for key, value in inputs.items():
+                    inputs[key] = value.to(self.device)
+
+                # zero the parameter gradients
+                optimizer.zero_grad()
+
+                # forward + backward + optimize
+                outputs = model(inputs)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+
+                # print statistics
+                running_loss += loss.item()
+
+            # evaluate model on validation set
+            val_metrics = self.eval_model(model, val_dataloader, criterion)
+            train_loss = running_loss / len(train_dataloader)
+            val_loss = val_metrics["loss"]
+            running_loss = 0.0
+            if logging:
+                run["loss/train"].append(train_loss)
+                run["loss/val"].append(val_loss)
+
+                for key, value in val_metrics.items():
+                    if key != "loss":
+                        run[f"val/{key}"].append(value)
+
+            if verbose:
+                self.print_val_metrics(epoch + 1, val_metrics)
+
+            # check early stopping
+            ES(val_loss, model)
+            if ES.early_stop:
+                break
+
+        # load best model
+        model.load_state_dict(
+            torch.load(Path(join(self.pl_checkpoint_path, "best_model.pt")))
+        )
+
+        # evaluate model on test dataset
+        test_metrics = self.eval_model(model, test_dataloader, criterion)
+        if verbose:
+            self.print_test_metrics(test_metrics)
+        if logging:
+            run["method"] = self.method_name
+            run["dataset"] = self.dataset
+            run["relation"] = self.relation_type
+            run["strategy"] = "passive learning"
+            run["config"] = config.__dict__
+            run["epochs"] = epoch
+
+            for key, value in test_metrics.items():
+                run["test/" + key] = value
+            run.stop()
+
+        return model
+
+    def set_al_metrics(self, baal_model: MyModelWrapperBilstm):
+        """
+        Configures the metrics that are to be computed during the active learning experiment
+
+        Args:
+            baal_model (MyModelWrapperBilstm): model wrapper
+
+        """
+        # accuracy
+        baal_model.add_metric(
+            name="acc",
+            initializer=lambda: Accuracy(task=self.task, average="micro").to(
+                self.device
+            ),
+        )
+
+        if self.dataset == "n2c2":
+            f1 = F1Score(num_classes=self.num_classes, ignore_index=0).to(self.device)
+            p = Precision(num_classes=self.num_classes, ignore_index=0).to(self.device)
+            r = Recall(num_classes=self.num_classes, ignore_index=0).to(self.device)
+            baal_model.add_metric(name="f1", initializer=lambda: f1)
+            baal_model.add_metric(name="p", initializer=lambda: p)
+            baal_model.add_metric(name="r", initializer=lambda: r)
+
+        else:  # self.dataset == "ddi":
+            # detection + classification metrics
+            cla_f1_micro = F1Score(
+                num_classes=self.num_classes, average="micro", ignore_index=0
+            ).to(self.device)
+
+            cla_p_micro = Precision(
+                num_classes=self.num_classes, average="micro", ignore_index=0
+            ).to(self.device)
+
+            cla_r_micro = Recall(
+                num_classes=self.num_classes, average="micro", ignore_index=0
+            ).to(self.device)
+
+            cla_f1_macro = F1Score(
+                num_classes=self.num_classes, average="macro", ignore_index=0
+            ).to(self.device)
+
+            cla_p_macro = Precision(
+                num_classes=self.num_classes, average="macro", ignore_index=0
+            ).to(self.device)
+
+            cla_r_macro = Recall(
+                num_classes=self.num_classes, average="macro", ignore_index=0
+            ).to(self.device)
+
+            baal_model.add_metric(name="micro_f1", initializer=lambda: cla_f1_micro)
+            baal_model.add_metric(name="micro_p", initializer=lambda: cla_p_micro)
+            baal_model.add_metric(name="micro_r", initializer=lambda: cla_r_micro)
+            baal_model.add_metric(name="macro_f1", initializer=lambda: cla_f1_macro)
+            baal_model.add_metric(name="macro_p", initializer=lambda: cla_p_macro)
+            baal_model.add_metric(name="macro_r", initializer=lambda: cla_r_macro)
+
+            # detection metrics
+            detect_f1 = DetectionF1Score().to(self.device)
+            detect_p = DetectionPrecision().to(self.device)
+            detect_r = DetectionRecall().to(self.device)
+
+            baal_model.add_metric(name="detect_f1", initializer=lambda: detect_f1)
+            baal_model.add_metric(name="detect_p", initializer=lambda: detect_p)
+            baal_model.add_metric(name="detect_r", initializer=lambda: detect_r)
+
+            # per class metrics
+            per_class_f1 = F1Score(num_classes=self.num_classes, average="none").to(
+                self.device
+            )
+
+            per_class_p = Precision(num_classes=self.num_classes, average="none").to(
+                self.device
+            )
+
+            per_class_r = Recall(num_classes=self.num_classes, average="none").to(
+                self.device
+            )
+
+            baal_model.add_metric(name="class_f1", initializer=lambda: per_class_f1)
+            baal_model.add_metric(name="class_p", initializer=lambda: per_class_p)
+            baal_model.add_metric(name="class_r", initializer=lambda: per_class_r)
+
+        return baal_model
+
+    def train_active_learning(
+        self,
+        query_strategy: BaalQueryStrategy,
+        config: BaalExperimentConfig,
+        verbose: bool = True,
+        logging: bool = True,
+    ):
+        """Trains the BiLSTM model using active learning
+
+        Args:
+            query_strategy (str): name of the query strategy to be used in the experiment.
+            config (BaalExperimentConfig): experiment configuration.
+            verbose (bool): determines if information is printed during trainig or not. Defaults to True.s
+            logging (bool): log the test metrics on Neptune. Defaults to True.
+        """
+        self._reset_trainer()
+
+        if logging:
+            run = neptune.init_run(project=NEPTUNE_PROJECT, api_token=NEPTUNE_API_TOKEN)
+
+        # setup querying
+        INIT_QUERY_SIZE = self.compute_init_q_size(config)
+        QUERY_SIZE = self.compute_q_size(config)
+        AL_STEPS = 2 # self.compute_al_steps(config)
+        
+        f_query_strategy = get_baal_query_strategy(
+            name=query_strategy.value,
+            shuffle_prop=config.shuffle_prop,
+            query_size=QUERY_SIZE,
+        )   
+
+
+        if verbose:
+            self.print_info_active_learning(
+                q_strategy=query_strategy.value,
+                pool_size=self.n_instances,
+                init_q_size=INIT_QUERY_SIZE,
+                q_size=QUERY_SIZE,
+            )
+
+        # setup active set
+        self.train_dataset.set_transform(self.transform)
+        self.test_dataset.set_transform(self.transform)
+        active_set = MyActiveLearningDatasetBilstm(self.train_dataset)
+        active_set.can_label = False
+        active_set.label_randomly(INIT_QUERY_SIZE)
+
+        # setup model
+        PATCH =  config.all_bayesian or (query_strategy == BaalQueryStrategy.BATCH_BALD)
+        if not PATCH: 
+            config.iterations = 1
+        model = self._init_model(PATCH)
+        model = model.to(self.device)
+        criterion = CrossEntropyLoss(self.compute_class_weights(active_set.labels))
+        optimizer = self._init_optimizer(model)
+
+        baal_model = MyModelWrapperBilstm(
+            model,
+            criterion,
+            replicate_in_memory=False,
+            min_train_passes=config.min_train_passes,
+        )
+        baal_model = self.set_al_metrics(baal_model)
+
+        # active loop
+        active_loop = MyActiveLearningLoop(
+            dataset=active_set,
+            get_probabilities=baal_model.predict_on_dataset,
+            heuristic=f_query_strategy,
+            query_size=QUERY_SIZE,
+            batch_size=config.batch_size,
+            iterations=config.iterations,
+            use_cuda=self.use_cuda,
+            verbose=False,
+            workers=2,
+            collate_fn=custom_collate,
+        )
+
+        # We will reset the weights at each active learning step so we make a copy.
+        init_weights = deepcopy(baal_model.state_dict())
+
+        if logging:
+            run["model"] = self.method_name
+            run["dataset"] = self.dataset
+            run["relation"] = self.relation_type
+            run["bayesian"] = config.all_bayesian or (
+                query_strategy == BaalQueryStrategy.BATCH_BALD
+            )
+            run["strategy"] = query_strategy.value
+            run["config"] = config.__dict__
+            run["annotation/intance_ann"].append(active_set.n_labelled / self.n_instances)
+            run["annotation/token_ann"].append(
+                active_set.n_labelled_tokens / self.n_tokens
+            )
+            run["annotation/char_ann"].append(
+                active_set.n_labelled_chars / self.n_characters
+            )
+
+        step_acc = []
+
+        # Active learning loop
+        for step in tqdm(range(AL_STEPS)):
+            init_step_time = time.time()
+
+            # Load the initial weights.
+            baal_model.load_state_dict(init_weights)
+
+            # Train the model on the currently labelled dataset.
+            init_train_time = time.time()
+            _ = baal_model.train_on_dataset(
+                dataset=active_set,
+                optimizer=optimizer,
+                batch_size=config.batch_size,
+                use_cuda=self.use_cuda,
+                epoch=config.max_epoch,
+                collate_fn=custom_collate,
+            )
+            train_time = time.time() - init_train_time
+
+            # test the model on the test set.
+            baal_model.test_on_dataset(
+                dataset=self.test_dataset,
+                batch_size=config.batch_size,
+                use_cuda=self.use_cuda,
+                average_predictions=config.iterations,
+                collate_fn=custom_collate,
+            )
+
+            if verbose:
+                self.print_al_iteration_metrics(step + 1, baal_model.get_metrics())
+
+            # query new instances to be labelled
+            init_query_time = time.time()
+            should_continue = active_loop.step()
+            query_time = time.time() - init_query_time
+            step_time = time.time() - init_step_time
+
+            if logging:
+                run["times/step_time"].append(step_time)
+                run["times/train_time"].append(train_time)
+                run["times/query_time"].append(query_time)
+                run["annotation/intance_ann"].append(
+                    active_set.n_labelled / self.n_instances
+                )
+                run["annotation/token_ann"].append(
+                    active_set.n_labelled_tokens / self.n_tokens
+                )
+                run["annotation/char_ann"].append(
+                    active_set.n_labelled_chars / self.n_characters
+                )
+
+            if not should_continue:
+                break
+
+            # adjust class weights
+            baal_model.criterion = CrossEntropyLoss(
+                self.compute_class_weights(active_set.labels)
+            )
+        # end of active learning loop
+
+        if logging:
+            for metrics in baal_model.active_learning_metrics.values():
+                for key, value in metrics.items():
+                    f_key = key.replace("test_", "test/").replace("train_", "train/")
+
+                    if "class" in key:
+                        for i, class_value in enumerate(value):
+                            run[f_key + "_" + str(i)].append(class_value)
+                    else:
+                        run[f_key].append(value)
+
+            run["train/step_acc"].extend(active_loop.step_acc)
+            run["train/step_score"].extend(active_loop.step_score)
+
+            run.stop()