Repository: dice-group/dice-embeddings (language: python)

Issue:
Refactoring Examples and Analysis of Experiments

Memory issues when running k fold cross validation with batching on a large dataset
For the dataset having the following stats:
```
Number of entities:51180
Number of relations:1
Number of triples on train set:26111
```
with 1 NVIDIA A100-SXM4 of 40GB and 1 CPU core with 340 GB memory on Nocta 2 cluster, the 10 fold batched cross validation fails by raising a memory kill error signal. 
Here is the dataset file [`commons_page_links_es.txt`](https://github.com/user-attachments/files/17965291/commons_page_links_es.txt). I am using the following command to run the training:
```
dicee --path_single_kg commons_page_links_es.txt --model Keci --num_epochs 200 --p 0 --q 1 --embedding_dim 256 --scoring_technique NegSample --batch_size 100_000 --optim Adopt --num_folds_for_cv 10
```

The batched evaluation works on small number of batches but it cannot handle large number of batches and gives a memory kill signal.


Discussion / hints:

Fix patch:
diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 00000000..619e419f
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,19 @@
+[run]
+omit =
+    tests/*
+    /tmp/*
+
+[report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    if self.debug:
+    if settings.DEBUG
+    raise AssertionError
+    raise NotImplementedError
+    if 0:
+    if __name__ == .__main__.:
+    if TYPE_CHECKING:
+    class .*\bProtocol\):
+    @(abc\.)?abstractmethod
+    pass
diff --git a/.github/workflows/github-actions-python-package.yml b/.github/workflows/github-actions-python-package.yml
index 98cbd177..b4cfee64 100644
--- a/.github/workflows/github-actions-python-package.yml
+++ b/.github/workflows/github-actions-python-package.yml
@@ -22,8 +22,13 @@ jobs:
 
       - name: Lint with ruff
         run: |
-          ruff  check dicee/ --select=E501 --line-length=200
+          ruff  check dicee/ --line-length=200
       - name: Test with pytest
         run: |
           wget https://files.dice-research.org/datasets/dice-embeddings/KGs.zip --no-check-certificate && unzip KGs.zip
           python -m pytest -p no:warnings -x
+      - name: Coverage report
+        run: |
+          pip install coverage
+          coverage run -m pytest -p no:warnings -x
+          coverage report -m 
diff --git a/README.md b/README.md
index 73192d29..7e685abf 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,11 @@
+[![Downloads](https://static.pepy.tech/badge/dicee)](https://pepy.tech/project/dicee)
+[![Downloads](https://img.shields.io/pypi/dm/dicee)](https://pypi.org/project/dicee/)
+[![Coverage](https://img.shields.io/badge/coverage-54%25-green)](https://dice-group.github.io/dice-embeddings/usage/main.html#coverage-report)
+[![Pypi](https://img.shields.io/badge/pypi-0.1.4-blue)](https://pypi.org/project/dicee/0.1.4/)
+[![Docs](https://img.shields.io/badge/documentation-0.1.4-yellow)](https://dice-group.github.io/dice-embeddings/index.html)
+
+![dicee_logo](docs/_static/images/dicee_logo.png)
+
 # DICE Embeddings: Hardware-agnostic Framework for Large-scale Knowledge Graph Embeddings
 
 Knowledge graph embedding research has mainly focused on learning continuous representations of knowledge graphs towards the link prediction problem. 
@@ -35,8 +43,7 @@ Deploy a pre-trained embedding model without writing a single line of code.
 ### Installation from Source
 ``` bash
 git clone https://github.com/dice-group/dice-embeddings.git
-conda create -n dice python=3.10.13 --no-default-packages && conda activate dice
-pip3 install -e .
+conda create -n dice python=3.10.13 --no-default-packages && conda activate dice && pip3 install -e .
 # or
 pip3 install -e .["dev"]
 ```
@@ -60,16 +67,80 @@ python -m pytest -p no:warnings --ff # to run the failures first and then the re
 ## Knowledge Graph Embedding Models
 <details> <summary> To see available Models</summary>
 
-1. TransE, DistMult, ComplEx, ConEx, QMult, OMult, ConvO, ConvQ, Keci
-2. All 44 models available in https://github.com/pykeen/pykeen#models
+* ```--model Decal | Keci | DualE | ComplEx | QMult | OMult | ConvQ | ConvO | ConEx | TransE | DistMult | Shallom```
+* ```--model Pykeen_QuatE | Pykeen_Mure ``` all embedding models available in https://github.com/pykeen/pykeen#models can be selected.
+
+Training and scoring techniques
+* ```--trainer torchCPUTrainer | PL | MP | torchDDP ```
+* ```--scoring_technique 1vsAll | KvsAll  | AllvsAll | KvsSample | NegSample ```
 
-> For more, please refer to `examples`.
 </details>
 
 ## How to Train
 <details> <summary> To see a code snippet </summary>
 
-To Train a KGE model (KECI) and evaluate it on the train, validation, and test sets of the UMLS benchmark dataset.
+#### Training Techniques
+
+A KGE model can be trained with a state-of-the-art training technique ```--trainer "torchCPUTrainer" | "PL" | "MP" | torchDDP ```
+```bash
+# CPU training
+dicee --dataset_dir "KGs/UMLS" --trainer "torchCPUTrainer" --scoring_technique KvsAll --model "Keci" --eval_model "train_val_test"
+# Distributed Data Parallelism
+dicee --dataset_dir "KGs/UMLS" --trainer "PL" --scoring_technique KvsAll --model "Keci" --eval_model "train_val_test"
+# Model Parallelism
+dicee --dataset_dir "KGs/UMLS" --trainer "MP" --scoring_technique KvsAll --model "Keci" --eval_model "train_val_test"
+# Distributed Data Parallelism in native torch
+OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc_per_node=gpu dicee --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test" --trainer "torchDDP" --scoring_technique KvsAll
+```
+A KGE model model can also be trained in multi-node multi-gpu DDP setting. 
+```bash
+torchrun --nnodes 2 --nproc_per_node=gpu  --node_rank 0 --rdzv_id 455 --rdzv_backend c10d --rdzv_endpoint=nebula  dicee --trainer "torchDDP" --dataset_dir "KGs/YAGO3-10"
+torchrun --nnodes 2 --nproc_per_node=gpu  --node_rank 1 --rdzv_id 455 --rdzv_backend c10d --rdzv_endpoint=nebula  dicee --trainer "torchDDP" --dataset_dir "KGs/YAGO3-10"
+```
+On large knowledge graphs, this configurations should be used.
+
+where the data is in the following form
+```bash
+$ head -3 KGs/UMLS/train.txt 
+acquired_abnormality    location_of     experimental_model_of_disease
+anatomical_abnormality  manifestation_of        physiologic_function
+alga    isa     entity
+
+$ head -3 KGs/YAGO3-10/valid.txt 
+Mikheil_Khutsishvili    playsFor        FC_Merani_Tbilisi
+Ebbw_Vale       isLocatedIn     Blaenau_Gwent
+Valenciennes    isLocatedIn     Nord-Pas-de-Calais
+```
+By default, ```--backend "pandas" --separator "\s+" ``` is used in ```pandas.read_csv(sep=args.separator)``` to obtain triples.
+You can choose a suitable backend for your knowledge graph ```--backend pandas | polars  | rdflib ```.
+On large knowledge graphs n-triples, ```--backend "polars" --separator " " ``` is a good option.
+**Apart from n-triples or standard link prediction dataset formats, we support ["owl", "nt", "turtle", "rdf/xml", "n3"]***.
+On other RDF knowledge graphs,  ```--backend "rdflib" ``` can be used. Note that knowledge graphs must not contain blank nodes or literals.
+Moreover, a KGE model can be also trained  by providing **an endpoint of a triple store**. 
+```bash
+dicee --sparql_endpoint "http://localhost:3030/mutagenesis/" --model Keci
+```
+
+#### Scoring Techniques
+
+We have implemented state-of-the-art scoring techniques to train a KGE model ```--scoring_technique 1vsAll | KvsAll  | AllvsAll | KvsSample | NegSample ```.
+```bash
+dicee --dataset_dir "KGs/YAGO3-10" --model Keci --trainer "torchCPUTrainer" --scoring_technique "NegSample" --neg_ratio 10 --num_epochs 10 --batch_size 10_000 --num_core 0 --eval_model None
+# Epoch:10: 100%|███████████| 10/10 [01:31<00:00,  9.11s/it, loss_step=0.09423, loss_epoch=0.07897]
+# Training Runtime: 1.520 minutes.
+dicee --dataset_dir "KGs/YAGO3-10" --model Keci --trainer "torchCPUTrainer" --scoring_technique "NegSample" --neg_ratio 10 --num_epochs 10 --batch_size 10_000 --num_core 10 --eval_model None
+# Epoch:10: 100%|███████████| 10/10 [00:58<00:00,  5.80s/it, loss_step=0.11909, loss_epoch=0.07991]
+# Training Runtime: 58.106 seconds.
+dicee --dataset_dir "KGs/YAGO3-10" --model Keci --trainer "torchCPUTrainer" --scoring_technique "NegSample" --neg_ratio 10 --num_epochs 10 --batch_size 10_000 --num_core 20 --eval_model None
+# Epoch:10: 100%|███████████| 10/10 [01:01<00:00,  6.16s/it, loss_step=0.10751, loss_epoch=0.06962]
+# Training Runtime: 1.029 minutes.
+dicee --dataset_dir "KGs/YAGO3-10" --model Keci --trainer "torchCPUTrainer" --scoring_technique "NegSample" --neg_ratio 10 --num_epochs 10 --batch_size 10_000 --num_core 50 --eval_model None
+# Epoch:10: 100%|███████████| 10/10 [01:08<00:00,  6.83s/it, loss_step=0.05347, loss_epoch=0.07003]
+# Training Runtime: 1.140 minutes.
+```
+Increasing the number of cores often (but not always) helps to decrease the runtimes on large knowledge graphs ```--num_core 4 --scoring_technique KvsSample | NegSample --neg_ratio 1``` 
+
+A KGE model can be also trained in a python script
 ```python
 from dicee.executer import Execute
 from dicee.config import Namespace
@@ -86,43 +157,9 @@ print(reports["Train"]["MRR"]) # => 0.9912
 print(reports["Test"]["MRR"]) # => 0.8155
 # See the Keci_UMLS folder embeddings and all other files
 ```
-where the data is in the following form
-```bash
-$ head -3 KGs/UMLS/train.txt 
-acquired_abnormality    location_of     experimental_model_of_disease
-anatomical_abnormality  manifestation_of        physiologic_function
-alga    isa     entity
-```
-A KGE model can also be trained from the command line
-```bash
-dicee --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-```
-dicee automatically detects available GPUs and trains a model with distributed data parallels technique.
-```bash
-# Train a model by only using the GPU-0
-CUDA_VISIBLE_DEVICES=0 dicee --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-# Train a model by only using GPU-1
-CUDA_VISIBLE_DEVICES=1 dicee --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-# Train a model by using all available GPUs
-dicee --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-```
-Under the hood, dicee executes the run.py script and uses [lightning](https://lightning.ai/) as a default trainer.
-```bash
-# Two equivalent executions
-# (1)
-dicee --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-# (2)
-CUDA_VISIBLE_DEVICES=0,1 python dicee/scripts/run.py --trainer PL --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-```
-Similarly, models can be easily trained with torchrun
-```bash
-torchrun --standalone --nnodes=1 --nproc_per_node=gpu dicee/scripts/run.py --trainer torchDDP --dataset_dir "KGs/UMLS" --model Keci --eval_model "train_val_test"
-```
-You can also train a model in multi-node multi-gpu setting.
-```bash
-torchrun --nnodes 2 --nproc_per_node=gpu  --node_rank 0 --rdzv_id 455 --rdzv_backend c10d --rdzv_endpoint=nebula  dicee/scripts/run.py --trainer torchDDP --dataset_dir KGs/UMLS
-torchrun --nnodes 2 --nproc_per_node=gpu  --node_rank 1 --rdzv_id 455 --rdzv_backend c10d --rdzv_endpoint=nebula dicee/scripts/run.py --trainer torchDDP --dataset_dir KGs/UMLS
-```
+
+#### Continual Learning
+
 Train a KGE model by providing the path of a single file and store all parameters under newly created directory
 called `KeciFamilyRun`.
 ```bash
@@ -135,17 +172,13 @@ _:1 <http://www.w3.org/1999/02/22-rdf-syntax-ns#type> <http://www.w3.org/2002/07
 <http://www.benchmark.org/family#hasChild> <http://www.w3.org/1999/02/22-rdf-syntax-ns#type> <http://www.w3.org/2002/07/owl#ObjectProperty> .
 <http://www.benchmark.org/family#hasParent> <http://www.w3.org/1999/02/22-rdf-syntax-ns#type> <http://www.w3.org/2002/07/owl#ObjectProperty> .
 ```
-**Continual Training:** the training phase of a pretrained model can be resumed.
-```bash
-dicee --continual_learning KeciFamilyRun --path_single_kg "KGs/Family/family-benchmark_rich_background.owl" --model Keci --path_to_store_single_run KeciFamilyRun --backend rdflib --eval_model None
-```
 
-**Apart from n-triples or standard link prediction dataset formats, we support ["owl", "nt", "turtle", "rdf/xml", "n3"]***.
-Moreover, a KGE model can be also trained  by providing **an endpoint of a triple store**.
+**Continual Training:** the training phase of a pretrained model can be resumed. The model will saved in the same directory ``` --continual_learning "KeciFamilyRun"```.
 ```bash
-dicee --sparql_endpoint "http://localhost:3030/mutagenesis/" --model Keci
+dicee --continual_learning "KeciFamilyRun" --path_single_kg "KGs/Family/family-benchmark_rich_background.owl" --model Keci --backend rdflib --eval_model None
 ```
-For more, please refer to `examples`.
+
+
 </details>
 
 ## Creating an Embedding Vector Database 
@@ -298,12 +331,239 @@ mure.predict_topk(h=["Mongolia"],r=["isLocatedIn"],topk=3)
 from dicee import KGE
 KGE(path='...').deploy(share=True,top_k=10)
 ```
+
 </details>
 
 <details> <summary> To see the interface of the webservice</summary>
 <img src="dicee/lp.png" alt="Italian Trulli">
 </details>
 
+
+## Link Prediction Benchmarks
+
+In the below, we provide a brief overview of the link prediction results. Results are sorted in descending order of the size of the respective dataset.
+
+#### YAGO3-10 ####
+
+<details> <summary> To see the results </summary>
+
+|                    |       |   MRR | Hits@1 | Hits@3 | Hits@10 |
+|--------------------|-------|------:|-------:|-------:|--------:|
+| ComplEx-KvsAll     | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-KvsAll     | val   | 0.374 |  0.308 |  0.402 |   0.501 |
+| ComplEx-KvsAll     | test  | 0.372 |  0.302 |  0.404 |   0.505 |
+| ComplEx-KvsAll-SWA | train | 0.998 |  0.997 |  1.000 |   1.000 |
+| ComplEx-KvsAll-SWA | val   | 0.345 |  0.279 |  0.372 |   0.474 |
+| ComplEx-KvsAll-SWA | test  | 0.341 |  0.272 |  0.374 |   0.474 |
+| ComplEx-KvsAll-ASWA | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-KvsAll-ASWA | val   | 0.404 |  0.335 |  0.448 |   0.531 |
+| ComplEx-KvsAll-ASWA | test  | 0.398 |  0.325 |  0.449 |   0.530 |
+| Keci-KvsAll        | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-KvsAll        | val   | 0.337 |  0.268 |  0.370 |   0.468 |
+| Keci-KvsAll        | test  | 0.343 |  0.274 |  0.376 |   0.343 |
+| Keci-KvsAll-SWA    | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-KvsAll-SWA    | val   | 0.325 |  0.253 |  0.358 |   0.459 |
+| Keci-KvsAll-SWA    | test  | 0.334 |  0.263 |  0.367 |   0.470 |
+| Keci-KvsAll-ASWA   | train | 0.978 |  0.969 |  0.985 |   0.991 |
+| Keci-KvsAll-ASWA   | val   | 0.400 |  0.324 |  0.439 |   0.540 |
+| Keci-KvsAll-ASWA   | test  | 0.394 |  0.317 |  0.439 |   0.539 |
+
+```--embedding_dim 256 --num_epochs 300 --batch_size 1024 --optim Adam 0.1``` leading to 31.6M params.
+Observations: A severe overfitting. ASWA improves the generalization better than SWA.
+
+</details>
+
+#### FB15k-237 ####
+
+<details> <summary> To see the results </summary>
+
+|                    |       |   MRR | Hits@1 | Hits@3 | Hits@10 |
+|--------------------|-------|------:|-------:|-------:|--------:|
+| ComplEx-KvsAll     | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-KvsAll     | val   | 0.197 |  0.140 |  0.211 |   0.307 |
+| ComplEx-KvsAll     | test  | 0.192 |  0.137 |  0.204 |   0.300 |
+| ComplEx-KvsAll-SWA | train | 0.911 |  0.875 |  0.938 |   0.970 |
+| ComplEx-KvsAll-SWA | val   | 0.169 |  0.121 |  0.178 |   0.264 |
+| ComplEx-KvsAll-SWA | test  | 0.166 |  0.118 |  0.176 |   0.261 |
+| ComplEx-KvsAll-ASWA | train | 0.780 | 0.719 | 0.822  |   0.886 |
+| ComplEx-KvsAll-ASWA | val   | 0.220 | 0.158 | 0.240  |   0.342 |
+| ComplEx-KvsAll-ASWA | test  | 0.217 | 0.155 |  0.234 |   0.337 |
+| Keci-KvsAll        | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-KvsAll        | val   | 0.158 |  0.107 |  0.166 |   0.259 |
+| Keci-KvsAll        | test  | 0.155 |  0.105 |  0.162 |   0.253 |
+| Keci-KvsAll-SWA    | train | 0.941 |  0.909 |  0.967 |   0.990 |
+| Keci-KvsAll-SWA    | val   | 0.188 |  0.133 |  0.200 |   0.298 |
+| Keci-KvsAll-SWA    | test  | 0.183 |  0.128 |  0.195 |   0.292 |
+| Keci-KvsAll-ASWA   | train | 0.745 |  0.666 |  0.799 |   0.886 |
+| Keci-KvsAll-ASWA   | val   | 0.221 |  0.158 |  0.237 |   0.346 |
+| Keci-KvsAll-ASWA   | test  | 0.216 |  0.153 |  0.233 |   0.342 |
+
+</details>
+
+
+#### WN18RR ####
+
+<details> <summary> To see the results </summary>
+
+|                    |       |   MRR | Hits@1 | Hits@3 | Hits@10 |
+|--------------------|-------|------:|-------:|-------:|--------:|
+| ComplEx-KvsAll     | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-KvsAll     | val   | 0.346 |  0.337 |  0.353 |   0.359 |
+| ComplEx-KvsAll     | test  | 0.343 |  0.333 |  0.349 |   0.357 |
+| ComplEx-KvsAll-SWA | train | 0.335 |  0.242 |  0.365 |   0.522 |
+| ComplEx-KvsAll-SWA | val   | 0.014 |  0.004 |  0.013 |   0.029 |
+| ComplEx-KvsAll-SWA | test  | 0.018 |  0.007 |  0.017 |   0.036 |
+| ComplEx-KvsAll-ASWA | train | 0.999 | 0.999 | 0.999  |   0.999 |
+| ComplEx-KvsAll-ASWA | val   | 0.352 | 0.348 | 0.355  |   0.359 |
+| ComplEx-KvsAll-ASWA | test  | 0.348 | 0.342 |  0.351 |   0.355 |
+| Keci-KvsAll        | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-KvsAll        | val   | 0.316 |  0.296 |  0.331 |   0.352 |
+| Keci-KvsAll        | test  | 0.308 |  0.285 |  0.326 |   0.346 |
+| Keci-KvsAll-SWA    | train | 0.608 |  0.516 |  0.657 |   0.787 |
+| Keci-KvsAll-SWA    | val   | 0.069 |  0.038 |  0.073 |   0.130 |
+| Keci-KvsAll-SWA    | test  | 0.061 |  0.032 |  0.064 |   0.119 |
+| Keci-KvsAll-ASWA   | train | 0.951 |  0.916 |  0.985 |   0.997 |
+| Keci-KvsAll-ASWA   | val   | 0.356 |  0.353 |  0.356 |   0.360 |
+| Keci-KvsAll-ASWA   | test  | 0.350 |  0.347 |  0.350 |   0.355 |
+
+</details>
+
+#### UMLS ####
+
+<details> <summary> To see the results </summary>
+
+|                       |       |   MRR | Hits@1 | Hits@3 | Hits@10 |
+|-----------------------|-------|------:|-------:|-------:|--------:|
+| ComplEx-KvsAll        | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-KvsAll        | val   | 0.684 |  0.557 |  0.771 |   0.928 |
+| ComplEx-KvsAll        | test  | 0.680 |  0.563 |  0.750 |   0.918 |
+| ComplEx-AllvsAll      | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-AllvsAll      | val   | 0.771 |  0.670 |  0.847 |   0.949 |
+| ComplEx-AllvsAll      | test  | 0.778 |  0.678 |  0.850 |   0.957 |
+| ComplEx-KvsAll-SWA    | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-KvsAll-SWA    | val   | 0.762 |  0.666 |  0.825 |   0.941 |
+| ComplEx-KvsAll-SWA    | test  | 0.757 |  0.653 |  0.833 |   0.939 |
+| ComplEx-AllvsAll-SWA  | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| ComplEx-AllvsAll-SWA  | val   | 0.817 |  0.736 |  0.879 |   0.953 |
+| ComplEx-AllvsAll-SWA  | test  | 0.827 |  0.748 |  0.883 |   0.967 |
+| ComplEx-KvsAll-ASWA   | train | 0.998 |  0.997 |  0.999 |   1.000 |
+| ComplEx-KvsAll-ASWA   | val   | 0.799 |  0.712 |  0.863 |   0.946 |
+| ComplEx-KvsAll-ASWA   | test  | 0.804 |  0.720 |  0.866 |   0.948 |
+| ComplEx-AllvsAll-ASWA | train | 0.998 |  0.997 |  0.998 |   0.999 |
+| ComplEx-AllvsAll-ASWA | val   | 0.879 |  0.824 |  0.926 |   0.964 |
+| ComplEx-AllvsAll-ASWA | test  | 0.877 |  0.819 |  0.924 |   0.971 |
+| Keci-KvsAll           | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-KvsAll           | val   | 0.538 |  0.401 |  0.595 |   0.829 |
+| Keci-KvsAll           | test  | 0.543 |  0.411 |  0.610 |   0.815 |
+| Keci-AllvsAll         | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-AllvsAll         | val   | 0.672 |  0.556 |  0.742 |   0.909 |
+| Keci-AllvsAll         | test  | 0.684 |  0.567 |  0.759 |   0.914 |
+| Keci-KvsAll-SWA       | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-KvsAll-SWA       | val   | 0.633 |  0.509 |  0.705 |   0.877 |
+| Keci-KvsAll-SWA       | test  | 0.628 |  0.498 |  0.710 |   0.868 |
+| Keci-AllvsAll-SWA     | train | 1.000 |  1.000 |  1.000 |   1.000 |
+| Keci-AllvsAll-SWA     | val   | 0.697 |  0.584 |  0.770 |   0.911 |
+| Keci-AllvsAll-SWA     | test  | 0.711 |  0.606 |  0.775 |   0.921 |
+| Keci-KvsAll-ASWA      | train | 0.996 |  0.993 |  0.999 |   1.000 |
+| Keci-KvsAll-ASWA      | val   | 0.767 |  0.668 |  0.836 |   0.944 |
+| Keci-KvsAll-ASWA      | test  | 0.762 |  0.660 |  0.830 |   0.949 |
+| Keci-AllvsAll-ASWA    | train | 0.998 |  0.997 |  0.999 |   1.000 |
+| Keci-AllvsAll-ASWA    | val   | 0.852 |  0.793 |  0.896 |   0.955 |
+| Keci-AllvsAll-ASWA    | test  | 0.848 |  0.787 |  0.886 |   0.951 |
+
+
+
+```--embedding_dim 256 --num_epochs 300 --batch_size 1024 --optim Adam 0.1``` leading to 58.1K params.
+Observations: 
++ A severe overfitting. 
++ AllvsAll improves the generalization more than KvsAll does
++ ASWA improves the generalization more than SWA does
+
+```bash
+dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "KvsSample" --embedding_dim 256 --num_epochs 100 --batch_size 32 --num_core 10
+# Epoch 99: 100%|███████████| 13/13 [00:00<00:00, 29.56it/s, loss_step=6.46e-6, loss_epoch=8.35e-6]
+# *** Save Trained Model ***
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 1.0, 'H@3': 1.0, 'H@10': 1.0, 'MRR': 1.0}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.33358895705521474, 'H@3': 0.5253067484662577, 'H@10': 0.7576687116564417, 'MRR': 0.46992150194876076}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+# {'H@1': 0.3320726172465961, 'H@3': 0.5098335854765507, 'H@10': 0.7594553706505295, 'MRR': 0.4633434701052234}
+```
+Increasing cores increases the runtimes if there is a preprocessing step at the batch generation.
+```bash
+dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "KvsAll" --embedding_dim 256 --num_epochs 100 --batch_size 32
+# Epoch 99: 100%|██████████| 13/13 [00:00<00:00, 101.94it/s, loss_step=8.11e-6, loss_epoch=8.92e-6]
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 1.0, 'H@3': 1.0, 'H@10': 1.0, 'MRR': 1.0}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.348159509202454, 'H@3': 0.5659509202453987, 'H@10': 0.7883435582822086, 'MRR': 0.4912162082105331}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+# {'H@1': 0.34568835098335854, 'H@3': 0.5544629349470499, 'H@10': 0.7776096822995462, 'MRR': 0.48692617590763265}
+```
+
+```bash
+dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "AllvsAll" --embedding_dim 256 --num_epochs 100 --batch_size 32
+# Epoch 99: 100%|██████████████| 98/98 [00:01<00:00, 88.95it/s, loss_step=0.000, loss_epoch=0.0655]
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 0.9976993865030674, 'H@3': 0.9997124233128835, 'H@10': 0.9999041411042945, 'MRR': 0.9987183437408705}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.3197852760736196, 'H@3': 0.5398773006134969, 'H@10': 0.7714723926380368, 'MRR': 0.46912531544840963}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+# {'H@1': 0.329803328290469, 'H@3': 0.5711043872919819, 'H@10': 0.7934947049924357, 'MRR': 0.4858500337837166}
+```
+In KvsAll and AllvsAll, a single data point **z=(x,y)** corresponds to a tuple of input indices **x** and multi-label output vector **y**.
+**x** is a tuple of indices of a unique entity and relation pair.
+**y** contains a binary vector of size of the number of unique entities.
+
+To mitigate the rate of overfitting, many regularization techniques can be applied ,e.g.,
+Stochastic Weight Averaging (SWA), Adaptive Stochastic Weight Averaging (ASWA), or Dropout.
+Use ```--swa``` to apply Stochastic Weight Averaging
+```bash
+dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "KvsAll" --embedding_dim 256 --num_epochs 100 --batch_size 32 --swa
+# Epoch 99: 100%|███████████| 13/13 [00:00<00:00, 85.61it/s, loss_step=8.11e-6, loss_epoch=8.92e-6]
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 1.0, 'H@3': 1.0, 'H@10': 1.0, 'MRR': 1.0}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.45858895705521474, 'H@3': 0.6510736196319018, 'H@10': 0.8458588957055214, 'MRR': 0.5845156794070833}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+# {'H@1': 0.4636913767019667, 'H@3': 0.651285930408472, 'H@10': 0.8456883509833586, 'MRR': 0.5877221440365971}
+# Total Runtime: 25.417 seconds
+```
+Use ```--adaptive_swa``` to apply Adaptive Stochastic Weight Averaging. Currently, ASWA should not be used with DDP on multi GPUs.
+We are working on it.
+```bash
+CUDA_VISIBLE_DEVICES=0 dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "KvsAll" --embedding_dim 256 --num_epochs 100 --batch_size 32 --adaptive_swa
+# Epoch 99: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 49/49 [00:00<00:00, 93.86it/s, loss_step=0.0978, loss_epoch=0.143]
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 0.9974118098159509, 'H@3': 0.9992331288343558, 'H@10': 0.9996165644171779, 'MRR': 0.9983922084274367}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.7668711656441718, 'H@3': 0.8696319018404908, 'H@10': 0.9440184049079755, 'MRR': 0.828767705987023}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+#{'H@1': 0.7844175491679274, 'H@3': 0.8888048411497731, 'H@10': 0.9546142208774584, 'MRR': 0.8460991515345323}
+```
+```bash
+CUDA_VISIBLE_DEVICES=0 dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "KvsAll" --embedding_dim 256 --input_dropout_rate 0.1 --num_epochs 100 --batch_size 32 --adaptive_swa
+# Epoch 99: 100%|██████████████████████████████████████████████████████████| 49/49 [00:00<00:00, 93.49it/s, loss_step=0.600, loss_epoch=0.553]
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 0.9970283742331288, 'H@3': 0.9992331288343558, 'H@10': 0.999808282208589, 'MRR': 0.9981489117237927}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.8473926380368099, 'H@3': 0.9049079754601227, 'H@10': 0.9470858895705522, 'MRR': 0.8839172788777631}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+# {'H@1': 0.8381240544629349, 'H@3': 0.9167927382753404, 'H@10': 0.9568835098335855, 'MRR': 0.8829572716873321}
+
+CUDA_VISIBLE_DEVICES=0 dicee --dataset_dir "KGs/UMLS" --model "Keci" --p 0 --q 1 --trainer "PL" --scoring_technique "KvsAll" --embedding_dim 256 --input_dropout_rate 0.2 --num_epochs 100 --batch_size 32 --adaptive_swa
+# Epoch 99: 100%|██████████████████████████████████████████████████████████| 49/49 [00:00<00:00, 94.43it/s, loss_step=0.108, loss_epoch=0.111]
+# Evaluate Keci on Train set: Evaluate Keci on Train set
+# {'H@1': 0.9818826687116564, 'H@3': 0.9942484662576687, 'H@10': 0.9972200920245399, 'MRR': 0.9885307022708297}
+# Evaluate Keci on Validation set: Evaluate Keci on Validation set
+# {'H@1': 0.8581288343558282, 'H@3': 0.9156441717791411, 'H@10': 0.9447852760736196, 'MRR': 0.8930935122236525}
+# Evaluate Keci on Test set: Evaluate Keci on Test set
+# {'H@1': 0.8494704992435703, 'H@3': 0.9334341906202723, 'H@10': 0.9667170953101362, 'MRR': 0.8959156201718665}
+```
+
+</details>
+
 ## Docker
 <details> <summary> Details</summary>
 To build the Docker image:
diff --git a/dicee/__init__.py b/dicee/__init__.py
index 1c0570cb..7c5741aa 100644
--- a/dicee/__init__.py
+++ b/dicee/__init__.py
@@ -4,4 +4,4 @@
 from .executer import Execute # noqa
 from .dataset_classes import * # noqa
 from .query_generator import QueryGenerator # noqa
-__version__ = '0.1.4'
+__version__ = '0.1.5'
diff --git a/dicee/__main__.py b/dicee/__main__.py
new file mode 100644
index 00000000..3ba749b9
--- /dev/null
+++ b/dicee/__main__.py
@@ -0,0 +1,6 @@
+# dicee/__main__.py
+
+from dicee.scripts.run import main # Import the main entry point of dicee
+
+if __name__ == "__main__":
+    main()  # Call the main function to execute the program logic
\ No newline at end of file
diff --git a/dicee/abstracts.py b/dicee/abstracts.py
index 26bf16b1..70838ac0 100644
--- a/dicee/abstracts.py
+++ b/dicee/abstracts.py
@@ -26,6 +26,8 @@ def __init__(self, args, callbacks):
         self.attributes = args
         self.callbacks = callbacks
         self.is_global_zero = True
+        self.global_rank=0
+        self.local_rank = 0
         # Set True to use Model summary callback of pl.
         torch.manual_seed(self.attributes.random_seed)
         torch.cuda.manual_seed_all(self.attributes.random_seed)
@@ -178,21 +180,13 @@ def __init__(self, path: str = None, url: str = None, construct_ensemble: bool =
             self.num_relations = len(self.relation_to_idx)
             self.entity_to_idx: dict
             self.relation_to_idx: dict
-            assert list(self.entity_to_idx.values()) == list(range(0, len(self.entity_to_idx)))
-            assert list(self.relation_to_idx.values()) == list(range(0, len(self.relation_to_idx)))
+            # 0, ....,
+            assert sorted(list(self.entity_to_idx.values())) == list(range(0, len(self.entity_to_idx)))
+            assert sorted(list(self.relation_to_idx.values())) == list(range(0, len(self.relation_to_idx)))
 
             self.idx_to_entity = {v: k for k, v in self.entity_to_idx.items()}
             self.idx_to_relations = {v: k for k, v in self.relation_to_idx.items()}
 
-        # See https://numpy.org/doc/stable/reference/generated/numpy.memmap.html
-        # @TODO: Ignore temporalryIf file exists
-        # if os.path.exists(self.path + '/train_set.npy'):
-        #    self.train_set = np.load(file=self.path + '/train_set.npy', mmap_mode='r')
-
-        # if apply_semantic_constraint:
-        #    (self.domain_constraints_per_rel, self.range_constraints_per_rel,
-        #     self.domain_per_rel, self.range_per_rel) = create_constraints(self.train_set)
-
     def get_eval_report(self) -> dict:
         return load_json(self.path + "/eval_report.json")
 
@@ -253,17 +247,6 @@ def get_padded_bpe_triple_representation(self, triples: List[List[str]]) -> Tupl
             padded_bpe_t.append(self.get_bpe_token_representation(str_o))
         return padded_bpe_h, padded_bpe_r, padded_bpe_t
 
-    def get_domain_of_relation(self, rel: str) -> List[str]:
-        x = [self.idx_to_entity[i] for i in self.domain_per_rel[self.relation_to_idx[rel]]]
-        res = set(x)
-        assert len(x) == len(res)
-        return res
-
-    def get_range_of_relation(self, rel: str) -> List[str]:
-        x = [self.idx_to_entity[i] for i in self.range_per_rel[self.relation_to_idx[rel]]]
-        res = set(x)
-        assert len(x) == len(res)
-        return res
 
     def set_model_train_mode(self) -> None:
         """
diff --git a/dicee/analyse_experiments.py b/dicee/analyse_experiments.py
index ebe61089..09130c3e 100644
--- a/dicee/analyse_experiments.py
+++ b/dicee/analyse_experiments.py
@@ -1,4 +1,7 @@
-""" This script should be moved to dicee/scripts"""
+""" This script should be moved to dicee/scripts
+Example:
+python dicee/analyse_experiments.py --dir Experiments --features "model" "trainMRR" "testMRR"
+"""
 import os
 import json
 import pandas as pd
@@ -120,19 +123,13 @@ def analyse(args):
         if os.path.isdir(full_path) is False:
             continue
 
-        
         with open(f'{full_path}/configuration.json', 'r') as f:
             config = json.load(f)
-            
-        try:
-            with open(f'{full_path}/report.json', 'r') as f:
-                report = json.load(f)
-                report = {i: report[i] for i in ['Runtime', 'NumParam']}
-            with open(f'{full_path}/eval_report.json', 'r') as f:
-                eval_report = json.load(f)
-        except FileNotFoundError:
-            print("NOT found")
-            continue
+        with open(f'{full_path}/report.json', 'r') as f:
+            report = json.load(f)
+            report = {i: report[i] for i in ['Runtime', 'NumParam']}
+        with open(f'{full_path}/eval_report.json', 'r') as f:
+            eval_report = json.load(f)
         config.update(eval_report)
         config.update(report)
         if "Train" in config:
@@ -160,10 +157,9 @@ def analyse(args):
     # print(df.columns)
     try:
         df_features = df[args.features]
-    except:
+    except KeyError:
         print(f"--features ({args.features}) is not a subset of {df.columns}")
-        exit(1)
-
+        raise KeyError
     print(df_features.to_latex(index=False, float_format="%.3f"))
     path_to_save = args.dir + '/summary.csv'
     df_features.to_csv(path_or_buf=path_to_save)
diff --git a/dicee/callbacks.py b/dicee/callbacks.py
index 34ccf635..edc75b99 100644
--- a/dicee/callbacks.py
+++ b/dicee/callbacks.py
@@ -166,9 +166,10 @@ def on_fit_end(self, trainer, model):
         if self.initial_eval_setting:
             # ADD this info back
             trainer.evaluator.args.eval_model = self.initial_eval_setting
-
-        param_ensemble = torch.load(f"{self.path}/aswa.pt", torch.device("cpu"))
-        model.load_state_dict(param_ensemble)
+        
+        if trainer.global_rank==trainer.local_rank==0:
+            param_ensemble = torch.load(f"{self.path}/aswa.pt", torch.device("cpu"))
+            model.load_state_dict(param_ensemble)
 
     @staticmethod
     def compute_mrr(trainer, model) -> float:
@@ -241,6 +242,10 @@ def decide(self, running_model_state_dict, ensemble_state_dict, val_running_mode
             return True
 
     def on_train_epoch_end(self, trainer, model):
+        
+        if (trainer.global_rank == trainer.local_rank == 0) is False:
+            return None
+
         # (1) Increment epoch counter
         self.epoch_count += 1
         # (2) Save the given eval setting if it is not saved.
diff --git a/dicee/config.py b/dicee/config.py
index 3d7921cc..5c5e69a8 100644
--- a/dicee/config.py
+++ b/dicee/config.py
@@ -50,6 +50,9 @@ def __init__(self, **kwargs):
         self.backend: str = "pandas"
         """Backend to read, process, and index input knowledge graph. pandas, polars and rdflib available"""
 
+        self.separator: str = "\s+"
+        """separator for extracting head, relation and tail from a triple"""
+
         self.trainer: str = 'torchCPUTrainer'
         """Trainer for knowledge graph embedding model"""
 
@@ -82,7 +85,6 @@ def __init__(self, **kwargs):
 
         self.label_smoothing_rate: float = 0.0
 
-
         self.num_core: int = 0
         """Number of CPUs to be used in the mini-batch loading process"""
 
@@ -136,6 +138,9 @@ def __init__(self, **kwargs):
         self.continual_learning=None
         "Path of a pretrained model size of LLM"
 
+        self.auto_batch_finding=False
+        "A flag for using auto batch finding"
+
     def __iter__(self):
         # Iterate
         for k, v in self.__dict__.items():
diff --git a/dicee/dataset_classes.py b/dicee/dataset_classes.py
index 9b6b38bd..1c48cf2e 100644
--- a/dicee/dataset_classes.py
+++ b/dicee/dataset_classes.py
@@ -4,7 +4,7 @@
 import pytorch_lightning as pl
 from typing import List, Tuple, Union
 from .static_preprocess_funcs import mapping_from_first_two_cols_to_third
-from .static_funcs import timeit, load_pickle
+from .static_funcs import timeit, load_term_mapping
 
 
 @timeit
@@ -13,8 +13,8 @@ def reload_dataset(path: str, form_of_labelling, scoring_technique, neg_ratio, l
     return construct_dataset(train_set=np.load(path + '/train_set.npy'),
                              valid_set=None,
                              test_set=None,
-                             entity_to_idx=load_pickle(file_path=path + '/entity_to_idx.p'),
-                             relation_to_idx=load_pickle(file_path=path + '/relation_to_idx.p'),
+                             entity_to_idx=load_term_mapping(file_path=path + '/entity_to_idx'),
+                             relation_to_idx=load_term_mapping(file_path=path + '/relation_to_idx'),
                              form_of_labelling=form_of_labelling,
                              scoring_technique=scoring_technique, neg_ratio=neg_ratio,
                              label_smoothing_rate=label_smoothing_rate)
@@ -63,9 +63,9 @@ def construct_dataset(*,
         if scoring_technique == '1vsAll':
             # Multi-class.
             train_set = OnevsAllDataset(train_set, entity_idxs=entity_to_idx)
-        elif scoring_technique == 'KvsSample':
-            # Multi-label.
-            train_set = KvsSampleDataset(train_set=train_set,
+        elif scoring_technique == '1vsSample':
+            # Dynamic Multi-class
+            train_set = OnevsSample(train_set=train_set,
                                          num_entities=len(entity_to_idx),
                                          num_relations=len(relation_to_idx),
                                          neg_sample_ratio=neg_ratio,
@@ -83,6 +83,14 @@ def construct_dataset(*,
                                  entity_idxs=entity_to_idx,
                                  relation_idxs=relation_to_idx,
                                  label_smoothing_rate=label_smoothing_rate)
+        elif scoring_technique == 'KvsSample':
+            # Dynamic Multi-label.
+            train_set = KvsSampleDataset(train_set,
+                               entity_idxs=entity_to_idx,
+                               relation_idxs=relation_to_idx,
+                               form=form_of_labelling,
+                                         neg_ratio=neg_ratio,
+                               label_smoothing_rate=label_smoothing_rate)
         else:
             raise ValueError(f'Invalid scoring technique : {scoring_technique}')
     elif form_of_labelling == 'RelationPrediction':
@@ -91,6 +99,7 @@ def construct_dataset(*,
                            form=form_of_labelling, label_smoothing_rate=label_smoothing_rate)
     else:
         raise KeyError('Illegal input.')
+    print(f"Number of datapoints: {len(train_set)}")
     return train_set
 
 
@@ -235,21 +244,19 @@ class OnevsAllDataset(torch.utils.data.Dataset):
 
     def __init__(self, train_set_idx: np.ndarray, entity_idxs):
         super().__init__()
-        assert isinstance(train_set_idx, np.ndarray)
+        assert isinstance(train_set_idx, np.memmap) or isinstance(train_set_idx, np.ndarray)
         assert len(train_set_idx) > 0
-        self.train_data = torch.LongTensor(train_set_idx)
+        self.train_data = train_set_idx
         self.target_dim = len(entity_idxs)
         self.collate_fn = None
-
     def __len__(self):
         return len(self.train_data)
 
     def __getitem__(self, idx):
         y_vec = torch.zeros(self.target_dim)
-        y_vec[self.train_data[idx, 2]] = 1
-        return self.train_data[idx, :2], y_vec
-
-
+        triple= torch.from_numpy(self.train_data[idx].copy()).long()
+        y_vec[triple[2]] = 1
+        return triple[:2], y_vec
 class KvsAll(torch.utils.data.Dataset):
     """ Creates a dataset for KvsAll training by inheriting from torch.utils.data.Dataset.
     Let D denote a dataset for KvsAll training and be defined as D:= {(x,y)_i}_i ^N, where
@@ -291,7 +298,7 @@ def __init__(self, train_set_idx: np.ndarray, entity_idxs, relation_idxs, form,
                  label_smoothing_rate: float = 0.0):
         super().__init__()
         assert len(train_set_idx) > 0
-        assert isinstance(train_set_idx, np.ndarray)
+        assert isinstance(train_set_idx, np.memmap) or isinstance(train_set_idx, np.ndarray)
         self.train_data = None
         self.train_target = None
         self.label_smoothing_rate = torch.tensor(label_smoothing_rate)
@@ -387,7 +394,7 @@ def __init__(self, train_set_idx: np.ndarray, entity_idxs, relation_idxs,
                  label_smoothing_rate=0.0):
         super().__init__()
         assert len(train_set_idx) > 0
-        assert isinstance(train_set_idx, np.ndarray)
+        assert isinstance(train_set_idx, np.memmap) or isinstance(train_set_idx, np.ndarray)
         self.train_data = None
         self.train_target = None
         self.label_smoothing_rate = torch.tensor(label_smoothing_rate)
@@ -430,6 +437,94 @@ def __getitem__(self, idx):
         return self.train_data[idx], y_vec
 
 
+class OnevsSample(torch.utils.data.Dataset):
+    """
+    A custom PyTorch Dataset class for knowledge graph embeddings, which includes
+    both positive and negative sampling for a given dataset for multi-class classification problem..
+
+    Args:
+        train_set (np.ndarray): A numpy array containing triples of knowledge graph data.
+            Each triple consists of (head_entity, relation, tail_entity).
+        num_entities (int): The number of unique entities in the knowledge graph.
+        num_relations (int): The number of unique relations in the knowledge graph.
+        neg_sample_ratio (int, optional): The number of negative samples to be generated
+            per positive sample. Must be a positive integer and less than num_entities.
+        label_smoothing_rate (float, optional): A label smoothing rate to apply to the
+            positive and negative labels. Defaults to 0.0.
+
+    Attributes:
+        train_data (torch.Tensor): The input data converted into a PyTorch tensor.
+        num_entities (int): Number of entities in the dataset.
+        num_relations (int): Number of relations in the dataset.
+        neg_sample_ratio (int): Ratio of negative samples to be drawn for each positive sample.
+        label_smoothing_rate (torch.Tensor): The smoothing factor applied to the labels.
+        collate_fn (function, optional): A function that can be used to collate data samples into
+            batches (set to None by default).
+    """
+    def __init__(self, train_set: np.ndarray, num_entities, num_relations, neg_sample_ratio: int = None,
+                 label_smoothing_rate: float = 0.0):
+        super().__init__()
+        # Input validation
+        assert isinstance(train_set, np.ndarray), "train_set must be a numpy array."
+        assert isinstance(neg_sample_ratio, int), "neg_sample_ratio must be an integer."
+        assert isinstance(num_entities, int) and num_entities > 0, "num_entities must be a positive integer."
+        assert isinstance(num_relations, int) and num_relations > 0, "num_relations must be a positive integer."
+        assert neg_sample_ratio < num_entities, (
+            f"Negative sample ratio {neg_sample_ratio} cannot be larger than the number of entities ({num_entities})."
+        )
+        assert neg_sample_ratio > 0, f"Negative sample ratio {neg_sample_ratio} must be greater than 0."
+
+        # Converting the input numpy array to a PyTorch tensor
+        self.train_data = torch.from_numpy(train_set).long()
+        self.num_entities = num_entities
+        self.num_relations = num_relations
+        self.neg_sample_ratio = neg_sample_ratio
+        self.label_smoothing_rate = torch.tensor(label_smoothing_rate)
+        self.collate_fn = None
+
+    def __len__(self):
+        """Returns the number of samples in the dataset."""
+        return len(self.train_data)
+
+    def __getitem__(self, idx):
+        """
+         Retrieves a single data sample from the dataset at the given index.
+
+         Args:
+             idx (int): The index of the sample to retrieve.
+
+         Returns:
+             tuple: A tuple consisting of:
+                 - x (torch.Tensor): The head and relation part of the triple.
+                 - y_idx (torch.Tensor): The concatenated indices of the true object (tail entity)
+                   and the indices of the negative samples.
+                 - y_vec (torch.Tensor): A vector containing the labels for the positive and
+                   negative samples, with label smoothing applied.
+         """
+        # Retrieve the triple (head, relation, tail) from the training data
+        triple = self.train_data[idx]
+        # Separate the head and relation (x) and the tail entity (y)
+        x = triple[:2]
+        y = triple[-1].unsqueeze(0)  # Tail entity
+
+        # Initialize weights for negative sampling
+        weights = torch.ones(self.num_entities)
+        # Set the weight of the true tail entity to zero
+        weights[y] = 0.0
+
+        # Sample negative examples from the entity set
+        negative_idx = torch.multinomial(weights, num_samples=self.neg_sample_ratio, replacement=False)
+        # Concatenate the true tail entity with the negative samples
+        y_idx = torch.cat((y, negative_idx), 0).long()
+
+        # Create a label vector with smoothing for the true and negative examples
+        y_vec = torch.cat(
+            (torch.ones(1) - self.label_smoothing_rate,  # Positive label with smoothing
+             torch.zeros(self.neg_sample_ratio) + self.label_smoothing_rate),  # Negative labels with smoothing
+            0)
+        return x, y_idx, y_vec
+
+
 class KvsSampleDataset(torch.utils.data.Dataset):
     """
     KvsSample a Dataset:
@@ -459,70 +554,50 @@ class KvsSampleDataset(torch.utils.data.Dataset):
        torch.utils.data.Dataset
        """
 
-    def __init__(self, train_set: np.ndarray, num_entities, num_relations, neg_sample_ratio: int = None,
-                 label_smoothing_rate: float = 0.0):
+    def __init__(self, train_set_idx: np.ndarray, entity_idxs, relation_idxs, form, store=None,
+                 neg_ratio=None,label_smoothing_rate: float = 0.0):
         super().__init__()
-        assert isinstance(train_set, np.ndarray)
-        assert isinstance(neg_sample_ratio, int)
-        self.train_data = train_set
-        self.num_entities = num_entities
-        self.num_relations = num_relations
-        self.neg_sample_ratio = neg_sample_ratio
+
+        assert len(train_set_idx) > 0
+        assert isinstance(train_set_idx, np.ndarray)
+        assert neg_ratio is not None
+        self.train_data = None
+        self.train_target = None
+        self.neg_ratio = neg_ratio
+        self.num_entities = len(entity_idxs)
         self.label_smoothing_rate = torch.tensor(label_smoothing_rate)
         self.collate_fn = None
+        store = mapping_from_first_two_cols_to_third(train_set_idx)
+        assert len(store) > 0
+        # Keys in store correspond to integer representation (index) of subject and predicate
+        # Values correspond to a list of integer representations of entities.
+        self.train_data = torch.LongTensor(list(store.keys()))
+        self.train_target = list(store.values())
+        self.max_num_of_classes = max([ len(i) for i in self.train_target]) + self.neg_ratio
 
-        if self.neg_sample_ratio == 0:
-            print(f'neg_sample_ratio is {neg_sample_ratio}. It will be set to 10.')
-            self.neg_sample_ratio = 10
-
-        print('Constructing training data...')
-        store = mapping_from_first_two_cols_to_third(train_set)
-        self.train_data = torch.IntTensor(list(store.keys()))
-        # https://pytorch.org/docs/stable/data.html#multi-process-data-loading
-        # TLDL; replace Python objects with non-refcounted representations such as Pandas, Numpy or PyArrow objects
-        # Unsure whether a list of numpy arrays are non-refcounted
-        self.train_target = list([np.array(i) for i in store.values()])
         del store
-        # @TODO: Investigate reference counts of using list of numpy arrays.
-        # import sys
-        # import gc
-        # print(sys.getrefcount(self.train_target))
-        # print(sys.getrefcount(self.train_target[0]))
-        # print(gc.get_referrers(self.train_target))
-        # print(gc.get_referrers(self.train_target[0]))
-
     def __len__(self):
-        assert len(self.train_data) == len(self.train_target)
         return len(self.train_data)
 
     def __getitem__(self, idx):
         # (1) Get i.th unique (head,relation) pair.
         x = self.train_data[idx]
         # (2) Get tail entities given (1).
-        positives_idx = self.train_target[idx]
-        num_positives = len(positives_idx)
-        # (3) Do we need to subsample (2) to create training data points of same size.
-        if num_positives < self.neg_sample_ratio:
-            # (3.1) Take all tail entities as positive examples
-            positives_idx = torch.IntTensor(positives_idx)
-            # (3.2) Generate more negative entities
-            negative_idx = torch.randint(low=0,
-                                         high=self.num_entities,
-                                         size=(self.neg_sample_ratio + self.neg_sample_ratio - num_positives,))
-        else:
-            # (3.1) Subsample positives without replacement.
-            positives_idx = torch.IntTensor(np.random.choice(positives_idx, size=self.neg_sample_ratio, replace=False))
-            # (3.2) Generate random entities.
-            negative_idx = torch.randint(low=0,
-                                         high=self.num_entities,
-                                         size=(self.neg_sample_ratio,))
-        # (5) Create selected indexes.
-        y_idx = torch.cat((positives_idx, negative_idx), 0)
-        # (6) Create binary labels.
-        y_vec = torch.cat((torch.ones(len(positives_idx)), torch.zeros(len(negative_idx))), 0)
+        y = self.train_target[idx]
+        num_positive_class =len(y)
+        num_negative_class = self.max_num_of_classes - num_positive_class
+        # Sample negatives
+        weights = torch.ones(self.num_entities)
+        weights[y] = 0.0
+        negative_idx = torch.multinomial(weights, num_samples=num_negative_class, replacement=True)
+
+        y_idx = torch.cat((torch.LongTensor(y), negative_idx), 0)
+        y_vec = torch.cat((torch.ones(num_positive_class), torch.zeros(num_negative_class)),0)
         return x, y_idx, y_vec
 
 
+
+
 class NegSampleDataset(torch.utils.data.Dataset):
     def __init__(self, train_set: np.ndarray, num_entities: int, num_relations: int, neg_sample_ratio: int = 1):
         assert isinstance(train_set, np.ndarray)
@@ -530,6 +605,7 @@ def __init__(self, train_set: np.ndarray, num_entities: int, num_relations: int,
         # TLDL; replace Python objects with non-refcounted representations such as Pandas, Numpy or PyArrow objects
         self.neg_sample_ratio = torch.tensor(
             neg_sample_ratio)
+        print("from numpy to torch")
         self.train_set = torch.from_numpy(train_set).unsqueeze(1)
         self.length = len(self.train_set)
         self.num_entities = torch.tensor(num_entities)
@@ -589,16 +665,17 @@ class TriplePredictionDataset(torch.utils.data.Dataset):
        torch.utils.data.Dataset
        """
 
-    @timeit
     def __init__(self, train_set: np.ndarray, num_entities: int, num_relations: int, neg_sample_ratio: int = 1,
                  label_smoothing_rate: float = 0.0):
         assert isinstance(train_set, np.ndarray)
         # https://pytorch.org/docs/stable/data.html#multi-process-data-loading
         # TLDL; replace Python objects with non-refcounted representations such as Pandas, Numpy or PyArrow objects
+        
         self.label_smoothing_rate = torch.tensor(label_smoothing_rate)
         self.neg_sample_ratio = torch.tensor(
             neg_sample_ratio)  # 0 Implies that we do not add negative samples. This is needed during testing and validation
-        self.train_set = torch.from_numpy(train_set)
+        #self.train_set = torch.from_numpy(train_set)
+        self.train_set = train_set
         assert num_entities >= max(self.train_set[:, 0]) and num_entities >= max(self.train_set[:, 2])
         self.length = len(self.train_set)
         self.num_entities = torch.tensor(num_entities)
@@ -608,7 +685,7 @@ def __len__(self):
         return self.length
 
     def __getitem__(self, idx):
-        return self.train_set[idx]
+        return torch.from_numpy(self.train_set[idx].copy()).long()
 
     def collate_fn(self, batch: List[torch.Tensor]):
         batch = torch.stack(batch, dim=0)
@@ -616,7 +693,7 @@ def collate_fn(self, batch: List[torch.Tensor]):
         size_of_batch, _ = batch.shape
         assert size_of_batch > 0
         label = torch.ones((size_of_batch,)) - self.label_smoothing_rate
-        corr_entities = torch.randint(0, high=self.num_entities, size=(size_of_batch * self.neg_sample_ratio,))
+        corr_entities = torch.randint(0, high=self.num_entities, size=(size_of_batch * self.neg_sample_ratio,),dtype=torch.long)
         if torch.rand(1) >= 0.5:
             # corrupt head
             r_head_corr = r.repeat(self.neg_sample_ratio, )
diff --git a/dicee/evaluator.py b/dicee/evaluator.py
index ab46393a..b005f20a 100644
--- a/dicee/evaluator.py
+++ b/dicee/evaluator.py
@@ -2,6 +2,8 @@
 import torch
 import numpy as np
 import json
+
+from . import load_json
 from .static_funcs import pickle
 from .static_funcs_training import evaluate_lp, evaluate_bpe_lp
 from typing import Tuple, List
@@ -57,16 +59,6 @@ def vocab_preparation(self, dataset) -> None:
         else:
             self.ee_vocab = dataset.ee_vocab.result()
 
-        """
-        if isinstance(dataset.constraints, tuple):
-            self.domain_constraints_per_rel, self.range_constraints_per_rel = dataset.constraints
-        else:
-            try:
-                self.domain_constraints_per_rel, self.range_constraints_per_rel = dataset.constraints.result()
-            except RuntimeError:
-                print('Domain constraint exception occurred')
-        """
-
         self.num_entities = dataset.num_entities
         self.num_relations = dataset.num_relations
         self.func_triple_to_bpe_representation = dataset.func_triple_to_bpe_representation
@@ -77,6 +69,7 @@ def vocab_preparation(self, dataset) -> None:
 
     # @timeit
     def eval(self, dataset: KG, trained_model, form_of_labelling, during_training=False) -> None:
+        assert isinstance(dataset, KG), "dataset must be KG"
         # @TODO: Why this reassigment ?
         self.during_training = during_training
         # (1) Exit, if the flag is not set
@@ -97,7 +90,7 @@ def eval(self, dataset: KG, trained_model, form_of_labelling, during_training=Fa
                                                                      trained_model=trained_model)
         elif self.args.scoring_technique in ["AllvsAll",
                                              "KvsAll",
-                                             'KvsSample',
+                                             '1vsSample',
                                              "1vsAll"] and self.args.byte_pair_encoding:
             if self.args.model != "BytE":
                 self.eval_with_bpe_vs_all(raw_train_set=dataset.raw_train_set,
@@ -116,7 +109,7 @@ def eval(self, dataset: KG, trained_model, form_of_labelling, during_training=Fa
                                                    valid_set=dataset.valid_set,
                                                    test_set=dataset.test_set,
                                                    trained_model=trained_model)
-        elif self.args.scoring_technique in ["AllvsAll", "KvsAll", 'KvsSample', "1vsAll"]:
+        elif self.args.scoring_technique in ["AllvsAll", "KvsAll", '1vsSample',"KvsSample", "1vsAll"]:
             self.eval_with_vs_all(train_set=dataset.train_set,
                                   valid_set=dataset.valid_set,
                                   test_set=dataset.test_set,
@@ -150,6 +143,9 @@ def __load_and_set_mappings(self):
         self.er_vocab = pickle.load(open(self.args.full_storage_path + "/er_vocab.p", "rb"))
         self.re_vocab = pickle.load(open(self.args.full_storage_path + "/re_vocab.p", "rb"))
         self.ee_vocab = pickle.load(open(self.args.full_storage_path + "/ee_vocab.p", "rb"))
+        report = load_json(self.args.full_storage_path + "/report.json")
+        self.num_entities = report["num_entities"]
+        self.num_relations = report["num_relations"]
 
     def eval_rank_of_head_and_tail_entity(self, *, train_set, valid_set=None, test_set=None, trained_model):
         # 4. Test model on the training dataset if it is needed.
@@ -346,11 +342,6 @@ def evaluate_lp_with_byte(self, model, triples: List[List[str]], info=None):
         import tiktoken
 
         enc = tiktoken.get_encoding("gpt2")
-
-        ranks = []
-        # Hit range
-        hits_range = [i for i in range(1, 11)]
-        hits = {i: [] for i in hits_range}
         if info and self.during_training is False:
             print(info + ':', end=' ')
         # Iterate over integer indexed triples in mini batch fashion
@@ -441,11 +432,17 @@ def evaluate_lp_bpe_k_vs_all(self, model, triples: List[List[str]], info=None, f
         return results
 
     def evaluate_lp(self, model, triple_idx, info: str):
-        return evaluate_lp(model, triple_idx, num_entities=self.num_entities, er_vocab=self.er_vocab,
+        assert self.num_entities is not None, "self.num_entities cannot be None"
+        assert self.er_vocab is not None, "self.er_vocab cannot be None"
+        assert self.re_vocab is not None, "self.re_vocab cannot be None"
+        return evaluate_lp(model, triple_idx,
+                           num_entities=self.num_entities,
+                           er_vocab=self.er_vocab,
                            re_vocab=self.re_vocab, info=info)
 
     def dummy_eval(self, trained_model, form_of_labelling: str):
-
+        assert trained_model is not None
+        # @TODO:CD: Why such naming! We need to document it better.
         if self.is_continual_training:
             self.__load_and_set_mappings()
 
@@ -456,13 +453,14 @@ def dummy_eval(self, trained_model, form_of_labelling: str):
                                                    valid_set=valid_set,
                                                    test_set=test_set,
                                                    trained_model=trained_model)
-        elif self.args.scoring_technique in ["AllvsAll",'KvsAll', 'KvsSample', '1vsAll']:
+        elif self.args.scoring_technique in ["AllvsAll",'KvsAll', '1vsSample', "KvsSample", '1vsAll']:
             self.eval_with_vs_all(train_set=train_set,
                                   valid_set=valid_set,
                                   test_set=test_set,
                                   trained_model=trained_model, form_of_labelling=form_of_labelling)
         else:
             raise ValueError(f'Invalid argument: {self.args.scoring_technique}')
+
         with open(self.args.full_storage_path + '/eval_report.json', 'w') as file_descriptor:
             json.dump(self.report, file_descriptor, indent=4)
 
@@ -474,7 +472,7 @@ def eval_with_data(self, dataset, trained_model, triple_idx: np.ndarray, form_of
             return self.evaluate_lp(trained_model, triple_idx,
                                     info=f'Evaluate {trained_model.name} on a given dataset', )
 
-        elif self.args.scoring_technique in ['KvsAll', 'KvsSample', '1vsAll', 'PvsAll', 'CCvsAll']:
+        elif self.args.scoring_technique in ['KvsAll', '1vsSample', '1vsAll', 'PvsAll', 'CCvsAll']:
             return self.evaluate_lp_k_vs_all(trained_model, triple_idx,
                                              info=f'Evaluate {trained_model.name} on a given dataset',
                                              form_of_labelling=form_of_labelling)
diff --git a/dicee/executer.py b/dicee/executer.py
index 7be7fbef..ecc4bf0d 100644
--- a/dicee/executer.py
+++ b/dicee/executer.py
@@ -6,13 +6,12 @@
 import os
 import datetime
 from pytorch_lightning import seed_everything
-
-from dicee.knowledge_graph import KG
-from dicee.evaluator import Evaluator
-# Avoid
-from dicee.static_preprocess_funcs import preprocesses_input_args
-from dicee.trainer import DICE_Trainer
-from dicee.static_funcs import timeit, continual_training_setup_executor, read_or_load_kg, load_json, store
+from .knowledge_graph import KG
+from .evaluator import Evaluator
+from .static_preprocess_funcs import preprocesses_input_args
+from .trainer import DICE_Trainer
+from .static_funcs import timeit, read_or_load_kg, load_json, store, create_experiment_folder
+import numpy as np
 
 logging.getLogger('pytorch_lightning').setLevel(0)
 warnings.filterwarnings(action="ignore", category=DeprecationWarning)
@@ -34,7 +33,7 @@ def __init__(self, args, continuous_training=False):
         # (3) Set the continual training flag
         self.is_continual_training = continuous_training
         # (4) Create an experiment folder or use the previous one
-        continual_training_setup_executor(self)
+        self.setup_executor()
         # (5) A variable is initialized for pytorch lightning trainer or DICE_Trainer()
         self.trainer = None
         self.trained_model = None
@@ -42,78 +41,26 @@ def __init__(self, args, continuous_training=False):
         self.knowledge_graph = None
         # (7) Store few data in memory for numerical results, e.g. runtime, H@1 etc.
         self.report = dict()
-        # (8) Create an object to carry out link prediction evaluations
-        self.evaluator = None  # e.g. Evaluator(self)
+        # (8) Create an object to carry out link prediction evaluations,  e.g. Evaluator(self)
+        self.evaluator = None
         # (9) Execution start time
         self.start_time = None
 
-    def read_or_load_kg(self):
-        print('*** Read or Load Knowledge Graph  ***')
-        start_time = time.time()
-        kg = KG(dataset_dir=self.args.dataset_dir,
-                byte_pair_encoding=self.args.byte_pair_encoding,
-                padding=True if self.args.byte_pair_encoding and self.args.model != "BytE" else False,
-                add_noise_rate=self.args.add_noise_rate,
-                sparql_endpoint=self.args.sparql_endpoint,
-                path_single_kg=self.args.path_single_kg,
-                add_reciprical=self.args.apply_reciprical_or_noise,
-                eval_model=self.args.eval_model,
-                read_only_few=self.args.read_only_few,
-                sample_triples_ratio=self.args.sample_triples_ratio,
-                path_for_serialization=self.args.full_storage_path,
-                path_for_deserialization=self.args.path_experiment_folder if hasattr(self.args,
-                                                                                     'path_experiment_folder') else None,
-                backend=self.args.backend,
-                training_technique=self.args.scoring_technique)
-        print(f'Preprocessing took: {time.time() - start_time:.3f} seconds')
-        # (2) Share some info about data for easy access.
-        print(kg.description_of_input)
-        return kg
-
-    def read_preprocess_index_serialize_data(self) -> None:
-        """ Read & Preprocess & Index & Serialize Input Data
-
-        (1) Read or load the data from disk into memory.
-        (2) Store the statistics of the data.
-
-        Parameter
-        ----------
-
-        Return
-        ----------
-        None
-
-        """
-        # (1) Read & Preprocess & Index & Serialize Input Data.
-        self.knowledge_graph = self.read_or_load_kg()
-
-        # (2) Store the stats and share parameters
-        self.args.num_entities = self.knowledge_graph.num_entities
-        self.args.num_relations = self.knowledge_graph.num_relations
-        self.args.num_tokens = self.knowledge_graph.num_tokens
-        self.args.max_length_subword_tokens = self.knowledge_graph.max_length_subword_tokens
-        self.args.ordered_bpe_entities=self.knowledge_graph.ordered_bpe_entities
-        self.report['num_train_triples'] = len(self.knowledge_graph.train_set)
-        self.report['num_entities'] = self.knowledge_graph.num_entities
-        self.report['num_relations'] = self.knowledge_graph.num_relations
-        self.report['num_relations'] = self.knowledge_graph.num_relations
-        self.report[
-            'max_length_subword_tokens'] = self.knowledge_graph.max_length_subword_tokens if self.knowledge_graph.max_length_subword_tokens else None
-
-        self.report['runtime_kg_loading'] = time.time() - self.start_time
-
-    def load_indexed_data(self) -> None:
-        """ Load the indexed data from disk into memory
-
-        Parameter
-        ----------
-
-        Return
-        ----------
-        None
-
-        """
-        self.knowledge_graph = read_or_load_kg(self.args, cls=KG)
+    def setup_executor(self) -> None:
+        if self.is_continual_training is False:
+            # Create a single directory containing KGE and all related data
+            if self.args.path_to_store_single_run is not None:
+                if os.path.exists(self.args.path_to_store_single_run):
+                    print(f"Deleting the existing directory of {self.args.path_to_store_single_run}")
+                    os.system(f'rm -rf {self.args.path_to_store_single_run}')
+                os.makedirs(self.args.path_to_store_single_run, exist_ok=False)
+                self.args.full_storage_path = self.args.path_to_store_single_run
+            else:
+                self.args.full_storage_path = create_experiment_folder(folder_name=self.args.storage_path)
+
+            with open(self.args.full_storage_path + '/configuration.json', 'w') as file_descriptor:
+                temp = vars(self.args)
+                json.dump(temp, file_descriptor, indent=3)
 
     @timeit
     def save_trained_model(self) -> None:
@@ -140,21 +87,19 @@ def save_trained_model(self) -> None:
         self.report.update(self.trained_model.mem_of_model())
         # (3) Store/Serialize Model for further use.
         if self.is_continual_training is False:
-            store(trainer=self.trainer,
-                  trained_model=self.trained_model,
+            store(trained_model=self.trained_model,
                   model_name='model',
-                  full_storage_path=self.storage_path,
+                  full_storage_path=self.args.full_storage_path,
                   save_embeddings_as_csv=self.args.save_embeddings_as_csv)
         else:
-            store(trainer=self.trainer,
-                  trained_model=self.trained_model,
-                  model_name='model_' + str(datetime.datetime.now()),
-                  full_storage_path=self.storage_path, save_embeddings_as_csv=self.args.save_embeddings_as_csv)
+            store(trained_model=self.trained_model,
+                  model_name='model', # + str(datetime.datetime.now()),
+                  full_storage_path=self.args.full_storage_path,
+                  save_embeddings_as_csv=self.args.save_embeddings_as_csv)
 
-        self.report['path_experiment_folder'] = self.storage_path
+        self.report['path_experiment_folder'] = self.args.full_storage_path
         self.report['num_entities'] = self.args.num_entities
         self.report['num_relations'] = self.args.num_relations
-        self.report['path_experiment_folder'] = self.storage_path
 
     def end(self, form_of_labelling: str) -> dict:
         """
@@ -181,13 +126,15 @@ def end(self, form_of_labelling: str) -> dict:
             self.write_report()
             return {**self.report}
         else:
-            self.evaluator.eval(dataset=self.knowledge_graph, trained_model=self.trained_model,
+            self.evaluator.eval(dataset=self.knowledge_graph,
+                                trained_model=self.trained_model,
                                 form_of_labelling=form_of_labelling)
             self.write_report()
             return {**self.report, **self.evaluator.report}
 
     def write_report(self) -> None:
         """ Report training related information in a report.json file """
+        # @TODO: Move to static funcs
         # Report total runtime.
         self.report['Runtime'] = time.time() - self.start_time
         print(f"Total Runtime: {self.report['Runtime']:.3f} seconds")
@@ -213,15 +160,45 @@ def start(self) -> dict:
         """
         self.start_time = time.time()
         print(f"Start time:{datetime.datetime.now()}")
-        # (1) Loading the Data
-        #  Load the indexed data from disk or read a raw data from disk into knowledge_graph attribute
-        self.load_indexed_data() if self.is_continual_training else self.read_preprocess_index_serialize_data()
+        # (1) Reload the memory-map of index knowledge graph stored as a numpy ndarray.
+        if self.args.path_to_store_single_run and os.path.exists(self.args.path_to_store_single_run+"/memory_map_train_set.npy"):
+            # (1.1) Read information about memory-map of KG.
+            with open(self.args.path_to_store_single_run+'/memory_map_details.json', 'r') as file_descriptor:
+                memory_map_details = json.load(file_descriptor)
+            self.knowledge_graph = np.memmap(self.args.path_to_store_single_run + '/memory_map_train_set.npy',
+                                             mode='r',
+                                             dtype=memory_map_details["dtype"],
+                                             shape=tuple(memory_map_details["shape"]))
+            self.args.num_entities = memory_map_details["num_entities"]
+            self.args.num_relations = memory_map_details["num_relations"]
+            self.args.num_tokens = None
+            self.args.max_length_subword_tokens = None
+            self.args.ordered_bpe_entities = None
+        else:
+            self.knowledge_graph = read_or_load_kg(self.args, cls=KG)
+            self.args.num_entities = self.knowledge_graph.num_entities
+            self.args.num_relations = self.knowledge_graph.num_relations
+            self.args.num_tokens = self.knowledge_graph.num_tokens
+            self.args.max_length_subword_tokens = self.knowledge_graph.max_length_subword_tokens
+            self.args.ordered_bpe_entities = self.knowledge_graph.ordered_bpe_entities
+            self.report['num_train_triples'] = len(self.knowledge_graph.train_set)
+            self.report['num_entities'] = self.knowledge_graph.num_entities
+            self.report['num_relations'] = self.knowledge_graph.num_relations
+            self.report['max_length_subword_tokens'] = self.knowledge_graph.max_length_subword_tokens if self.knowledge_graph.max_length_subword_tokens else None
+            self.report['runtime_kg_loading'] = time.time() - self.start_time
+            data={"shape":tuple(self.knowledge_graph.train_set.shape),
+                  "dtype":self.knowledge_graph.train_set.dtype.str,
+                  "num_entities":self.knowledge_graph.num_entities,
+                  "num_relations":self.knowledge_graph.num_relations}
+            with open(self.args.full_storage_path + '/memory_map_details.json', 'w') as file_descriptor:
+                json.dump(data, file_descriptor, indent=4)
+
         # (2) Create an evaluator object.
         self.evaluator = Evaluator(args=self.args)
         # (3) Create a trainer object.
         self.trainer = DICE_Trainer(args=self.args,
                                     is_continual_training=self.is_continual_training,
-                                    storage_path=self.storage_path,
+                                    storage_path=self.args.full_storage_path,
                                     evaluator=self.evaluator)
         # (4) Start the training
         self.trained_model, form_of_labelling = self.trainer.start(knowledge_graph=self.knowledge_graph)
@@ -242,7 +219,7 @@ class ContinuousExecute(Execute):
 
     def __init__(self, args):
         # (1) Current input configuration.
-        assert os.path.exists(args.continual_learning)
+        assert os.path.exists(args.continual_learning), f"Path doesn't exist {args.continual_learning}"
         assert os.path.isfile(args.continual_learning + '/configuration.json')
         # (2) Load previous input configuration.
         previous_args = load_json(args.continual_learning + '/configuration.json')
@@ -279,11 +256,26 @@ def continual_start(self) -> dict:
 
         """
         # (1)
-        self.trainer = DICE_Trainer(args=self.args, is_continual_training=True,
+        self.trainer = DICE_Trainer(args=self.args,
+                                    is_continual_training=True,
                                     storage_path=self.args.continual_learning)
         # (2)
-        self.trained_model, form_of_labelling = self.trainer.continual_start()
 
+        assert os.path.exists(f"{self.args.continual_learning}/memory_map_train_set.npy")
+        # (1) Reload the memory-map of index knowledge graph stored as a numpy ndarray.
+        with open(f"{self.args.continual_learning}/memory_map_details.json", 'r') as file_descriptor:
+            memory_map_details = json.load(file_descriptor)
+        knowledge_graph = np.memmap(f"{self.args.continual_learning}/memory_map_train_set.npy",
+                                         mode='r',
+                                         dtype=memory_map_details["dtype"],
+                                         shape=tuple(memory_map_details["shape"]))
+        self.args.num_entities = memory_map_details["num_entities"]
+        self.args.num_relations = memory_map_details["num_relations"]
+        self.args.num_tokens = None
+        self.args.max_length_subword_tokens = None
+        self.args.ordered_bpe_entities = None
+
+        self.trained_model, form_of_labelling = self.trainer.continual_start(knowledge_graph)
         # (5) Store trained model.
         self.save_trained_model()
         # (6) Eval model.
@@ -292,4 +284,4 @@ def continual_start(self) -> dict:
         else:
             self.evaluator = Evaluator(args=self.args, is_continual_training=True)
             self.evaluator.dummy_eval(self.trained_model, form_of_labelling)
-            return {**self.report, **self.evaluator.report}
\ No newline at end of file
+            return {**self.report, **self.evaluator.report}
diff --git a/dicee/knowledge_graph.py b/dicee/knowledge_graph.py
index 8a410773..25aa3259 100644
--- a/dicee/knowledge_graph.py
+++ b/dicee/knowledge_graph.py
@@ -1,7 +1,8 @@
 from typing import List
 from .read_preprocess_save_load_kg import ReadFromDisk, PreprocessKG, LoadSaveToDisk
 import sys
-
+import pandas as pd
+import polars as pl
 class KG:
     """ Knowledge Graph """
 
@@ -12,10 +13,10 @@ def __init__(self, dataset_dir: str = None,
                  sparql_endpoint: str = None,
                  path_single_kg: str = None,
                  path_for_deserialization: str = None,
-                 add_reciprical: bool = None, eval_model: str = None,
+                 add_reciprocal: bool = None, eval_model: str = None,
                  read_only_few: int = None, sample_triples_ratio: float = None,
                  path_for_serialization: str = None,
-                 entity_to_idx=None, relation_to_idx=None, backend=None, training_technique: str = None):
+                 entity_to_idx=None, relation_to_idx=None, backend=None, training_technique: str = None, separator:str=None):
         """
         :param dataset_dir: A path of a folder containing train.txt, valid.txt, test.text
         :param byte_pair_encoding: Apply Byte pair encoding.
@@ -25,7 +26,7 @@ def __init__(self, dataset_dir: str = None,
         :param path_single_kg: The path of a single file containing the input knowledge graph
         :param path_for_deserialization: A path of a folder containing previously parsed data
         :param num_core: Number of subprocesses used for data loading
-        :param add_reciprical: A flag for applying reciprocal data augmentation technique
+        :param add_reciprocal: A flag for applying reciprocal data augmentation technique
         :param eval_model: A flag indicating whether evaluation will be applied.
         If no eval, then entity relation mappings will be deleted to free memory.
         :param add_noise_rate: Add say 10% noise in the input data
@@ -33,15 +34,16 @@ def __init__(self, dataset_dir: str = None,
         :param training_technique
         """
         self.dataset_dir = dataset_dir
+        self.sparql_endpoint = sparql_endpoint
+        self.path_single_kg = path_single_kg
+
         self.byte_pair_encoding = byte_pair_encoding
         self.ordered_shaped_bpe_tokens = None
-        self.sparql_endpoint = sparql_endpoint
         self.add_noise_rate = add_noise_rate
         self.num_entities = None
         self.num_relations = None
-        self.path_single_kg = path_single_kg
         self.path_for_deserialization = path_for_deserialization
-        self.add_reciprical = add_reciprical
+        self.add_reciprocal = add_reciprocal
         self.eval_model = eval_model
 
         self.read_only_few = read_only_few
@@ -69,21 +71,33 @@ def __init__(self, dataset_dir: str = None,
         self.target_dim = None
         self.train_target_indices = None
         self.ordered_bpe_entities = None
+        self.separator=separator
 
         if self.path_for_deserialization is None:
+            # Read a knowledge graph into memory
             ReadFromDisk(kg=self).start()
+            # Map a knowledge graph into integer indexed.
             PreprocessKG(kg=self).start()
+            # Saving.
             LoadSaveToDisk(kg=self).save()
-
         else:
             LoadSaveToDisk(kg=self).load()
+        assert len(self.train_set) > 0, "Training set is empty"
+        self.description_of_input=None
+        self.describe()
 
-        assert len(self.train_set) > 0
-
-        self._describe()
+        if self.entity_to_idx is not None:
+            assert isinstance(self.entity_to_idx, dict) or isinstance(self.entity_to_idx, pd.DataFrame) or isinstance(self.entity_to_idx,
+                                                                      pl.DataFrame), f"entity_to_idx must be a dict or a pandas/polars DataFrame: {type(self.entity_to_idx)}"
+            # TODO:CD: Why do we need to create this inverse mapping at this point?
+            if isinstance(self.entity_to_idx, dict):
+                self.idx_to_entity = {v: k for k, v in self.entity_to_idx.items()}
+                self.idx_to_relations = {v: k for k, v in self.relation_to_idx.items()}
+            else:
+                pass
 
-    def _describe(self) -> None:
-        self.description_of_input = f'\n------------------- Description of Dataset {self.dataset_dir} -------------------'
+    def describe(self) -> None:
+        self.description_of_input = f'\n------------------- Description of Dataset {self.dataset_dir if isinstance(self.dataset_dir, str) else self.sparql_endpoint if isinstance(self.sparql_endpoint, str) else self.path_single_kg} -------------------'
         if self.byte_pair_encoding:
             self.description_of_input += f'\nNumber of tokens:{self.num_tokens}' \
                                          f'\nNumber of max sequence of sub-words: {self.max_length_subword_tokens}' \
@@ -113,6 +127,18 @@ def entities_str(self) -> List:
     def relations_str(self) -> List:
         return list(self.relation_to_idx.keys())
 
+    def exists(self,h:str,r:str,t:str):
+        # Row to check for existence
+        row_to_check = {'subject': self.entity_to_idx[h], 'relation': self.relation_to_idx[r], 'object': self.entity_to_idx[t]}
+        # Check if the row exists
+        return ((self.raw_train_set == pd.Series(row_to_check)).all(axis=1)).any()
+
+    def __iter__(self):
+        for h, r, t in self.raw_train_set.to_numpy().tolist():
+            yield self.idx_to_entity[h], self.idx_to_relations[r], self.idx_to_entity[t]
+    def __len__(self):
+        return len(self.raw_train_set)
+
     def func_triple_to_bpe_representation(self, triple: List[str]):
         result = []
 
diff --git a/dicee/knowledge_graph_embeddings.py b/dicee/knowledge_graph_embeddings.py
index 18b3ca89..03d55b8e 100644
--- a/dicee/knowledge_graph_embeddings.py
+++ b/dicee/knowledge_graph_embeddings.py
@@ -9,9 +9,6 @@
 from .static_funcs_training import evaluate_lp
 import numpy as np
 import sys
-
-# import gradio as gr
-
 import traceback
 
 
@@ -19,8 +16,7 @@ class KGE(BaseInteractiveKGE):
     """ Knowledge Graph Embedding Class for interactive usage of pre-trained models"""
 
     def __init__(self, path=None, url=None, construct_ensemble=False,
-                 model_name=None,
-                 apply_semantic_constraint=False):
+                 model_name=None):
         super().__init__(path=path, url=url, construct_ensemble=construct_ensemble, model_name=model_name)
 
     def __str__(self):
@@ -345,7 +341,9 @@ def predict(self, *, h: Union[List[str], str] = None, r: Union[List[str], str] =
         else:
             return torch.sigmoid(scores)
 
-    def predict_topk(self, *, h: List[str] = None, r: List[str] = None, t: List[str] = None,
+    def predict_topk(self, *, h: Union[str, List[str]] = None,
+                     r: Union[str, List[str]] = None,
+                     t: Union[str, List[str]] = None,
                      topk: int = 10, within: List[str] = None):
         """
         Predict missing item in a given triple.
@@ -354,15 +352,15 @@ def predict_topk(self, *, h: List[str] = None, r: List[str] = None, t: List[str]
 
         Parameter
         ---------
-        head_entity: List[str]
+        head_entity: Union[str, List[str]]
 
         String representation of selected entities.
 
-        relation: List[str]
+        relation: Union[str, List[str]]
 
         String representation of selected relations.
 
-        tail_entity: List[str]
+        tail_entity: Union[str, List[str]]
 
         String representation of selected entities.
 
@@ -379,22 +377,17 @@ def predict_topk(self, *, h: List[str] = None, r: List[str] = None, t: List[str]
 
         # (1) Sanity checking.
         if h is not None:
-            assert isinstance(h, list)
+            assert isinstance(h, list) or isinstance(h,str)
         if r is not None:
-            assert isinstance(r, list)
+            assert isinstance(r, list) or isinstance(r,str)
         if t is not None:
-            assert isinstance(t, list)
+            assert isinstance(t, list) or isinstance(t,str)
         # (2) Predict missing head entity given a relation and a tail entity.
         if h is None:
             assert r is not None
             assert t is not None
             # ? r, t
             scores = self.predict_missing_head_entity(r, t, within=within).flatten()
-            if self.apply_semantic_constraint:
-                # filter the scores
-                for th, i in enumerate(r):
-                    scores[self.domain_constraints_per_rel[self.relation_to_idx[i]]] = -torch.inf
-
             sort_scores, sort_idxs = torch.topk(scores, topk)
             return [(self.idx_to_entity[idx_top_entity], scores.item()) for idx_top_entity, scores in
                     zip(sort_idxs.tolist(), torch.sigmoid(sort_scores))]
@@ -415,10 +408,6 @@ def predict_topk(self, *, h: List[str] = None, r: List[str] = None, t: List[str]
             assert r is not None
             # h r ?t
             scores = self.predict_missing_tail_entity(h, r, within=within).flatten()
-            if self.apply_semantic_constraint:
-                # filter the scores
-                for th, i in enumerate(r):
-                    scores[self.range_constraints_per_rel[self.relation_to_idx[i]]] = -torch.inf
             sort_scores, sort_idxs = torch.topk(scores, topk)
             return [(self.idx_to_entity[idx_top_entity], scores.item()) for idx_top_entity, scores in
                     zip(sort_idxs.tolist(), torch.sigmoid(sort_scores))]
diff --git a/dicee/models/__init__.py b/dicee/models/__init__.py
index a8240a92..99da892f 100644
--- a/dicee/models/__init__.py
+++ b/dicee/models/__init__.py
@@ -3,7 +3,7 @@
 from .complex import * # noqa
 from .quaternion import * # noqa
 from .octonion import * # noqa
-from .clifford import Keci, KeciBase, CMult, DeCaL # noqa
+from .clifford import Keci, KeciBase, DeCaL # noqa
 from .pykeen_models import * # noqa
 from .function_space import * # noqa
-from .dualE import DualE
+from .dualE import DualE # noqa
diff --git a/dicee/models/adopt.py b/dicee/models/adopt.py
new file mode 100644
index 00000000..54cc59d9
--- /dev/null
+++ b/dicee/models/adopt.py
@@ -0,0 +1,536 @@
+# CD: copy pasted from https://raw.githubusercontent.com/iShohei220/adopt/refs/heads/main/adopt.py
+# mypy: allow-untyped-decorators
+# mypy: allow-untyped-defs
+from typing import cast, Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from torch.optim.optimizer import (
+    _capturable_doc, # noqa: F401
+    _default_to_fused_or_foreach,
+    _device_dtype_check_for_fused,
+    _differentiable_doc, # noqa: F401
+    _disable_dynamo_if_unsupported,
+    _foreach_doc, # noqa: F401
+    _fused_doc, # noqa: F401
+    _get_capturable_supported_devices,
+    _get_scalar_dtype,
+    _get_value, # noqa: F401
+    _maximize_doc, # noqa: F401
+    _stack_if_compiling, # noqa: F401
+    _use_grad_for_differentiable,
+    _view_as_real,
+    DeviceDict, # noqa: F401
+    Optimizer,
+    ParamsT,
+)
+
+
+__all__ = ["ADOPT", "adopt"]
+
+
+class ADOPT(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.9999),
+        eps: float = 1e-6,
+        clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25,
+        weight_decay: float = 0.0,
+        decouple: bool = False,
+        *,
+        foreach: Optional[bool] = None,
+        maximize: bool = False,
+        capturable: bool = False,
+        differentiable: bool = False,
+        fused: Optional[bool] = None,
+    ):
+        if isinstance(lr, Tensor):
+            if foreach and not capturable:
+                raise ValueError(
+                    "lr as a Tensor is not supported for capturable=False and foreach=True"
+                )
+            if lr.numel() != 1:
+                raise ValueError("Tensor lr must be 1-element")
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        self.clip_lambda = clip_lambda
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            decouple=decouple,
+            maximize=maximize,
+            foreach=foreach,
+            capturable=capturable,
+            differentiable=differentiable,
+            fused=fused,
+        )
+        super().__init__(params, defaults)
+
+        if fused:
+            # TODO: support fused
+            raise RuntimeError("`fused` is not currently supported")
+
+            if differentiable:
+                raise RuntimeError("`fused` does not support `differentiable`")
+            self._step_supports_amp_scaling = True
+            # TODO(crcrpar): [low prec params & their higher prec copy]
+            # Support AMP with FP16/BF16 model params which would need
+            # higher prec copy of params to do update math in higher prec to
+            # alleviate the loss of information.
+            if foreach:
+                raise RuntimeError("`fused` and `foreach` cannot be `True` together.")
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("maximize", False)
+            group.setdefault("foreach", None)
+            group.setdefault("capturable", False)
+            group.setdefault("differentiable", False)
+            fused = group.setdefault("fused", None)
+            for p in group["params"]:
+                p_state = self.state.get(p, [])
+                if len(p_state) != 0 and not torch.is_tensor(p_state["step"]):
+                    step_val = float(p_state["step"])
+                    p_state["step"] = (
+                        torch.tensor(
+                            step_val,
+                            dtype=_get_scalar_dtype(is_fused=fused),
+                            device=p.device,
+                        )
+                        if group["capturable"] or group["fused"]
+                        else torch.tensor(step_val, dtype=_get_scalar_dtype())
+                    )
+
+    def _init_group(
+        self,
+        group,
+        params_with_grad,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        state_steps,
+    ):
+        has_complex = False
+        for p in group["params"]:
+            if p.grad is not None:
+                has_complex |= torch.is_complex(p)
+                params_with_grad.append(p)
+                if p.grad.is_sparse:
+                    raise RuntimeError(
+                        "ADOPT does not support sparse gradients"
+                    )
+                grads.append(p.grad)
+
+                state = self.state[p]
+                # Lazy state initialization
+                if len(state) == 0:
+                    if group["fused"]:
+                        _device_dtype_check_for_fused(p)
+                    # note(crcrpar): [special device hosting for step]
+                    # Deliberately host `step` on CPU if both capturable and fused are off.
+                    # This is because kernel launches are costly on CUDA and XLA.
+                    state["step"] = (
+                        torch.zeros(
+                            (),
+                            dtype=_get_scalar_dtype(is_fused=group["fused"]),
+                            device=p.device,
+                        )
+                        if group["capturable"] or group["fused"]
+                        else torch.tensor(0.0, dtype=_get_scalar_dtype())
+                    )
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+
+                exp_avgs.append(state["exp_avg"])
+                exp_avg_sqs.append(state["exp_avg_sq"])
+
+                if group["differentiable"] and state["step"].requires_grad:
+                    raise RuntimeError(
+                        "`requires_grad` is not supported for `step` in differentiable mode"
+                    )
+
+                # Foreach without capturable does not support a tensor lr
+                if (
+                    group["foreach"]
+                    and torch.is_tensor(group["lr"])
+                    and not group["capturable"]
+                ):
+                    raise RuntimeError(
+                        "lr as a Tensor is not supported for capturable=False and foreach=True"
+                    )
+
+                state_steps.append(state["step"])
+        return has_complex
+
+    @_use_grad_for_differentiable
+    def step(self, closure=None):
+        """Perform a single optimization step.
+
+        Args:
+            closure (Callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        self._cuda_graph_capture_health_check()
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad: List[Tensor] = []
+            grads: List[Tensor] = []
+            exp_avgs: List[Tensor] = []
+            exp_avg_sqs: List[Tensor] = []
+            state_steps: List[Tensor] = []
+            beta1, beta2 = group["betas"]
+
+            has_complex = self._init_group(
+                group,
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+            )
+
+            adopt(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                has_complex=has_complex,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                clip_lambda=self.clip_lambda,
+                weight_decay=group["weight_decay"],
+                decouple=group["decouple"],
+                eps=group["eps"],
+                maximize=group["maximize"],
+                foreach=group["foreach"],
+                capturable=group["capturable"],
+                differentiable=group["differentiable"],
+                fused=group["fused"],
+                grad_scale=getattr(self, "grad_scale", None),
+                found_inf=getattr(self, "found_inf", None),
+            )
+
+        return loss
+
+
+def _single_tensor_adopt(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    grad_scale: Optional[Tensor],
+    found_inf: Optional[Tensor],
+    *,
+    has_complex: bool,
+    beta1: float,
+    beta2: float,
+    lr: Union[float, Tensor],
+    clip_lambda: Optional[Callable[[int], float]],
+    weight_decay: float,
+    decouple: bool,
+    eps: float,
+    maximize: bool,
+    capturable: bool,
+    differentiable: bool,
+):
+    assert grad_scale is None and found_inf is None
+
+    if torch.jit.is_scripting():
+        # this assert is due to JIT being dumb and not realizing that the ops below
+        # have overloads to handle both float and Tensor lrs, so we just assert it's
+        # a float since most people using JIT are using floats
+        assert isinstance(lr, float)
+
+    for i, param in enumerate(params):
+        grad = grads[i] if not maximize else -grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        step_t = state_steps[i]
+
+        # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable]
+        if not torch._utils.is_compiling() and capturable:
+            capturable_supported_devices = _get_capturable_supported_devices()
+            assert (
+                param.device.type == step_t.device.type
+                and param.device.type in capturable_supported_devices
+            ), f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}."
+
+        step = step_t if capturable or differentiable else _get_value(step_t)
+
+        if weight_decay != 0 and not decouple:
+            grad = grad.add(param, alpha=weight_decay)
+
+        if torch.is_complex(param):
+            grad = torch.view_as_real(grad)
+            if exp_avg is not None:
+                exp_avg = torch.view_as_real(exp_avg)
+            if exp_avg_sq is not None:
+                exp_avg_sq = torch.view_as_real(exp_avg_sq)
+            param = torch.view_as_real(param)
+
+        if step == 0:
+            exp_avg_sq.addcmul_(grad, grad.conj())
+            # update step
+            step_t += 1
+            continue
+
+        if weight_decay != 0 and decouple:
+            param.add_(param, alpha=-lr*weight_decay)
+
+        denom = torch.clamp(exp_avg_sq.sqrt(), eps)
+        normed_grad = grad.div(denom)
+        if clip_lambda is not None:
+            clip = clip_lambda(step)
+            normed_grad.clamp_(-clip, clip)
+
+        exp_avg.lerp_(normed_grad, 1 - beta1)
+
+        param.add_(exp_avg, alpha=-lr)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
+
+        # update step
+        step_t += 1
+
+
+def _multi_tensor_adopt(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    grad_scale: Optional[Tensor],
+    found_inf: Optional[Tensor],
+    *,
+    has_complex: bool,
+    beta1: float,
+    beta2: float,
+    lr: Union[float, Tensor],
+    clip_lambda: Optional[Callable[[int], float]],
+    weight_decay: float,
+    decouple: bool,
+    eps: float,
+    maximize: bool,
+    capturable: bool,
+    differentiable: bool,
+):
+    if len(params) == 0:
+        return
+
+    if isinstance(lr, Tensor) and not capturable:
+        raise RuntimeError(
+            "lr as a Tensor is not supported for capturable=False and foreach=True"
+        )
+
+    # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable]
+    if not torch._utils.is_compiling() and capturable:
+        capturable_supported_devices = _get_capturable_supported_devices(
+            supports_xla=False
+        )
+        assert all(
+            p.device.type == step.device.type
+            and p.device.type in capturable_supported_devices
+            for p, step in zip(params, state_steps)
+        ), f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}."
+
+    assert grad_scale is None and found_inf is None
+
+    assert not differentiable, "_foreach ops don't support autograd"
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]  # type: ignore[list-item]
+    )
+    for (
+        device_params_,
+        device_grads_,
+        device_exp_avgs_,
+        device_exp_avg_sqs_,
+        device_state_steps_,
+    ), _ in grouped_tensors.values():
+        device_params = cast(List[Tensor], device_params_)
+        device_grads = cast(List[Tensor], device_grads_)
+        device_exp_avgs = cast(List[Tensor], device_exp_avgs_)
+        device_exp_avg_sqs = cast(List[Tensor], device_exp_avg_sqs_)
+        device_state_steps = cast(List[Tensor], device_state_steps_)
+
+        # Handle complex parameters
+        if has_complex:
+            _view_as_real(
+                device_params, device_grads, device_exp_avgs, device_exp_avg_sqs
+            )
+
+        if maximize:
+            device_grads = torch._foreach_neg(device_grads)  # type: ignore[assignment]
+
+        if weight_decay != 0 and not decouple:
+            # Re-use the intermediate memory (device_grads) already allocated for maximize
+            if maximize:
+                torch._foreach_add_(device_grads, device_params, alpha=weight_decay)
+            else:
+                device_grads = torch._foreach_add(  # type: ignore[assignment]
+                    device_grads, device_params, alpha=weight_decay
+                )
+
+        if device_state_steps[0] == 0:
+            torch._foreach_addcmul_(device_exp_avg_sqs, device_grads, device_grads)
+
+            # Update steps
+            # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over
+            # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just
+            # wrapped it once now. The alpha is required to assure we go to the right overload.
+            if not torch._utils.is_compiling() and device_state_steps[0].is_cpu:
+                torch._foreach_add_(
+                    device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0
+                )
+            else:
+                torch._foreach_add_(device_state_steps, 1)
+
+            continue
+
+        if weight_decay != 0 and decouple:
+            torch._foreach_add_(device_params, device_params, alpha=-lr*weight_decay)
+
+        exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs)
+        torch._foreach_maximum_(exp_avg_sq_sqrt, eps)
+
+        normed_grad = torch._foreach_div(device_grads, exp_avg_sq_sqrt)
+        if clip_lambda is not None:
+            clip = clip_lambda(device_state_steps[0])
+            torch._foreach_maximum_(normed_grad, -clip)
+            torch._foreach_minimum_(normed_grad, clip)
+
+        torch._foreach_lerp_(device_exp_avgs, normed_grad, 1 - beta1)
+
+        torch._foreach_add_(device_params, device_exp_avgs, alpha=-lr)
+        torch._foreach_mul_(device_exp_avg_sqs, beta2)
+        torch._foreach_addcmul_(
+            device_exp_avg_sqs, device_grads, device_grads, value=1 - beta2
+        )
+
+        # Update steps
+        # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over
+        # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just
+        # wrapped it once now. The alpha is required to assure we go to the right overload.
+        if not torch._utils.is_compiling() and device_state_steps[0].is_cpu:
+            torch._foreach_add_(
+                device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0
+            )
+        else:
+            torch._foreach_add_(device_state_steps, 1)
+
+
+@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adopt)
+def adopt(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627
+    # setting this as kwarg for now as functional API is compiled by torch/distributed/optim
+    foreach: Optional[bool] = None,
+    capturable: bool = False,
+    differentiable: bool = False,
+    fused: Optional[bool] = None,
+    grad_scale: Optional[Tensor] = None,
+    found_inf: Optional[Tensor] = None,
+    has_complex: bool = False,
+    *,
+    beta1: float,
+    beta2: float,
+    lr: Union[float, Tensor],
+    clip_lambda: Optional[Callable[[int], float]],
+    weight_decay: float,
+    decouple: bool,
+    eps: float,
+    maximize: bool,
+):
+    r"""Functional API that performs ADOPT algorithm computation.
+
+    """
+    # Respect when the user inputs False/True for foreach or fused. We only want to change
+    # the default when neither have been user-specified. Note that we default to foreach
+    # and pass False to use_fused. This is not a mistake--we want to give the fused impl
+    # bake-in time before making it the default, even if it is typically faster.
+    if fused is None and foreach is None:
+        _, foreach = _default_to_fused_or_foreach(
+            params, differentiable, use_fused=False
+        )
+        # Do not flip on foreach for the unsupported case where lr is a Tensor and capturable=False.
+        if foreach and isinstance(lr, Tensor) and not capturable:
+            foreach = False
+    if fused is None:
+        fused = False
+    if foreach is None:
+        foreach = False
+
+    # this check is slow during compilation, so we skip it
+    # if it's strictly needed we can add this check back in dynamo
+    if not torch._utils.is_compiling() and not all(
+        isinstance(t, torch.Tensor) for t in state_steps
+    ):
+        raise RuntimeError(
+            "API has changed, `state_steps` argument must contain a list of singleton tensors"
+        )
+
+    if foreach and torch.jit.is_scripting():
+        raise RuntimeError("torch.jit.script not supported with foreach optimizers")
+    if fused and torch.jit.is_scripting():
+        raise RuntimeError("torch.jit.script not supported with fused optimizers")
+
+    if fused and not torch.jit.is_scripting():
+        func = _fused_adopt # noqa: F821
+    elif foreach and not torch.jit.is_scripting():
+        func = _multi_tensor_adopt
+    else:
+        func = _single_tensor_adopt
+
+    func(
+        params,
+        grads,
+        exp_avgs,
+        exp_avg_sqs,
+        state_steps,
+        has_complex=has_complex,
+        beta1=beta1,
+        beta2=beta2,
+        lr=lr,
+        clip_lambda=clip_lambda,
+        weight_decay=weight_decay,
+        decouple=decouple,
+        eps=eps,
+        maximize=maximize,
+        capturable=capturable,
+        differentiable=differentiable,
+        grad_scale=grad_scale,
+        found_inf=found_inf,
+    )
\ No newline at end of file
diff --git a/dicee/models/base_model.py b/dicee/models/base_model.py
index c6a51890..aed94331 100644
--- a/dicee/models/base_model.py
+++ b/dicee/models/base_model.py
@@ -4,6 +4,7 @@
 import torch
 from torch import nn
 from torch.nn import functional as F
+from .adopt import ADOPT
 
 class BaseKGELightning(pl.LightningModule):
     def __init__(self, *args, **kwargs):
@@ -22,8 +23,17 @@ def mem_of_model(self) -> Dict:
         return {'EstimatedSizeMB': (num_params + buffer_size) / 1024 ** 2, 'NumParam': num_params}
 
     def training_step(self, batch, batch_idx=None):
-        x_batch, y_batch = batch
-        yhat_batch = self.forward(x_batch)
+        if len(batch)==2:
+            # Default
+            x_batch, y_batch = batch
+            yhat_batch = self.forward(x_batch)
+        elif len(batch)==3:
+            # KvsSample or 1vsSample
+            x_batch, y_select, y_batch = batch
+            yhat_batch = self.forward((x_batch,y_select))
+        else:
+            raise RuntimeError("Invalid batch received.")
+
         loss_batch = self.loss_function(yhat_batch, y_batch)
         self.training_step_outputs.append(loss_batch.item())
         self.log("loss",
@@ -82,9 +92,12 @@ def configure_optimizers(self, parameters=None):
             self.selected_optimizer = torch.optim.SGD(params=parameters, lr=self.learning_rate,
                                                       momentum=0, dampening=0, weight_decay=self.weight_decay,
                                                       nesterov=False)
+
         elif self.optimizer_name == 'Adam':
             self.selected_optimizer = torch.optim.Adam(parameters, lr=self.learning_rate,
                                                        weight_decay=self.weight_decay)
+        elif self.optimizer_name == 'Adopt':
+            self.selected_optimizer = ADOPT(parameters, lr=self.learning_rate)
         elif self.optimizer_name == 'AdamW':
             self.selected_optimizer = torch.optim.AdamW(parameters, lr=self.learning_rate,
                                                        weight_decay=self.weight_decay)
@@ -100,7 +113,8 @@ def configure_optimizers(self, parameters=None):
                                                        lr=self.learning_rate, lambd=0.0001, alpha=0.75,
                                                        weight_decay=self.weight_decay)
         else:
-            raise KeyError()
+            raise KeyError(f"{self.optimizer_name} is not found!")
+        print(self.selected_optimizer)
         return self.selected_optimizer
 
 
@@ -282,11 +296,8 @@ def init_params_with_sanity_checking(self):
             self.normalizer_class = IdentityClass
         else:
             raise NotImplementedError()
-        if self.args.get("optim") in ['AdamW', 'Adam', 'SGD']:
-            self.optimizer_name = self.args['optim']
-        else:
-            print(f'--optim (***{self.args.get("optim")}***) not found')
-            self.optimizer_name = 'Adam'
+
+        self.optimizer_name = self.args.get('optim',None)
 
         if self.args.get("init_param") is None:
             self.param_init = IdentityClass
diff --git a/dicee/models/clifford.py b/dicee/models/clifford.py
index 570090f7..2657bd88 100644
--- a/dicee/models/clifford.py
+++ b/dicee/models/clifford.py
@@ -1,230 +1,5 @@
 from .base_model import BaseKGE
 import torch
-from torch.nn import functional as F
-from torch import nn
-
-
-class CMult(BaseKGE):
-    """
-    Cl_(0,0) => Real Numbers
-
-
-    Cl_(0,1) =>
-                A multivector \mathbf{a} = a_0 + a_1 e_1
-                A multivector \mathbf{b} = b_0 + b_1 e_1
-
-                multiplication is isomorphic to the product of two complex numbers
-
-                \mathbf{a} \times \mathbf{b} = a_0 b_0 + a_0b_1 e1 + a_1 b_1 e_1 e_1
-                                             = (a_0 b_0 - a_1 b_1) + (a_0 b_1 + a_1 b_0) e_1
-    Cl_(2,0) =>
-                A multivector \mathbf{a} = a_0 + a_1 e_1 + a_2 e_2 + a_{12} e_1 e_2
-                A multivector \mathbf{b} = b_0 + b_1 e_1 + b_2 e_2 + b_{12} e_1 e_2
-
-                \mathbf{a} \times \mathbf{b} = a_0b_0 + a_0b_1 e_1 + a_0b_2e_2 + a_0 b_12 e_1 e_2
-                                            + a_1 b_0 e_1 + a_1b_1 e_1_e1 ..
-
-    Cl_(0,2) => Quaternions
-    """
-
-    def __init__(self, args):
-        super().__init__(args)
-        self.name = 'CMult'
-        self.entity_embeddings = torch.nn.Embedding(self.num_entities, self.embedding_dim)
-        self.relation_embeddings = torch.nn.Embedding(self.num_relations, self.embedding_dim)
-        self.param_init(self.entity_embeddings.weight.data), self.param_init(self.relation_embeddings.weight.data)
-        self.p = self.args['p']
-        self.q = self.args['q']
-        if self.p is None:
-            self.p = 0
-        if self.q is None:
-            self.q = 0
-        print(f'\tp:{self.p}\tq:{self.q}')
-
-    def clifford_mul(self, x: torch.FloatTensor, y: torch.FloatTensor, p: int, q: int) -> tuple:
-        """
-        Clifford multiplication Cl_{p,q} (\mathbb{R})
-
-        ei ^2 = +1     for i =< i =< p
-        ej ^2 = -1     for p < j =< p+q
-        ei ej = -eje1  for i \neq j
-
-        Parameter
-        ---------
-        x: torch.FloatTensor with (n,d) shape
-
-        y: torch.FloatTensor with (n,d) shape
-
-        p: a non-negative integer p>= 0
-        q: a non-negative integer q>= 0
-
-
-
-        Returns
-        -------
-
-        """
-
-        if p == q == 0:
-            return x * y
-        elif (p == 1 and q == 0) or (p == 0 and q == 1):
-            # {1,e1} e_i^2 = +1 for i
-            a0, a1 = torch.hsplit(x, 2)
-            b0, b1 = torch.hsplit(y, 2)
-            if p == 1 and q == 0:
-                ab0 = a0 * b0 + a1 * b1
-                ab1 = a0 * b1 + a1 * b0
-            else:
-                ab0 = a0 * b0 - a1 * b1
-                ab1 = a0 * b1 + a1 * b0
-            return ab0, ab1
-        elif (p == 2 and q == 0) or (p == 0 and q == 2):
-            a0, a1, a2, a12 = torch.hsplit(x, 4)
-            b0, b1, b2, b12 = torch.hsplit(y, 4)
-            if p == 2 and q == 0:
-                ab0 = a0 * b0 + a1 * b1 + a2 * b2 - a12 * b12
-                ab1 = a0 * b1 + a1 * b0 - a2 * b12 + a12 * b2
-                ab2 = a0 * b2 + a1 * b12 + a2 * b0 - a12 * b1
-                ab12 = a0 * b12 + a1 * b2 - a2 * b1 + a12 * b0
-            else:
-                ab0 = a0 * b0 - a1 * b1 - a2 * b2 - a12 * b12
-                ab1 = a0 * b1 + a1 * b0 + a2 * b12 - a12 * b2
-                ab2 = a0 * b2 - a1 * b12 + a2 * b0 + a12 * b1
-                ab12 = a0 * b12 + a1 * b2 - a2 * b1 + a12 * b0
-            return ab0, ab1, ab2, ab12
-        elif p == 1 and q == 1:
-            a0, a1, a2, a12 = torch.hsplit(x, 4)
-            b0, b1, b2, b12 = torch.hsplit(y, 4)
-
-            ab0 = a0 * b0 + a1 * b1 - a2 * b2 + a12 * b12
-            ab1 = a0 * b1 + a1 * b0 + a2 * b12 - a12 * b2
-            ab2 = a0 * b2 + a1 * b12 + a2 * b0 - a12 * b1
-            ab12 = a0 * b12 + a1 * b2 - a2 * b1 + a12 * b0
-            return ab0, ab1, ab2, ab12
-        elif p == 3 and q == 0:
-            # cl3,0 no 0,3
-            a0, a1, a2, a3, a12, a13, a23, a123 = torch.hsplit(x, 8)
-            b0, b1, b2, b3, b12, b13, b23, b123 = torch.hsplit(y, 8)
-
-            ab0 = a0 * b0 + a1 * b1 + a2 * b2 + a3 * b3 - a12 * b12 - a13 * b13 - a23 * b23 - a123 * b123
-            ab1 = a0 * b1 + a1 * b0 - a2 * b12 - a3 * b13 + a12 * b2 + a13 * b3 - a23 * b123 - a123 * b23
-            ab2 = a0 * b2 + a1 * b12 + a2 * b0 - a3 * b23 - a12 * b1 + a13 * b123 + a23 * b3 + a123 * b13
-            ab3 = a0 * b3 + a1 * b13 + a2 * b23 + a3 * b0 - a12 * b123 - a13 * b1 - a23 * b2 - a123 * b12
-            ab12 = a0 * b12 + a1 * b2 - a2 * b1 + a3 * b123 + a12 * b0 - a13 * b23 + a23 * b13 + a123 * b3
-            ab13 = a0 * b13 + a1 * b3 - a2 * b123 - a3 * b1 + a12 * b23 + a13 * b0 - a23 * b12 - a123 * b2
-            ab23 = a0 * b23 + a1 * b123 + a2 * b3 - a3 * b2 - a12 * b13 - a13 * b12 + a23 * b0 + a123 * b1
-            ab123 = a0 * b123 + a1 * b23 - a2 * b13 + a3 * b12 + a12 * b3 - a13 * b2 + a23 * b1 + a123 * b0
-            return ab0, ab1, ab2, ab3, ab12, ab13, ab23, ab123
-        else:
-            raise NotImplementedError
-
-    def score(self, head_ent_emb, rel_ent_emb, tail_ent_emb):
-        ab = self.clifford_mul(x=head_ent_emb, y=rel_ent_emb, p=self.p, q=self.q)
-
-        if self.p == self.q == 0:
-            return torch.einsum('bd,bd->b', ab, tail_ent_emb)
-        elif (self.p == 1 and self.q == 0) or (self.p == 0 and self.q == 1):
-            ab0, ab1 = ab
-            c0, c1 = torch.hsplit(tail_ent_emb, 2)
-            return torch.einsum('bd,bd->b', ab0, c0) + torch.einsum('bd,bd->b', ab1, c1)
-        elif (self.p == 2 and self.q == 0) or (self.p == 0 and self.q == 2):
-            ab0, ab1, ab2, ab12 = ab
-            c0, c1, c2, c12 = torch.hsplit(tail_ent_emb, 4)
-            return torch.einsum('bd,bd->b', ab0, c0) \
-                + torch.einsum('bd,bd->b', ab1, c1) \
-                + torch.einsum('bd,bd->b', ab2, c2) \
-                + torch.einsum('bd,bd->b', ab12, c12)
-        else:
-            raise NotImplementedError
-
-    def forward_triples(self, x: torch.LongTensor) -> torch.FloatTensor:
-        """
-        Compute batch triple scores
-
-        Parameter
-        ---------
-        x: torch.LongTensor with shape n by 3
-
-
-        Returns
-        -------
-        torch.LongTensor with shape n
-
-        """
-
-        # (1) Retrieve real-valued embedding vectors.
-        head_ent_emb, rel_ent_emb, tail_ent_emb = self.get_triple_representation(x)
-        ab = self.clifford_mul(x=head_ent_emb, y=rel_ent_emb, p=self.p, q=self.q)
-
-        if self.p == self.q == 0:
-            return torch.einsum('bd,bd->b', ab, tail_ent_emb)
-        elif (self.p == 1 and self.q == 0) or (self.p == 0 and self.q == 1):
-            ab0, ab1 = ab
-            c0, c1 = torch.hsplit(tail_ent_emb, 2)
-            return torch.einsum('bd,bd->b', ab0, c0) + torch.einsum('bd,bd->b', ab1, c1)
-        elif (self.p == 2 and self.q == 0) or (self.p == 0 and self.q == 2):
-            ab0, ab1, ab2, ab12 = ab
-            c0, c1, c2, c12 = torch.hsplit(tail_ent_emb, 4)
-            return torch.einsum('bd,bd->b', ab0, c0) \
-                + torch.einsum('bd,bd->b', ab1, c1) \
-                + torch.einsum('bd,bd->b', ab2, c2) \
-                + torch.einsum('bd,bd->b', ab12, c12)
-        else:
-            raise NotImplementedError
-
-    def forward_k_vs_all(self, x: torch.Tensor) -> torch.FloatTensor:
-        """
-        Compute batch KvsAll triple scores
-
-        Parameter
-        ---------
-        x: torch.LongTensor with shape n by 3
-
-
-        Returns
-        -------
-        torch.LongTensor with shape n
-
-        """
-        # (1) Retrieve embedding vectors of heads and relations.
-        head_ent_emb, rel_ent_emb = self.get_head_relation_representation(x)
-        # (2) CL multiply (1).
-        ab = self.clifford_mul(x=head_ent_emb, y=rel_ent_emb, p=self.p, q=self.q)
-        # (3) Inner product of (2) and all entity embeddings.
-        if self.p == self.q == 0:
-            return torch.mm(ab, self.entity_embeddings.weight.transpose(1, 0))
-        elif (self.p == 1 and self.q == 0) or (self.p == 0 and self.q == 1):
-            ab0, ab1 = ab
-            c0, c1 = torch.hsplit(self.entity_embeddings.weight, 2)
-            return torch.mm(ab0, c0.transpose(1, 0)) + torch.mm(ab1, c1.transpose(1, 0))
-        elif (self.p == 2 and self.q == 0) or (self.p == 0 and self.q == 2):
-            ab0, ab1, ab2, ab12 = ab
-            c0, c1, c2, c12 = torch.hsplit(self.entity_embeddings.weight, 4)
-            return torch.mm(ab0, c0.transpose(1, 0)) + \
-                torch.mm(ab1, c1.transpose(1, 0)) + torch.mm(ab2, c2.transpose(1, 0)) + torch.mm(
-                    ab12, c12.transpose(1, 0))
-        elif self.p == 3 and self.q == 0:
-
-            ab0, ab1, ab2, ab3, ab12, ab13, ab23, ab123 = ab
-            c0, c1, c2, c3, c12, c13, c23, c123 = torch.hsplit(self.entity_embeddings.weight, 8)
-
-            return torch.mm(ab0, c0.transpose(1, 0)) \
-                + torch.mm(ab1, c1.transpose(1, 0)) \
-                + torch.mm(ab2, c2.transpose(1, 0)) \
-                + torch.mm(ab3, c3.transpose(1, 0)) + \
-                torch.mm(ab12, c3.transpose(1, 0)) + torch.mm(ab13, c13.transpose(1, 0)) \
-                + torch.mm(ab23, c23.transpose(1, 0)) + torch.mm(ab123, c123.transpose(1, 0))
-        elif self.p == 1 and self.q == 1:
-            ab0, ab1, ab2, ab12 = ab
-            c0, c1, c2, c12 = torch.hsplit(self.entity_embeddings.weight, 4)
-            return torch.mm(ab0, c0.transpose(1, 0)) + torch.mm(ab1, c1.transpose(1, 0)) + \
-                torch.mm(ab2, c2.transpose(1, 0)) + torch.mm(
-                    ab12, c12.transpose(1, 0))
-
-        else:
-            raise NotImplementedError
-
-
 class Keci(BaseKGE):
     def __init__(self, args):
         super().__init__(args)
@@ -244,6 +19,7 @@ def __init__(self, args):
         self.r = int(self.r)
         self.requires_grad_for_interactions = True
         # Initialize parameters for dimension scaling
+        # TODO:Do we need coefficients for the real part ?
         if self.p > 0:
             self.p_coefficients = torch.nn.Embedding(num_embeddings=1, embedding_dim=self.p)
             torch.nn.init.zeros_(self.p_coefficients.weight)
@@ -326,7 +102,7 @@ def compute_sigma_pq(self, *, hp, hq, rp, rq):
         assert sigma_pq.shape[1:] == (self.r, self.p, self.q)
         return sigma_pq
 
-    def apply_coefficients(self, h0, hp, hq, r0, rp, rq):
+    def apply_coefficients(self, hp, hq, rp, rq):
         """ Multiplying a base vector with its scalar coefficient """
         if self.p > 0:
             hp = hp * self.p_coefficients.weight
@@ -334,7 +110,7 @@ def apply_coefficients(self, h0, hp, hq, r0, rp, rq):
         if self.q > 0:
             hq = hq * self.q_coefficients.weight
             rq = rq * self.q_coefficients.weight
-        return h0, hp, hq, r0, rp, rq
+        return hp, hq, rp, rq
 
     def clifford_multiplication(self, h0, hp, hq, r0, rp, rq):
         """ Compute our CL multiplication
@@ -386,6 +162,7 @@ def clifford_multiplication(self, h0, hp, hq, r0, rp, rq):
     def construct_cl_multivector(self, x: torch.FloatTensor, r: int, p: int, q: int) -> tuple[
         torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
         """
+
         Construct a batch of multivectors Cl_{p,q}(\mathbb{R}^d)
 
         Parameter
@@ -397,6 +174,7 @@ def construct_cl_multivector(self, x: torch.FloatTensor, r: int, p: int, q: int)
         a0: torch.FloatTensor with (n,r) shape
         ap: torch.FloatTensor with (n,r,p) shape
         aq: torch.FloatTensor with (n,r,q) shape
+
         """
         batch_size, d = x.shape
         # (1) A_{n \times k}: take the first k columns
@@ -482,7 +260,7 @@ def k_vs_all_score(self, bpe_head_ent_emb, bpe_rel_ent_emb, E):
         h0, hp, hq = self.construct_cl_multivector(bpe_head_ent_emb, r=self.r, p=self.p, q=self.q)
         r0, rp, rq = self.construct_cl_multivector(bpe_rel_ent_emb, r=self.r, p=self.p, q=self.q)
 
-        h0, hp, hq, h0, rp, rq = self.apply_coefficients(h0, hp, hq, h0, rp, rq)
+        hp, hq, rp, rq = self.apply_coefficients(hp, hq, rp, rq)
         # (3.1) Extract real part
         t0 = E[:, :self.r]
 
@@ -550,61 +328,122 @@ def forward_k_vs_all(self, x: torch.Tensor) -> torch.FloatTensor:
         E = self.entity_embeddings.weight
         return self.k_vs_all_score(head_ent_emb, rel_ent_emb, E)
 
-    def forward_k_vs_sample(self, x: torch.LongTensor, target_entity_idx: torch.LongTensor) -> torch.FloatTensor:
+    def construct_batch_selected_cl_multivector(self, x: torch.FloatTensor, r: int, p: int, q: int) -> tuple[
+        torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
         """
-        Kvsall training
+        Construct a batch of batchs multivectors Cl_{p,q}(\mathbb{R}^d)
 
-        (1) Retrieve real-valued embedding vectors for heads and relations \mathbb{R}^d .
-        (2) Construct head entity and relation embeddings according to Cl_{p,q}(\mathbb{R}^d) .
-        (3) Perform Cl multiplication
-        (4) Inner product of (3) and all entity embeddings
+        Parameter
+        ---------
+        x: torch.FloatTensor with (n,k, d) shape
+
+        Returns
+        -------
+        a0: torch.FloatTensor with (n,k, m) shape
+        ap: torch.FloatTensor with (n,k, m, p) shape
+        aq: torch.FloatTensor with (n,k, m, q) shape
+        """
+        batch_size, k, d = x.shape
+
+        # (1) Take the first m columns for each k
+        a0 = x[:, :, :r].view(batch_size, k, r)
 
+        # (2) B_{n \times p}, C_{n \times q}: take the self.k * self.p columns after the k. column
+        if p > 0:
+            ap = x[:, :, r: r + (r * p)].view(batch_size, k, r, p)
+        else:
+            ap = torch.zeros((batch_size, k, r, p), device=self.device)
+        if q > 0:
+            # (3) B_{n \times p}, C_{n \times q}: take the last self.r * self.q .
+            aq = x[:, :, -(r * q):].view(batch_size, k, r, q)
+        else:
+            aq = torch.zeros((batch_size, k, r, q), device=self.device)
+        return a0, ap, aq
+
+    def forward_k_vs_sample(self, x: torch.LongTensor, target_entity_idx: torch.LongTensor) -> torch.FloatTensor:
+        """
         Parameter
         ---------
         x: torch.LongTensor with (n,2) shape
 
+        target_entity_idx: torch.LongTensor with (n, k ) shape k denotes the selected number of examples.
+
         Returns
         -------
-        torch.FloatTensor with (n, |E|) shape
+        torch.FloatTensor with (n, k) shape
         """
         # (1) Retrieve real-valued embedding vectors.
-        head_ent_emb, rel_emb = self.get_head_relation_representation(x)
-        # (2) Construct multi-vector in Cl_{p,q} (\mathbb{R}^d) for head entities.
-        a0, ap, aq = self.construct_cl_multivector(head_ent_emb, r=self.r, p=self.p, q=self.q)
-        # (2) Construct multi-vector in Cl_{p,q} (\mathbb{R}^d) for relations.
-        b0, bp, bq = self.construct_cl_multivector(rel_emb, r=self.r, p=self.p, q=self.q)
-
-        # (4) Clifford multiplication of (2) and (3).
-        # AB_pp, AB_qq, AB_pq
-        # AB_0, AB_p, AB_q, AB_pp, AB_qq, AB_pq = self.clifford_mul_reduced_interactions(a0, ap, aq, b0, bp, bq)
-        AB_0, AB_p, AB_q, AB_pp, AB_qq, AB_pq = self.clifford_mul(a0, ap, aq, b0, bp, bq)
-
-        # b e r
-        selected_tail_entity_embeddings = self.entity_embeddings(target_entity_idx)
-        # (7) Inner product of AB_0 and a0 of all entities.
-        A_score = torch.einsum('br,ber->be', AB_0, selected_tail_entity_embeddings[:, :self.r])
-
-        # (8) Inner product of AB_p and ap of all entities.
+        # (b, d), (b, d)
+        head_ent_emb, rel_ent_emb = self.get_head_relation_representation(x)
+        # (2) Construct multi-vector embeddings in Cl_{p,q} (\mathbb{R}^d) for head entities and relations
+        # (b, m), (b, m, p), (b, m, q)
+        h0, hp, hq = self.construct_cl_multivector(head_ent_emb, r=self.r, p=self.p, q=self.q)
+        # (b, m), (b, m, p), (b, m, q)
+        r0, rp, rq = self.construct_cl_multivector(rel_ent_emb, r=self.r, p=self.p, q=self.q)
+        hp, hq, rp, rq = self.apply_coefficients(hp, hq, rp, rq)
+
+
+        # (3) (b, k, d) Retrieve real-valued embedding vectors of selected entities.
+        E = self.entity_embeddings(target_entity_idx)
+        # (4) Construct multi-vector embeddings in Cl_{p,q} (\mathbb{R}^d) for head entities and relations
+        # (b, k, m), (b, k, m, p), (b, k, m, q)
+        t0, tp, tq = self.construct_batch_selected_cl_multivector(E, r=self.r, p=self.p, q=self.q)
+
+        # (4) Batch vector matrix multiplications
+        # Equivalent computations
+        #                           h0*r0@t0.transpose(1,2)
+        #                           torch.einsum('bm, bmk -> bk', h0 * r0, t0.transpose(1, 2))
+        #                           torch.einsum('bm, bkm -> bk', h0 * r0, t0)
+        h0r0t0 = torch.einsum('bm, bkm -> bk', h0 * r0, t0)
+
+        # (5) Compute a triple score based on interactions described by the bases of p {e_1, ..., e_p}. Eq. 21
         if self.p > 0:
-            B_score = torch.einsum('brp,berp->be', AB_p,
-                                   selected_tail_entity_embeddings[:, self.r: self.r + (self.r * self.p)]
-                                   .view(self.num_entities, self.r, self.p))
+            raise NotImplementedError("Sample with p>0 for Keci not implemented")
+            """
+            # Second term in Eq.16
+            hp_rp_t0 = torch.einsum('brp, br  -> b', hp * rp, t0)
+            # Eq. 17
+            # b=e
+            h0_rp_tp = torch.einsum('brp, erp -> b', torch.einsum('br,  brp -> brp', h0, rp), tp)
+            hp_r0_tp = torch.einsum('brp, erp -> b', torch.einsum('brp, br  -> brp', hp, r0), tp)
+            score_p = hp_rp_t0 + h0_rp_tp + hp_r0_tp
+            """
         else:
-            B_score = 0
-        # (9) Inner product of AB_q and aq of all entities.
+            score_p = 0
+
+        # (6) Compute a triple score based on interactions described by the bases of q {e_{p+1}, ..., e_{p+q}}. Eq. 22
         if self.q > 0:
-            C_score = torch.einsum('brq,berq->be', AB_q,
-                                   selected_tail_entity_embeddings[:, -(self.r * self.q):]
-                                   .view(self.num_entities, self.r, self.q))
+            # \sum_{j=p+1}^{p+q} (h_j r_j t_0) : Third parth of the in Eq 16.
+            # Equivalent computation
+            # torch.einsum('bmq, bkm -> bk', hq*rq, t0) => (hq * rq).transpose(1,2) @ t0.transpose(1,2)
+            hq_rq_t0 = torch.einsum('bmq, bkm -> bk', hq * rq, t0)
+
+            # Eq. 18. Batch elementwise matrix matrix multiplication: bmq -> bkmq
+            rq_tq=torch.einsum('bmq, bkmq -> bkmq', rq, tq)
+            h0_rq_tq = torch.einsum('bm, bkmq  -> bk', h0, rq_tq)
+            hq_tq=torch.einsum('bmq, bkmq -> bkmq',hq, tq)
+            r0_hq_tq = torch.einsum('bm, bkmq  -> bk', r0, hq_tq)
+            score_q = - hq_rq_t0 + (h0_rq_tq + r0_hq_tq)
+        else:
+            score_q = 0
+
+        if self.p >= 2:
+            sigma_pp = torch.sum(self.compute_sigma_pp(hp, rp), dim=[1, 2]).unsqueeze(-1)
         else:
-            C_score = 0
-        # (10) Aggregate (7,8,9).
-        A_B_C_score = A_score + B_score + C_score
-        # (11) Compute inner products of AB_pp, AB_qq, AB_pq and respective identity matrices of all entities.
-        D_E_F_score = (torch.einsum('brpp->b', AB_pp) + torch.einsum('brqq->b', AB_qq) + torch.einsum('brpq->b', AB_pq))
-        D_E_F_score = D_E_F_score.view(len(head_ent_emb), 1)
-        # (12) Score
-        return A_B_C_score + D_E_F_score
+            sigma_pp = 0
+
+        if self.q >= 2:
+            sigma_qq = torch.sum(self.compute_sigma_qq(hq, rq), dim=[1, 2]).unsqueeze(-1)
+        else:
+            sigma_qq = 0
+
+        if self.p >= 2 and self.q >= 2:
+            sigma_pq = torch.sum(self.compute_sigma_pq(hp=hp, hq=hq, rp=rp, rq=rq), dim=[1, 2, 3]).unsqueeze(-1)
+        else:
+            sigma_pq = 0
+        return h0r0t0 + score_p + score_q + sigma_pp + sigma_qq + sigma_pq
+
+
 
     def score(self, h, r, t):
         # (2) Construct multi-vector in Cl_{p,q} (\mathbb{R}^d) for head entities and relations
@@ -615,7 +454,7 @@ def score(self, h, r, t):
         if self.q > 0:
             self.q_coefficients = self.q_coefficients.to(h0.device, non_blocking=True)
 
-        h0, hp, hq, h0, rp, rq = self.apply_coefficients(h0, hp, hq, h0, rp, rq)
+        hp, hq, rp, rq = self.apply_coefficients(hp, hq, rp, rq)
         # (4) Compute a triple score based on interactions described by the basis 1. Eq. 20
         h0r0t0 = torch.einsum('br, br -> b', h0 * r0, t0)
 
@@ -676,7 +515,7 @@ def forward_triples(self, x: torch.Tensor) -> torch.FloatTensor:
         h0, hp, hq = self.construct_cl_multivector(head_ent_emb, r=self.r, p=self.p, q=self.q)
         r0, rp, rq = self.construct_cl_multivector(rel_ent_emb, r=self.r, p=self.p, q=self.q)
         t0, tp, tq = self.construct_cl_multivector(tail_ent_emb, r=self.r, p=self.p, q=self.q)
-        h0, hp, hq, h0, rp, rq = self.apply_coefficients(h0, hp, hq, h0, rp, rq)
+        hp, hq, rp, rq = self.apply_coefficients( hp, hq, rp, rq)
         # (4) Compute a triple score based on interactions described by the basis 1. Eq. 20
         h0r0t0 = torch.einsum('br, br -> b', h0 * r0, t0)
 
diff --git a/dicee/models/complex.py b/dicee/models/complex.py
index 68960f38..032d0247 100644
--- a/dicee/models/complex.py
+++ b/dicee/models/complex.py
@@ -266,4 +266,23 @@ def k_vs_all_score(emb_h: torch.FloatTensor, emb_r: torch.FloatTensor, emb_E: to
     def forward_k_vs_all(self, x: torch.LongTensor) -> torch.FloatTensor:
         # (1) Retrieve embeddings & Apply Dropout & Normalization.
         head_ent_emb, rel_ent_emb = self.get_head_relation_representation(x)
-        return self.k_vs_all_score(head_ent_emb,rel_ent_emb,self.entity_embeddings.weight)
\ No newline at end of file
+        return self.k_vs_all_score(head_ent_emb,rel_ent_emb,self.entity_embeddings.weight)
+
+    def forward_k_vs_sample(self, x: torch.LongTensor, target_entity_idx: torch.LongTensor):
+        # (b,2d), (b,2d)
+        emb_h, emb_r = self.get_head_relation_representation(x)
+        # (b,k,2d)
+        emb_T = self.entity_embeddings(target_entity_idx)
+        # (b,d), (b,d)
+        emb_head_real, emb_head_imag = torch.hsplit(emb_h, 2)
+        # (b,d), (b,d)
+        emb_rel_real, emb_rel_imag = torch.hsplit(emb_r, 2)
+        # (b,k,d), (b,k,d)
+        emb_tail_real, emb_tail_imag = torch.split(emb_T, self.embedding_dim // 2, dim=-1)
+        # Compute hermitian inner product on embedding vectors.
+        real_real_real = torch.einsum("bd, bkd -> bk",emb_head_real * emb_rel_real, emb_tail_real)
+        real_imag_imag = torch.einsum("bd, bkd -> bk",emb_head_real * emb_rel_imag, emb_tail_imag)
+        imag_real_imag = torch.einsum("bd, bkd -> bk",emb_head_imag * emb_rel_real, emb_tail_imag)
+        imag_imag_real = torch.einsum("bd, bkd -> bk",emb_head_imag * emb_rel_imag, emb_tail_real)
+        return real_real_real + real_imag_imag + imag_real_imag - imag_imag_real
+
diff --git a/dicee/models/ensemble.py b/dicee/models/ensemble.py
new file mode 100644
index 00000000..b99022c0
--- /dev/null
+++ b/dicee/models/ensemble.py
@@ -0,0 +1,131 @@
+import torch
+import copy
+from typing import List
+class EnsembleKGE:
+    def __init__(self, seed_model=None,pretrained_models:List=None):
+
+        if seed_model is not None:
+            self.models = []
+            self.optimizers = []
+            self.loss_history = []
+            for i in range(torch.cuda.device_count()):
+                i_model=copy.deepcopy(seed_model)
+                # TODO: Why we cant send the compile model to cpu ?
+                #i_model = torch.compile(i_model)
+                i_model.to(torch.device(f"cuda:{i}"))
+                self.optimizers.append(i_model.configure_optimizers())
+                self.models.append(i_model)
+        else:
+            assert pretrained_models is not None
+            self.models = pretrained_models
+            self.optimizers = []
+            self.loss_history = []
+
+            for i in range(torch.cuda.device_count()):
+                self.models[i].to(torch.device(f"cuda:{i}"))
+                self.optimizers.append(self.models[i].configure_optimizers())
+            # Maybe use the original model's name ?
+        self.name=self.models[0].name
+        self.train_mode=True
+    def named_children(self):
+        return self.models[0].named_children()
+    @property
+    def example_input_array(self):
+        return self.models[0].example_input_array
+    @property
+    def _trainer(self):
+        return self.models[0]._trainer
+
+    def parameters(self):
+        return [ x  for i in self.models for x in i.parameters()]
+        # return self.models[0].parameters()
+    def modules(self):
+        return [x for i in self.models for x in i.modules()]
+
+    def __iter__(self):
+        return (i for i in self.models)
+
+    def __len__(self):
+        return len(self.models)
+
+    def eval(self):
+        for model in self.models:
+            model.eval()
+        self.train_mode=False
+    def to(self,device):
+        for i in range(len(self.models)):
+            if device == "cpu":
+                self.models[i].cpu()
+            else:
+                raise NotImplementedError
+
+
+    def mem_of_model(self):
+        mem_of_ensemble={'EstimatedSizeMB': 0, 'NumParam': 0}
+        for i in self.models:
+            for k,v in i.mem_of_model().items():
+                mem_of_ensemble[k] += v
+        return mem_of_ensemble
+    def __call__(self,x_batch):
+
+        if self.train_mode is False:
+            yhat=0
+            for gpu_id, model in enumerate(self.models):
+                yhat += model(x_batch)
+            return yhat / len(self.models)
+        else:
+            for opt in self.optimizers:
+                opt.zero_grad()
+            yhat=None
+            for gpu_id, model in enumerate(self.models):
+                # Move batch into the GPU where the i.th model resides
+                if isinstance(x_batch, tuple):
+                    x_batch=(x_batch[0].to(f"cuda:{gpu_id}"),x_batch[1].to(f"cuda:{gpu_id}"))
+                else:
+                    x_batch=x_batch.to(f"cuda:{gpu_id}")
+                if yhat is None:
+                    yhat=model(x_batch)
+                else:
+                    yhat+=model(x_batch).to("cuda:0")
+            return yhat/len(self.models)
+    
+    def step(self):
+        for opt in self.optimizers:
+            opt.step()
+
+    def get_embeddings(self):
+        entity_embeddings=[]
+        relation_embeddings=[]
+        # () Iterate
+        for trained_model in self.models:
+            entity_emb, relation_ebm = trained_model.get_embeddings()
+            entity_embeddings.append(entity_emb)
+            if relation_ebm is not None:
+                relation_embeddings.append(relation_ebm)
+        # () Concat the embedding vectors horizontally.
+        entity_embeddings=torch.cat(entity_embeddings,dim=1)
+        if relation_embeddings:
+            relation_embeddings=torch.cat(relation_embeddings,dim=1)
+        else:
+            relation_embeddings=None
+
+        return entity_embeddings, relation_embeddings
+
+    """
+    def __getattr__(self, name):
+        # Create a function that will call the same attribute/method on each model
+        def method(*args, **kwargs):
+            results = []
+            for model in self.models:
+                attr = getattr(model, name)
+                if callable(attr):
+                    # If it's a method, call it with provided arguments
+                    results.append(attr(*args, **kwargs))
+                else:
+                    # If it's an attribute, just get its value
+                    results.append(attr)
+            return results
+        return method
+    """
+    def __str__(self):
+        return f"EnsembleKGE of {len(self.models)} {self.models[0]}"
diff --git a/dicee/models/function_space.py b/dicee/models/function_space.py
index fc991ac1..57bfa6f1 100644
--- a/dicee/models/function_space.py
+++ b/dicee/models/function_space.py
@@ -18,7 +18,6 @@ def __init__(self, args):
         self.gamma = torch.randn(self.k, self.num_sample)  # N(0,1)
         # Lazy import
         from scipy.special import roots_legendre
-        from scipy.special import roots_legendre
         roots, weights = roots_legendre(self.num_sample)
         self.roots = torch.from_numpy(roots).repeat(self.k, 1).float()  # shape self.k by self.n
         self.weights = torch.from_numpy(weights).reshape(1, -1).float()  # shape 1 by self.n
@@ -81,6 +80,8 @@ def __init__(self, args):
         self.param_init(self.entity_embeddings.weight.data), self.param_init(self.relation_embeddings.weight.data)
         self.k = int(np.sqrt(self.embedding_dim // 2))
         self.num_sample = 250
+        # Lazy import
+        from scipy.special import roots_legendre
         roots, weights = roots_legendre(self.num_sample)
         self.roots = torch.from_numpy(roots).repeat(self.k, 1).float()  # shape self.k by self.n
         self.weights = torch.from_numpy(weights).reshape(1, -1).float()  # shape 1 by self.n
diff --git a/dicee/models/pykeen_models.py b/dicee/models/pykeen_models.py
index e2aeacdb..7c218691 100644
--- a/dicee/models/pykeen_models.py
+++ b/dicee/models/pykeen_models.py
@@ -57,7 +57,7 @@ def __init__(self, args: dict):
         try:
             # lazy import
             from pykeen.models import model_resolver
-        except:
+        except ImportError:
             print(traceback.format_exc())
             print("Pykeen does not work with pytorch>2.0.0. Current pytorch version:",torch.__version__)
             exit(1)
diff --git a/dicee/models/real.py b/dicee/models/real.py
index fccb9c15..8d32ec66 100644
--- a/dicee/models/real.py
+++ b/dicee/models/real.py
@@ -33,10 +33,14 @@ def forward_k_vs_all(self, x: torch.LongTensor):
         return self.k_vs_all_score(emb_h=emb_head, emb_r=emb_rel, emb_E=self.entity_embeddings.weight)
 
     def forward_k_vs_sample(self, x: torch.LongTensor, target_entity_idx: torch.LongTensor):
+        # (b,d),     (b,d)
         emb_head_real, emb_rel_real = self.get_head_relation_representation(x)
-        hr = self.hidden_dropout(self.hidden_normalizer(emb_head_real * emb_rel_real)).unsqueeze(1)
-        t = self.entity_embeddings(target_entity_idx).transpose(1, 2)
-        return torch.bmm(hr, t).squeeze(1)
+        # (b, d)
+        hr = torch.einsum('bd, bd -> bd', emb_head_real, emb_rel_real)
+        # (b, k, d)
+        t = self.entity_embeddings(target_entity_idx)
+        return torch.einsum('bd, bkd -> bk', hr, t)
+
 
     def score(self, h, r, t):
         return (self.hidden_dropout(self.hidden_normalizer(h * r)) * t).sum(dim=1)
diff --git a/dicee/models/transformers.py b/dicee/models/transformers.py
index 518f3161..56832fdc 100644
--- a/dicee/models/transformers.py
+++ b/dicee/models/transformers.py
@@ -1,6 +1,3 @@
-from .base_model import BaseKGE
-import torch
-
 """
 Full definition of a GPT Language Model, all of it in this single file.
 References:
@@ -9,11 +6,10 @@
 2) huggingface/transformers PyTorch implementation:
 https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
 """
-
+from .base_model import BaseKGE
 import math
 import inspect
 from dataclasses import dataclass
-
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
diff --git a/dicee/query_generator.py b/dicee/query_generator.py
index 14597142..f9e1cccd 100644
--- a/dicee/query_generator.py
+++ b/dicee/query_generator.py
@@ -1,6 +1,5 @@
 from collections import defaultdict
 from typing import Union, Dict, List, Tuple
-
 import numpy as np
 import random
 import os
@@ -491,6 +490,8 @@ def save_queries(self, query_type: str, gen_num: int, save_path: str):
         name_to_save = f'{self.mode}-{query_type}'
         if not os.path.isdir(save_path):
             os.makedirs(save_path)
+        # TODO: CD: Deprecate the pickle usage for data serialization.
+        # TODO: CD: yet since this files are small we can get away with them
         with open(f'{save_path}/{name_to_save}-queries.pkl', 'wb') as f:
             pickle.dump(unmapped_queries, f)
         with open(f'{save_path}/{name_to_save}-easy-answers.pkl', 'wb') as f:
diff --git a/dicee/read_preprocess_save_load_kg/preprocess.py b/dicee/read_preprocess_save_load_kg/preprocess.py
index 9bf895b1..fde7473a 100644
--- a/dicee/read_preprocess_save_load_kg/preprocess.py
+++ b/dicee/read_preprocess_save_load_kg/preprocess.py
@@ -1,13 +1,12 @@
 import pandas as pd
 import polars as pl
-from .util import timeit, index_triples_with_pandas, dataset_sanity_checking
+from .util import timeit, pandas_dataframe_indexer, dataset_sanity_checking
 from dicee.static_funcs import numpy_data_type_changer
-from .util import get_er_vocab, get_re_vocab, get_ee_vocab, create_constraints, apply_reciprical_or_noise
+from .util import get_er_vocab, get_re_vocab, get_ee_vocab, apply_reciprical_or_noise, polars_dataframe_indexer
 import numpy as np
 import concurrent
 from typing import List, Tuple
 from typing import Union
-import itertools
 
 
 class PreprocessKG:
@@ -39,6 +38,9 @@ def start(self) -> None:
         else:
             raise KeyError(f'{self.kg.backend} not found')
 
+        # TODO: Still we keep the raw and index data in memory
+        # print(self.kg.raw_train_set)
+        # print(self.kg.train_set)
         if self.kg.eval_model:
             if self.kg.byte_pair_encoding:
                 data = []
@@ -59,14 +61,10 @@ def start(self) -> None:
             self.kg.re_vocab = executor.submit(get_re_vocab, data, self.kg.path_for_serialization + '/re_vocab.p')
             self.kg.ee_vocab = executor.submit(get_ee_vocab, data, self.kg.path_for_serialization + '/ee_vocab.p')
 
-            self.kg.constraints = executor.submit(create_constraints, self.kg.train_set,
-                                                  self.kg.path_for_serialization + '/constraints.p')
-            self.kg.domain_constraints_per_rel, self.kg.range_constraints_per_rel = None, None
-
         # string containing
         assert isinstance(self.kg.raw_train_set, pd.DataFrame) or isinstance(self.kg.raw_train_set, pl.DataFrame)
 
-        print("Creating dataset...")
+        # print("Creating dataset...")
         if self.kg.byte_pair_encoding and self.kg.padding:
             assert isinstance(self.kg.train_set, list)
             assert isinstance(self.kg.train_set[0], tuple)
@@ -149,20 +147,13 @@ def start(self) -> None:
                 # print(self.kg.enc.decode(x))
                 triples.extend(x)
             self.kg.train_set = np.array(triples)
-
         else:
-            """No need to do anything. We create datasets for other models in the pyorch dataset construction"""
-            # @TODO: Either we should move the all pytorch dataset construciton into here
-            # Or we should move the byte pair encoding data into
-            print('Finding suitable integer type for the index...')
-            self.kg.train_set = numpy_data_type_changer(self.kg.train_set,
-                                                        num=max(self.kg.num_entities, self.kg.num_relations))
-            if self.kg.valid_set is not None:
-                self.kg.valid_set = numpy_data_type_changer(self.kg.valid_set,
-                                                            num=max(self.kg.num_entities, self.kg.num_relations))
-            if self.kg.test_set is not None:
-                self.kg.test_set = numpy_data_type_changer(self.kg.test_set,
-                                                           num=max(self.kg.num_entities, self.kg.num_relations))
+            # No need to keep the raw data in memory
+            # TODO: If BPE used, no need to clearn data for the time being,.
+            self.kg.raw_train_set = None
+            self.kg.raw_valid_set = None
+            self.kg.raw_test_set  = None
+
 
     @staticmethod
     def __replace_values_df(df: pd.DataFrame = None, f=None) -> Union[
@@ -230,13 +221,13 @@ def preprocess_with_byte_pair_encoding(self):
         assert isinstance(self.kg.raw_train_set, pd.DataFrame)
         assert self.kg.raw_train_set.columns.tolist() == ['subject', 'relation', 'object']
         # (1)  Add recipriocal or noisy triples into raw_train_set, raw_valid_set, raw_test_set
-        self.kg.raw_train_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprical,
+        self.kg.raw_train_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                           eval_model=self.kg.eval_model,
                                                           df=self.kg.raw_train_set, info="Train")
-        self.kg.raw_valid_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprical,
+        self.kg.raw_valid_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                           eval_model=self.kg.eval_model,
                                                           df=self.kg.raw_valid_set, info="Validation")
-        self.kg.raw_test_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprical,
+        self.kg.raw_test_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                          eval_model=self.kg.eval_model,
                                                          df=self.kg.raw_test_set, info="Test")
         # (2) Transformation from DataFrame to list of tuples.
@@ -304,13 +295,13 @@ def preprocess_with_pandas(self) -> None:
         None
         """
         # (1)  Add recipriocal or noisy triples.
-        self.kg.raw_train_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprical,
+        self.kg.raw_train_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                           eval_model=self.kg.eval_model,
                                                           df=self.kg.raw_train_set, info="Train")
-        self.kg.raw_valid_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprical,
+        self.kg.raw_valid_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                           eval_model=self.kg.eval_model,
                                                           df=self.kg.raw_valid_set, info="Validation")
-        self.kg.raw_test_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprical,
+        self.kg.raw_test_set = apply_reciprical_or_noise(add_reciprical=self.kg.add_reciprocal,
                                                          eval_model=self.kg.eval_model,
                                                          df=self.kg.raw_test_set, info="Test")
 
@@ -319,16 +310,17 @@ def preprocess_with_pandas(self) -> None:
         self.kg.num_entities, self.kg.num_relations = len(self.kg.entity_to_idx), len(self.kg.relation_to_idx)
 
         # (3) Index datasets
-        self.kg.train_set = index_triples_with_pandas(self.kg.raw_train_set,
+        self.kg.train_set = pandas_dataframe_indexer(self.kg.raw_train_set,
                                                       self.kg.entity_to_idx,
                                                       self.kg.relation_to_idx)
         assert isinstance(self.kg.train_set, pd.core.frame.DataFrame)
         self.kg.train_set = self.kg.train_set.values
         self.kg.train_set = numpy_data_type_changer(self.kg.train_set,
                                                     num=max(self.kg.num_entities, self.kg.num_relations))
+
         dataset_sanity_checking(self.kg.train_set, self.kg.num_entities, self.kg.num_relations)
         if self.kg.raw_valid_set is not None:
-            self.kg.valid_set = index_triples_with_pandas(self.kg.raw_valid_set, self.kg.entity_to_idx,
+            self.kg.valid_set = pandas_dataframe_indexer(self.kg.raw_valid_set, self.kg.entity_to_idx,
                                                           self.kg.relation_to_idx)
             self.kg.valid_set = self.kg.valid_set.values
             dataset_sanity_checking(self.kg.valid_set, self.kg.num_entities, self.kg.num_relations)
@@ -336,7 +328,7 @@ def preprocess_with_pandas(self) -> None:
                                                         num=max(self.kg.num_entities, self.kg.num_relations))
 
         if self.kg.raw_test_set is not None:
-            self.kg.test_set = index_triples_with_pandas(self.kg.raw_test_set, self.kg.entity_to_idx,
+            self.kg.test_set = pandas_dataframe_indexer(self.kg.raw_test_set, self.kg.entity_to_idx,
                                                          self.kg.relation_to_idx)
             # To numpy
             self.kg.test_set = self.kg.test_set.values
@@ -355,27 +347,27 @@ def preprocess_with_polars(self) -> None:
         print(f'*** Preprocessing Train Data:{self.kg.raw_train_set.shape} with Polars ***')
 
         # (1) Add reciprocal triples, e.g. KG:= {(s,p,o)} union {(o,p_inverse,s)}
-        if self.kg.add_reciprical and self.kg.eval_model:
+        if self.kg.add_reciprocal and self.kg.eval_model:
             def adding_reciprocal_triples():
                 """ Add reciprocal triples """
                 # (1.1) Add reciprocal triples into training set
                 self.kg.raw_train_set.extend(self.kg.raw_train_set.select([
                     pl.col("object").alias('subject'),
-                    pl.col("relation").apply(lambda x: x + '_inverse'),
+                    pl.col("relation")+'_inverse',
                     pl.col("subject").alias('object')
                 ]))
                 if self.kg.raw_valid_set is not None:
                     # (1.2) Add reciprocal triples into valid_set set.
                     self.kg.raw_valid_set.extend(self.kg.raw_valid_set.select([
                         pl.col("object").alias('subject'),
-                        pl.col("relation").apply(lambda x: x + '_inverse'),
+                        pl.col("relation")+'_inverse',
                         pl.col("subject").alias('object')
                     ]))
                 if self.kg.raw_test_set is not None:
                     # (1.2) Add reciprocal triples into test set.
                     self.kg.raw_test_set.extend(self.kg.raw_test_set.select([
                         pl.col("object").alias('subject'),
-                        pl.col("relation").apply(lambda x: x + '_inverse'),
+                        pl.col("relation")+'_inverse',
                         pl.col("subject").alias('object')
                     ]))
 
@@ -389,7 +381,6 @@ def adding_reciprocal_triples():
             raise TypeError(f"{type(self.kg.raw_train_set)}")
         assert isinstance(self.kg.raw_valid_set, pl.DataFrame) or self.kg.raw_valid_set is None
         assert isinstance(self.kg.raw_test_set, pl.DataFrame) or self.kg.raw_test_set is None
-
         def concat_splits(train, val, test):
             x = [train]
             if val is not None:
@@ -397,38 +388,39 @@ def concat_splits(train, val, test):
             if test is not None:
                 x.append(test)
             return pl.concat(x)
-
         print('Concat Splits...')
         df_str_kg = concat_splits(self.kg.raw_train_set, self.kg.raw_valid_set, self.kg.raw_test_set)
-
-        print('Entity Indexing...')
-        self.kg.entity_to_idx = pl.concat((df_str_kg['subject'],
-                                           df_str_kg['object'])).unique(maintain_order=True).rename('entity')
+        # () Select unique subject entities.
+        print("Collecting subject entities...")
+        subjects = df_str_kg.select(pl.col("subject").unique(maintain_order=True).alias("entity"))
+        print(f"Unique number of subjects:{len(subjects)}")
+        # () Select unique object entities.
+        print("Collecting object entities...")
+        objects = df_str_kg.select(pl.col("object").unique(maintain_order=True).alias("entity"))
+        # () Select unique entities.
+        self.kg.entity_to_idx = pl.concat([subjects, objects], how="vertical").unique(maintain_order=True)
+        self.kg.entity_to_idx = self.kg.entity_to_idx.with_row_index("index").select(["index", "entity"])
+        # () Write unique entities with indices.
         print('Relation Indexing...')
-        self.kg.relation_to_idx = df_str_kg['relation'].unique(maintain_order=True)
-        print('Creating index for entities...')
-        self.kg.entity_to_idx = {ent: idx for idx, ent in enumerate(self.kg.entity_to_idx.to_list())}
-        print('Creating index for relations...')
-        self.kg.relation_to_idx = {rel: idx for idx, rel in enumerate(self.kg.relation_to_idx.to_list())}
-        self.kg.num_entities, self.kg.num_relations = len(self.kg.entity_to_idx), len(self.kg.relation_to_idx)
-
+        self.kg.relation_to_idx = df_str_kg.select(pl.col("relation").unique(maintain_order=True)).with_row_index(
+            "index").select(["index", "relation"])
+        del df_str_kg
         print(f'Indexing Training Data {self.kg.raw_train_set.shape}...')
-        self.kg.train_set = self.kg.raw_train_set.with_columns(
-            pl.col("subject").map_dict(self.kg.entity_to_idx).alias("subject"),
-            pl.col("relation").map_dict(self.kg.relation_to_idx).alias("relation"),
-            pl.col("object").map_dict(self.kg.entity_to_idx).alias("object")).to_numpy()
+        self.kg.train_set=polars_dataframe_indexer(self.kg.raw_train_set, self.kg.entity_to_idx, self.kg.relation_to_idx).to_numpy()
+
         if self.kg.raw_valid_set is not None:
             print(f'Indexing Val Data {self.kg.raw_valid_set.shape}...')
-            self.kg.valid_set = self.kg.raw_valid_set.with_columns(
-                pl.col("subject").map_dict(self.kg.entity_to_idx).alias("subject"),
-                pl.col("relation").map_dict(self.kg.relation_to_idx).alias("relation"),
-                pl.col("object").map_dict(self.kg.entity_to_idx).alias("object")).to_numpy()
+            self.kg.valid_set = polars_dataframe_indexer(self.kg.raw_valid_set, self.kg.entity_to_idx, self.kg.relation_to_idx).to_numpy()
+
         if self.kg.raw_test_set is not None:
             print(f'Indexing Test Data {self.kg.raw_test_set.shape}...')
-            self.kg.test_set = self.kg.raw_test_set.with_columns(
-                pl.col("subject").map_dict(self.kg.entity_to_idx).alias("subject"),
-                pl.col("relation").map_dict(self.kg.relation_to_idx).alias("relation"),
-                pl.col("object").map_dict(self.kg.entity_to_idx).alias("object")).to_numpy()
+            self.kg.test_set = polars_dataframe_indexer(self.kg.raw_test_set, self.kg.entity_to_idx, self.kg.relation_to_idx).to_numpy()
+
+        self.kg.num_entities, self.kg.num_relations = len(self.kg.entity_to_idx), len(self.kg.relation_to_idx)
+        """
+        self.kg.entity_to_idx=dict(zip(self.kg.entity_to_idx["entity"].to_list(), self.kg.entity_to_idx["index"].to_list()))
+        self.kg.relation_to_idx=dict(zip(self.kg.relation_to_idx["relation"].to_list(), self.kg.relation_to_idx["index"].to_list()))
+        """
         print(f'*** Preprocessing Train Data:{self.kg.train_set.shape} with Polars DONE ***')
 
     def sequential_vocabulary_construction(self) -> None:
@@ -439,18 +431,10 @@ def sequential_vocabulary_construction(self) -> None:
                     => the index is integer and
                     => a single column is string (e.g. URI)
         """
-        try:
-            assert isinstance(self.kg.raw_train_set, pd.DataFrame)
-        except AssertionError:
-            raise AssertionError
-            print(type(self.kg.raw_train_set))
-            print('HEREE')
-            exit(1)
+        assert isinstance(self.kg.raw_train_set, pd.DataFrame)
         assert isinstance(self.kg.raw_valid_set, pd.DataFrame) or self.kg.raw_valid_set is None
         assert isinstance(self.kg.raw_test_set, pd.DataFrame) or self.kg.raw_test_set is None
 
-        # (4) Remove triples from (1).
-        self.remove_triples_from_train_with_condition()
         # Concatenate dataframes.
         print('Concatenating data to obtain index...')
         x = [self.kg.raw_train_set]
@@ -465,41 +449,13 @@ def sequential_vocabulary_construction(self) -> None:
         # ravel('K') => Return a contiguous flattened array.
         # ‘K’ means to read the elements in the order they occur in memory,
         # except for reversing the data when strides are negative.
-        ordered_list = pd.unique(df_str_kg[['subject', 'object']].values.ravel('K')).tolist()
-        self.kg.entity_to_idx = {k: i for i, k in enumerate(ordered_list)}
+        # ordered_list = pd.unique(df_str_kg[['subject', 'object']].values.ravel('K')).tolist()
+        # self.kg.entity_to_idx = {k: i for i, k in enumerate(ordered_list)}
+        # Instead of dict, storing it in a pandas dataframe
+        self.kg.entity_to_idx=pd.concat((df_str_kg['subject'],df_str_kg['object'])).to_frame("entity").drop_duplicates(keep="first",ignore_index=True)
         # 5. Create a bijection mapping  from relations to integer indexes.
-        ordered_list = pd.unique(df_str_kg['relation'].values.ravel('K')).tolist()
-        self.kg.relation_to_idx = {k: i for i, k in enumerate(ordered_list)}
-        del ordered_list
-
-    def remove_triples_from_train_with_condition(self):
-        if None:
-            # self.kg.min_freq_for_vocab is not
-            assert isinstance(self.kg.min_freq_for_vocab, int)
-            assert self.kg.min_freq_for_vocab > 0
-            print(
-                f'[5 / 14] Dropping triples having infrequent entities or relations (>{self.kg.min_freq_for_vocab})...',
-                end=' ')
-            num_triples = self.kg.raw_train_set.size
-            print('Total num triples:', num_triples, end=' ')
-            # Compute entity frequency: index is URI, val is number of occurrences.
-            entity_frequency = pd.concat(
-                [self.kg.raw_train_set['subject'], self.kg.raw_train_set['object']]).value_counts()
-            relation_frequency = self.kg.raw_train_set['relation'].value_counts()
-
-            # low_frequency_entities index and values are the same URIs: dask.dataframe.core.DataFrame
-            low_frequency_entities = entity_frequency[
-                entity_frequency <= self.kg.min_freq_for_vocab].index.values
-            low_frequency_relation = relation_frequency[
-                relation_frequency <= self.kg.min_freq_for_vocab].index.values
-            # If triple contains subject that is in low_freq, set False do not select
-            self.kg.raw_train_set = self.kg.raw_train_set[
-                ~self.kg.raw_train_set['subject'].isin(low_frequency_entities)]
-            # If triple contains object that is in low_freq, set False do not select
-            self.kg.raw_train_set = self.kg.raw_train_set[~self.kg.raw_train_set['object'].isin(low_frequency_entities)]
-            # If triple contains relation that is in low_freq, set False do not select
-            self.kg.raw_train_set = self.kg.raw_train_set[
-                ~self.kg.raw_train_set['relation'].isin(low_frequency_relation)]
-            # print('\t after dropping:', df_str_kg.size.compute(scheduler=scheduler_flag))
-            print('\t after dropping:', self.kg.raw_train_set.size)  # .compute(scheduler=scheduler_flag))
-            del low_frequency_entities
+        # ordered_list = pd.unique(df_str_kg['relation'].values.ravel('K')).tolist()
+        # self.kg.relation_to_idx = {k: i for i, k in enumerate(ordered_list)}
+        self.kg.relation_to_idx = df_str_kg['relation'].to_frame("relation").drop_duplicates(keep="first", ignore_index=True)
+
+        # del ordered_list
\ No newline at end of file
diff --git a/dicee/read_preprocess_save_load_kg/read_from_disk.py b/dicee/read_preprocess_save_load_kg/read_from_disk.py
index c9a43e97..5099d43c 100644
--- a/dicee/read_preprocess_save_load_kg/read_from_disk.py
+++ b/dicee/read_preprocess_save_load_kg/read_from_disk.py
@@ -24,14 +24,14 @@ def start(self) -> None:
         -------
         None
         """
+
         if self.kg.path_single_kg is not None:
             self.kg.raw_train_set = read_from_disk(self.kg.path_single_kg,
                                                    self.kg.read_only_few,
                                                    self.kg.sample_triples_ratio,
-                                                   backend=self.kg.backend)
+                                                   backend=self.kg.backend,separator=self.kg.separator)
             if self.kg.add_noise_rate:
                 self.add_noisy_triples_into_training()
-
             self.kg.raw_valid_set = None
             self.kg.raw_test_set = None
         elif self.kg.sparql_endpoint is not None:
@@ -42,13 +42,13 @@ def start(self) -> None:
             for i in glob.glob(self.kg.dataset_dir + '/*'):
                 if 'train' in i:
                     self.kg.raw_train_set = read_from_disk(i, self.kg.read_only_few, self.kg.sample_triples_ratio,
-                                                           backend=self.kg.backend)
+                                                       backend=self.kg.backend, separator=self.kg.separator)
                     if self.kg.add_noise_rate:
                         self.add_noisy_triples_into_training()
                 elif 'test' in i and self.kg.eval_model is not None:
-                    self.kg.raw_test_set = read_from_disk(i, backend=self.kg.backend)
+                    self.kg.raw_test_set = read_from_disk(i, backend=self.kg.backend, separator=self.kg.separator)
                 elif 'valid' in i and self.kg.eval_model is not None:
-                    self.kg.raw_valid_set = read_from_disk(i, backend=self.kg.backend)
+                    self.kg.raw_valid_set = read_from_disk(i, backend=self.kg.backend, separator=self.kg.separator)
                 else:
                     print(f'Not processed data: {i}')
         else:
diff --git a/dicee/read_preprocess_save_load_kg/save_load_disk.py b/dicee/read_preprocess_save_load_kg/save_load_disk.py
index c6fc6e54..b527b721 100644
--- a/dicee/read_preprocess_save_load_kg/save_load_disk.py
+++ b/dicee/read_preprocess_save_load_kg/save_load_disk.py
@@ -1,4 +1,6 @@
 import numpy as np
+import polars
+import pandas
 from .util import load_pickle, load_numpy_ndarray
 import os
 from dicee.static_funcs import save_pickle, save_numpy_ndarray
@@ -21,13 +23,19 @@ def save(self):
             save_pickle(data=self.kg.ordered_bpe_entities, file_path=self.kg.path_for_serialization + '/ordered_bpe_entities.p')
             save_pickle(data=self.kg.ordered_bpe_relations, file_path=self.kg.path_for_serialization + '/ordered_bpe_relations.p')
         else:
-            assert isinstance(self.kg.entity_to_idx, dict)
-            assert isinstance(self.kg.relation_to_idx, dict)
             assert isinstance(self.kg.train_set, np.ndarray)
-
             # (1) Save dictionary mappings into disk
-            save_pickle(data=self.kg.entity_to_idx, file_path=self.kg.path_for_serialization + '/entity_to_idx.p')
-            save_pickle(data=self.kg.relation_to_idx, file_path=self.kg.path_for_serialization + '/relation_to_idx.p')
+            if isinstance(self.kg.entity_to_idx, dict):
+                save_pickle(data=self.kg.entity_to_idx, file_path=self.kg.path_for_serialization + '/entity_to_idx.p')
+                save_pickle(data=self.kg.relation_to_idx, file_path=self.kg.path_for_serialization + '/relation_to_idx.p')
+            elif isinstance(self.kg.entity_to_idx, polars.DataFrame):
+                self.kg.entity_to_idx.write_csv(file=self.kg.path_for_serialization + "/entity_to_idx.csv", include_header=True)
+                self.kg.relation_to_idx.write_csv(file=self.kg.path_for_serialization + "/relation_to_idx.csv", include_header=True)
+            elif isinstance(self.kg.entity_to_idx, pandas.DataFrame):
+                self.kg.entity_to_idx.to_csv(path_or_buf=self.kg.path_for_serialization + "/entity_to_idx.csv", header=True)
+                self.kg.relation_to_idx.to_csv(path_or_buf=self.kg.path_for_serialization + "/relation_to_idx.csv", header=True)
+            else:
+                raise RuntimeError("Unexpected type for entity_to_idx or relation_to_idx")
 
             save_numpy_ndarray(data=self.kg.train_set, file_path=self.kg.path_for_serialization + '/train_set.npy')
             if self.kg.valid_set is not None:
diff --git a/dicee/read_preprocess_save_load_kg/util.py b/dicee/read_preprocess_save_load_kg/util.py
index 03665102..1ce8a223 100644
--- a/dicee/read_preprocess_save_load_kg/util.py
+++ b/dicee/read_preprocess_save_load_kg/util.py
@@ -9,6 +9,115 @@
 import os
 import psutil
 import requests
+from typing import Tuple
+
+def polars_dataframe_indexer(df_polars:polars.DataFrame, idx_entity:polars.DataFrame, idx_relation:polars.DataFrame)->polars.DataFrame:
+    """
+     Replaces 'subject', 'relation', and 'object' columns in the input Polars DataFrame with their corresponding index values
+     from the entity and relation index DataFrames.
+
+     This function processes the DataFrame in three main steps:
+     1. Replace the 'relation' values with the corresponding index from `idx_relation`.
+     2. Replace the 'subject' values with the corresponding index from `idx_entity`.
+     3. Replace the 'object' values with the corresponding index from `idx_entity`.
+
+     Parameters:
+     -----------
+     df_polars : polars.DataFrame
+         The input Polars DataFrame containing columns: 'subject', 'relation', and 'object'.
+
+     idx_entity : polars.DataFrame
+         A Polars DataFrame that contains the mapping between entity names and their corresponding indices.
+         Must have columns: 'entity' and 'index'.
+
+     idx_relation : polars.DataFrame
+         A Polars DataFrame that contains the mapping between relation names and their corresponding indices.
+         Must have columns: 'relation' and 'index'.
+
+     Returns:
+     --------
+     polars.DataFrame
+         A DataFrame with the 'subject', 'relation', and 'object' columns replaced by their corresponding indices.
+
+     Example Usage:
+     --------------
+     >>> df_polars = pl.DataFrame({
+             "subject": ["Alice", "Bob", "Charlie"],
+             "relation": ["knows", "works_with", "lives_in"],
+             "object": ["Dave", "Eve", "Frank"]
+         })
+     >>> idx_entity = pl.DataFrame({
+             "entity": ["Alice", "Bob", "Charlie", "Dave", "Eve", "Frank"],
+             "index": [0, 1, 2, 3, 4, 5]
+         })
+     >>> idx_relation = pl.DataFrame({
+             "relation": ["knows", "works_with", "lives_in"],
+             "index": [0, 1, 2]
+         })
+     >>> polars_dataframe_indexer(df_polars, idx_entity, idx_relation)
+
+     Steps:
+     ------
+     1. Join the input DataFrame `df_polars` on the 'relation' column with `idx_relation` to replace the relations with their indices.
+     2. Join on 'subject' to replace it with the corresponding entity index using a left join on `idx_entity`.
+     3. Join on 'object' to replace it with the corresponding entity index using a left join on `idx_entity`.
+     4. Select only the 'subject', 'relation', and 'object' columns to return the final result.
+     """
+    assert isinstance(df_polars, polars.DataFrame)
+    assert isinstance(idx_entity, polars.DataFrame)
+    assert isinstance(idx_relation, polars.DataFrame)
+
+    # Step : Join on 'relation' to replace relation with its index
+    df_merged = df_polars.join(idx_relation, on="relation", how="left")
+    df_merged = df_merged.select([polars.col("subject"), polars.col("index").alias("relation"), polars.col("object")])
+    # Step :  Consider Left Table on subject and Right Table on entity with the left join
+    # Returns all rows from the left table, and the matched rows from the right table
+    df_merged = df_merged.join(idx_entity, left_on="subject", right_on="entity", how="left")
+    df_merged = df_merged.drop("subject").rename({"index": "subject"})
+    # Step 3: Join on 'object' to replace object with its index
+    df_final = df_merged.join(idx_entity, left_on="object", right_on="entity", how="left")
+    df_final = df_final.drop("object").rename({"index": "object"})
+    # Step 4: Select the desired columns
+    df_final = df_final.select([polars.col("subject"), polars.col("relation"), polars.col("object")])
+    return df_final
+
+
+def pandas_dataframe_indexer(df_pandas: pd.DataFrame, idx_entity: pd.DataFrame, idx_relation: pd.DataFrame) -> pd.DataFrame:
+    """
+    Replaces 'subject', 'relation', and 'object' columns in the input Pandas DataFrame with their corresponding index values
+    from the entity and relation index DataFrames.
+
+    Parameters:
+    -----------
+    df_pandas : pd.DataFrame
+        The input Pandas DataFrame containing columns: 'subject', 'relation', and 'object'.
+
+    idx_entity : pd.DataFrame
+        A Pandas DataFrame that contains the mapping between entity names and their corresponding indices.
+        Must have columns: 'entity' and 'index'.
+
+    idx_relation : pd.DataFrame
+        A Pandas DataFrame that contains the mapping between relation names and their corresponding indices.
+        Must have columns: 'relation' and 'index'.
+
+    Returns:
+    --------
+    pd.DataFrame
+        A DataFrame with the 'subject', 'relation', and 'object' columns replaced by their corresponding indices.
+    """
+    assert isinstance(df_pandas, pd.DataFrame)
+    assert isinstance(idx_entity, pd.DataFrame)
+    assert isinstance(idx_relation, pd.DataFrame)
+
+    # Create a dictionary that maps entities to their indices
+    entity_to_index = pd.Series(idx_entity.index, index=idx_entity['entity']).to_dict()
+    df_pandas['subject'] = df_pandas['subject'].map(entity_to_index)
+    df_pandas['object'] = df_pandas['object'].map(entity_to_index)
+    del entity_to_index
+    relation_to_index = pd.Series(idx_relation.index, index=idx_relation['relation']).to_dict()
+    df_pandas['relation'] = df_pandas['relation'].map(relation_to_index)
+    del relation_to_index
+    return df_pandas
 
 
 def apply_reciprical_or_noise(add_reciprical: bool, eval_model: str, df: object = None, info: str = None):
@@ -40,32 +149,30 @@ def timeit_wrapper(*args, **kwargs):
 
 
 @timeit
-def read_with_polars(data_path, read_only_few: int = None, sample_triples_ratio: float = None) -> polars.DataFrame:
+def read_with_polars(data_path, read_only_few: int = None, sample_triples_ratio: float = None, separator:str=None) -> polars.DataFrame:
     """ Load and Preprocess via Polars """
+    assert separator is not None, "separator cannot be None"
     print(f'*** Reading {data_path} with Polars ***')
     # (1) Load the data.
-    if True:#data_path[-3:] in [".tar.gz",'txt', 'csv']:
-        print('Reading with polars.read_csv with sep **t** ...')
-        # TODO: if byte_pair_encoding=True, we should not use "\s+" as seperator I guess
+    #try:
+    if ".zst" in data_path:
+        df= polars.read_csv(data_path,n_rows=None if read_only_few is None else read_only_few)
+    else:
         df = polars.read_csv(data_path,
                              has_header=False,
                              low_memory=False,
                              n_rows=None if read_only_few is None else read_only_few,
                              columns=[0, 1, 2],
-                             dtypes=[polars.Utf8],  # str
+                             dtypes=[polars.String],
                              new_columns=['subject', 'relation', 'object'],
-                             separator="\t")  # \s+ doesn't work for polars
-    else:
-        if read_only_few is None:
-            df = polars.read_parquet(data_path, use_pyarrow=True)
-        else:
-            df = polars.read_parquet(data_path, n_rows=read_only_few)
+                             separator=separator)
+    #except ValueError as err:
+    #    raise ValueError(f"{err}\nYou may want to use a different separator.")
     # (2) Sample from (1).
     if sample_triples_ratio:
         print(f'Subsampling {sample_triples_ratio} of input data {df.shape}...')
         df = df.sample(frac=sample_triples_ratio)
         print(df.shape)
-
     # (3) Type heuristic prediction: If KG is an RDF KG, remove all triples where subject is not <?>.
     h = df.head().to_pandas()
     if sum(h["subject"].str.startswith('<')) + sum(h["relation"].str.startswith('<')) > 2:
@@ -75,13 +182,13 @@ def read_with_polars(data_path, read_only_few: int = None, sample_triples_ratio:
 
 
 @timeit
-def read_with_pandas(data_path, read_only_few: int = None, sample_triples_ratio: float = None):
+def read_with_pandas(data_path, read_only_few: int = None, sample_triples_ratio: float = None,separator:str=None):
+    assert separator is not None, "separator cannot be None"
     print(f'*** Reading {data_path} with Pandas ***')
-    if data_path[-3:] in ["ttl", 'txt', 'csv', 'zst']:
+    if data_path[-3:] in [".nt","ttl", 'txt', 'csv', 'zst']:
         print('Reading with pandas.read_csv with sep ** s+ ** ...')
-        # TODO: if byte_pair_encoding=True, we should not use "\s+" as seperator I guess
         df = pd.read_csv(data_path,
-                         sep="\s+",
+                         sep=separator,#"\s+",
                          header=None,
                          nrows=None if read_only_few is None else read_only_few,
                          usecols=[0, 1, 2],
@@ -111,8 +218,10 @@ def read_with_pandas(data_path, read_only_few: int = None, sample_triples_ratio:
 
 
 def read_from_disk(data_path: str, read_only_few: int = None,
-                   sample_triples_ratio: float = None, backend=None):
-    assert backend
+                   sample_triples_ratio: float = None, backend:str=None,separator:str=None)\
+        ->Tuple[polars.DataFrame,pd.DataFrame]:
+    assert backend is not None, "backend cannot be None"
+    assert separator is not None, f"separator cannot be None. Currently {separator}"
     # If path exits
     if glob.glob(data_path):
         # (1) Detect data format
@@ -120,20 +229,15 @@ def read_from_disk(data_path: str, read_only_few: int = None,
         if dformat in ["ttl", "owl", "turtle", "rdf/xml"] and backend != "rdflib":
             raise RuntimeError(
                 f"Data with **{dformat}** format cannot be read via --backend pandas or polars. Use --backend rdflib")
-
         if backend == 'pandas':
-            return read_with_pandas(data_path, read_only_few, sample_triples_ratio)
+            return read_with_pandas(data_path, read_only_few, sample_triples_ratio, separator)
         elif backend == 'polars':
-            return read_with_polars(data_path, read_only_few, sample_triples_ratio)
+            return read_with_polars(data_path, read_only_few, sample_triples_ratio, separator)
         elif backend == "rdflib":
             # Lazy import
             from rdflib import Graph
-
-            try:
-                assert dformat in ["ttl", "owl", "nt", "turtle", "rdf/xml", "n3", " n-triples"]
-            except AssertionError:
-                raise AssertionError(f"--backend {backend} and dataformat **{dformat}** is not matching. "
-                                     f"Use --backend pandas")
+            assert dformat in ["ttl", "owl", "nt", "turtle", "rdf/xml", "n3", " n-triples"],\
+                f"--backend {backend} and dataformat **{dformat}** is not matching. Use --backend pandas"
             return pd.DataFrame(data=[(str(s), str(p), str(o)) for s, p, o in Graph().parse(data_path)],
                                 columns=['subject', 'relation', 'object'], dtype=str)
         else:
@@ -315,25 +419,6 @@ def create_recipriocal_triples(x):
         x['subject'].to_frame(name='object'))], ignore_index=True)
 
 
-def index_triples_with_pandas(train_set, entity_to_idx: dict, relation_to_idx: dict) -> pd.core.frame.DataFrame:
-    """
-    :param train_set: pandas dataframe
-    :param entity_to_idx: a mapping from str to integer index
-    :param relation_to_idx: a mapping from str to integer index
-    :param num_core: number of cores to be used
-    :return: indexed triples, i.e., pandas dataframe
-    """
-    n, d = train_set.shape
-    train_set['subject'] = train_set['subject'].apply(lambda x: entity_to_idx.get(x))
-    train_set['relation'] = train_set['relation'].apply(lambda x: relation_to_idx.get(x))
-    train_set['object'] = train_set['object'].apply(lambda x: entity_to_idx.get(x))
-    # train_set = train_set.dropna(inplace=True)
-    if isinstance(train_set, pd.core.frame.DataFrame):
-        assert (n, d) == train_set.shape
-    else:
-        raise KeyError('Wrong type training data')
-    return train_set
-
 
 def dataset_sanity_checking(train_set: np.ndarray, num_entities: int, num_relations: int) -> None:
     """
diff --git a/dicee/sanity_checkers.py b/dicee/sanity_checkers.py
index 47642cd8..1e1463aa 100644
--- a/dicee/sanity_checkers.py
+++ b/dicee/sanity_checkers.py
@@ -32,11 +32,11 @@ def validate_knowledge_graph(args):
 
     elif args.path_single_kg is not None:
         if args.sparql_endpoint is not None or args.path_single_kg is not None:
-            print(f'The dataset_dir and sparql_endpoint arguments '
-                  f'must be None if path_single_kg is given.'
-                  f'***{args.dataset_dir}***\n'
-                  f'***{args.sparql_endpoint}***\n'
-                  f'These two parameters are set to None.')
+            #print(f'The dataset_dir and sparql_endpoint arguments '
+            #      f'must be None if path_single_kg is given.'
+            #      f'***{args.dataset_dir}***\n'
+            #      f'***{args.sparql_endpoint}***\n'
+            #      f'These two parameters are set to None.')
             args.dataset_dir = None
             args.sparql_endpoint = None
 
@@ -61,11 +61,11 @@ def validate_knowledge_graph(args):
                 f"Use --path_single_kg **folder/dataset.format**, if you have a single file.")
 
         if args.sparql_endpoint is not None or args.path_single_kg is not None:
-            print(f'The sparql_endpoint and path_single_kg arguments '
-                  f'must be None if dataset_dir is given.'
-                  f'***{args.sparql_endpoint}***\n'
-                  f'***{args.path_single_kg}***\n'
-                  f'These two parameters are set to None.')
+            #print(f'The sparql_endpoint and path_single_kg arguments '
+            #      f'must be None if dataset_dir is given.'
+            #      f'***{args.sparql_endpoint}***\n'
+            #      f'***{args.path_single_kg}***\n'
+            #      f'These two parameters are set to None.')
             args.sparql_endpoint = None
             args.path_single_kg = None
 
@@ -80,20 +80,11 @@ def validate_knowledge_graph(args):
 
 
 def sanity_checking_with_arguments(args):
-    try:
-        assert args.embedding_dim > 0
-    except AssertionError:
-        raise AssertionError(f'embedding_dim must be strictly positive. Currently:{args.embedding_dim}')
-
-    if args.scoring_technique not in ["AllvsAll", "KvsSample", "KvsAll", "NegSample", "1vsAll", "Pyke", "Sentence"]:
-        raise KeyError(f'Invalid training strategy => {args.scoring_technique}.')
-
-    assert args.learning_rate > 0
+    assert args.embedding_dim > 0,f"embedding_dim must be strictly positive. Currently:{args.embedding_dim}"
+    assert args.scoring_technique in ["AllvsAll", "1vsSample", "KvsSample","KvsAll", "NegSample", "1vsAll","Pyke", "Sentence"], f"Invalid training strategy => {args.scoring_technique}."
+    assert args.learning_rate > 0, f"Learning rate must be greater than 0. Currently:{args.learning_rate}"
     if args.num_folds_for_cv is None:
         args.num_folds_for_cv = 0
-    try:
-        assert args.num_folds_for_cv >= 0
-    except AssertionError:
-        raise AssertionError(f'num_folds_for_cv can not be negative. Currently:{args.num_folds_for_cv}')
+    assert args.num_folds_for_cv >= 0,f"num_folds_for_cv can not be negative. Currently:{args.num_folds_for_cv}"
     validate_knowledge_graph(args)
 
diff --git a/dicee/scripts/run.py b/dicee/scripts/run.py
index 3b95efaf..f0a6fbbe 100755
--- a/dicee/scripts/run.py
+++ b/dicee/scripts/run.py
@@ -4,20 +4,20 @@
 
 def get_default_arguments(description=None):
     """ Extends pytorch_lightning Trainer's arguments with ours """
-    # From "pytorch-lightning==1.6.4" to "lightning>=2.1.3",  'Trainer' has no attribute 'add_argparse_args'
-    # parser = pl.Trainer.add_argparse_args(argparse.ArgumentParser(add_help=False))
     parser = argparse.ArgumentParser(add_help=False)
     # Default Trainer param https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#methods
     # Knowledge graph related arguments
-    parser.add_argument("--dataset_dir", type=str, default="KGs/UMLS",
+    parser.add_argument("--dataset_dir", type=str, default=None,
                         help="The path of a folder containing train.txt, and/or valid.txt and/or test.txt"
                              ",e.g., KGs/UMLS")
     parser.add_argument("--sparql_endpoint", type=str, default=None,
                         help="An endpoint of a triple store, e.g. 'http://localhost:3030/mutagenesis/'. ")
-    parser.add_argument("--path_single_kg", type=str, default=None,
+    # TODO: Deprecate --path_single_kg
+    # TODO: --dataset_dir either be a single KG file or a folder.
+    parser.add_argument("--path_single_kg", type=str, default=None,#"/home/cdemir/Desktop/Softwares/dice-embeddings/dice.nt",
                         help="Path of a file corresponding to the input knowledge graph")
     # Saved files related arguments
-    parser.add_argument("--path_to_store_single_run", type=str, default=None,
+    parser.add_argument("--path_to_store_single_run", type=str, default=None,#"DBpedia",
                         help="A single directory created that contains related data about embeddings.")
     parser.add_argument("--storage_path", type=str, default='Experiments',
                         help="A directory named with time of execution under --storage_path "
@@ -27,6 +27,8 @@ def get_default_arguments(description=None):
     parser.add_argument("--backend", type=str, default="pandas",
                         choices=["pandas", "polars", "rdflib"],
                         help='Backend for loading, preprocessing, indexing input knowledge graph.')
+    parser.add_argument("--separator", type=str, default="\s+",
+                        help='Pandas \s+, t for \t polars works with the last two.')
     # Model related arguments
     parser.add_argument("--model", type=str,
                         default="Keci",
@@ -40,26 +42,26 @@ def get_default_arguments(description=None):
                         help="Available knowledge graph embedding models. "
                              "To use other knowledge graph embedding models available in python, e.g.,"
                              "**Pykeen_BoxE** and add this into choices")
-    parser.add_argument('--optim', type=str, default='Adam',
+    parser.add_argument('--optim', type=str, default='Adopt',
                         help='An optimizer',
-                        choices=['Adam', 'AdamW', 'SGD',"NAdam", "Adagrad", "ASGD"])
-    parser.add_argument('--embedding_dim', type=int, default=256,
+                        choices=['Adam', 'AdamW', 'SGD',"NAdam", "Adagrad", "ASGD", "Adopt"])
+    parser.add_argument('--embedding_dim', type=int, default=32,
                         help='Number of dimensions for an embedding vector. ')
-    parser.add_argument("--num_epochs", type=int, default=100, help='Number of epochs for training. ')
-    parser.add_argument('--batch_size', type=int, default=1024,
+    parser.add_argument("--num_epochs", type=int, default=10, help='Number of epochs for training. ')
+    parser.add_argument('--batch_size', type=int, default=32,
                         help='Mini batch size. If None, automatic batch finder is applied')
-    parser.add_argument("--lr", type=float, default=0.01)
+    parser.add_argument("--lr", type=float, default=0.1)
     parser.add_argument('--callbacks', type=json.loads,
                         default={},
                         help='{"PPE":{ "last_percent_to_consider": 10}}'
                              '"Perturb": {"level": "out", "ratio": 0.2, "method": "RN", "scaler": 0.3}')
     parser.add_argument("--trainer", type=str, default='PL',
-                        choices=['torchCPUTrainer', 'PL', 'torchDDP'],
-                        help='PL (pytorch lightning trainer), torchDDP (custom ddp), torchCPUTrainer (custom cpu only)')
-    parser.add_argument('--scoring_technique', default="KvsAll",
+                        choices=['torchCPUTrainer', 'PL', 'torchDDP', "TP"],
+                        help='PL (pytorch lightning trainer), torchDDP (custom ddp), torchCPUTrainer (custom cpu only), MP (Model Paralelisim)')
+    parser.add_argument('--scoring_technique', default="NegSample",
                         help="Training technique for knowledge graph embedding model",
-                        choices=["AllvsAll", "KvsAll", "1vsAll", "NegSample", "KvsSample"])
-    parser.add_argument('--neg_ratio', type=int, default=1,
+                        choices=["AllvsAll", "KvsAll", "1vsAll", "NegSample", "1vsSample", "KvsSample"])
+    parser.add_argument('--neg_ratio', type=int, default=2,
                         help='The number of negative triples generated per positive triple.')
     parser.add_argument('--weight_decay', type=float, default=0.0, help='L2 penalty e.g.(0.00001)')
     parser.add_argument('--input_dropout_rate', type=float, default=0.0)
@@ -121,6 +123,9 @@ def get_default_arguments(description=None):
     parser.add_argument("--swa",
                         action="store_true",
                         help="Stochastic weight averaging")
+    parser.add_argument("--auto_batch_finding",
+                        action="store_true",
+                        help="Find a batch size fitting in GPUs. Only available for TP trainer")
     parser.add_argument('--degree', type=int, default=0,
                         help='degree for polynomial embeddings')
 
diff --git a/dicee/static_funcs.py b/dicee/static_funcs.py
index 9ff1ad26..8685c0cb 100644
--- a/dicee/static_funcs.py
+++ b/dicee/static_funcs.py
@@ -2,7 +2,7 @@
 import torch
 import datetime
 from typing import Tuple, List
-from .models import CMult, Pyke, DistMult, KeciBase, Keci, TransE, DeCaL, DualE,\
+from .models import Pyke, DistMult, KeciBase, Keci, TransE, DeCaL, DualE,\
     ComplEx, AConEx, AConvO, AConvQ, ConvQ, ConvO, ConEx, QMult, OMult, Shallom, LFMult
 from .models.pykeen_models import PykeenKGE
 from .models.transformers import BytE
@@ -16,8 +16,10 @@
 from .models.base_model import BaseKGE
 import pickle
 from collections import defaultdict
-
+import polars as pl
 import requests
+import csv
+from .models.ensemble import EnsembleKGE
 
 def create_recipriocal_triples(x):
     """
@@ -75,17 +77,23 @@ def timeit_wrapper(*args, **kwargs):
 
     return timeit_wrapper
 
-
+# TODO:CD: Deprecate the pickle usage for data serialization.
 def save_pickle(*, data: object=None, file_path=str):
     if data:
         pickle.dump(data, open(file_path, "wb"))
     else:
         print("Input data is None. Nothing to save.")
-
+# TODO:CD: Deprecate the pickle usage for data serialization.
 def load_pickle(file_path=str):
     with open(file_path, 'rb') as f:
         return pickle.load(f)
 
+def load_term_mapping(file_path=str):
+    try:
+        return load_pickle(file_path=file_path+".p")
+    except FileNotFoundError:
+        print(f"python file not found\t{file_path} with .p extension")
+    return pl.read_csv(file_path + ".csv")
 
 # @TODO: Could these funcs can be merged?
 def select_model(args: dict, is_continual_training: bool = None, storage_path: str = None):
@@ -94,20 +102,40 @@ def select_model(args: dict, is_continual_training: bool = None, storage_path: s
     assert isinstance(is_continual_training, bool)
     assert isinstance(storage_path, str)
     if is_continual_training:
-        print('Loading pre-trained model...')
-        model, _ = intialize_model(args)
-        try:
-            weights = torch.load(storage_path + '/model.pt', torch.device('cpu'))
-            model.load_state_dict(weights)
-            for parameter in model.parameters():
-                parameter.requires_grad = True
-            model.train()
-        except FileNotFoundError:
-            print(f"{storage_path}/model.pt is not found. The model will be trained with random weights")
-        return model, _
+        # Check whether we have tensor parallelized KGE.
+        files_under_storage_path = [f for f in os.listdir(storage_path) if os.path.isfile(os.path.join(storage_path, f))]
+        num_of_partial_models_for_tensor_parallel= len([ i for i in files_under_storage_path if "partial" in i ])
+        if num_of_partial_models_for_tensor_parallel >= 1:
+            models=[]
+            labelling_flag=None
+            for i in range(num_of_partial_models_for_tensor_parallel):
+                model, labelling_flag = intialize_model(args)
+                weights = torch.load(storage_path + f'/model_partial_{i}.pt', torch.device('cpu'),weights_only=False)
+                model.load_state_dict(weights)
+                for parameter in model.parameters():
+                    parameter.requires_grad = True
+                model.train()
+                models.append(model)
+            return EnsembleKGE(pretrained_models=models), labelling_flag
+        else:
+            print('Loading pre-trained model...')
+            model, labelling_flag = intialize_model(args)
+            try:
+                weights = torch.load(storage_path + '/model.pt', torch.device('cpu'))
+                model.load_state_dict(weights)
+                for parameter in model.parameters():
+                    parameter.requires_grad = True
+                model.train()
+            except FileNotFoundError as e:
+                print(f"{storage_path}/model.pt is not found. The model will be trained with random weights")
+                raise e
+            return model, labelling_flag
     else:
-        return intialize_model(args)
+        model, labelling_flag= intialize_model(args)
+        if args["trainer"]=="TP":
+            model=EnsembleKGE(seed_model=model)
 
+    return model, labelling_flag
 
 def load_model(path_of_experiment_folder: str, model_name='model.pt',verbose=0) -> Tuple[object, Tuple[dict, dict]]:
     """ Load weights and initialize pytorch module from namespace arguments"""
@@ -117,6 +145,7 @@ def load_model(path_of_experiment_folder: str, model_name='model.pt',verbose=0)
     # (1) Load weights..
     weights = torch.load(path_of_experiment_folder + f'/{model_name}', torch.device('cpu'))
     configs = load_json(path_of_experiment_folder + '/configuration.json')
+    reports = load_json(path_of_experiment_folder + '/report.json')
 
     if configs.get("byte_pair_encoding", None):
         num_tokens, ent_dim = weights['token_embeddings.weight'].shape
@@ -128,9 +157,9 @@ def load_model(path_of_experiment_folder: str, model_name='model.pt',verbose=0)
         configs["num_tokens"] = num_tokens
         configs["max_length_subword_tokens"] = report["max_length_subword_tokens"]
     else:
-        num_ent, ent_dim = weights['entity_embeddings.weight'].shape
-        num_rel, rel_dim = weights['relation_embeddings.weight'].shape
-        assert ent_dim == rel_dim
+
+        num_ent = reports["num_entities"]
+        num_rel = reports["num_relations"]
         # Update the training configuration
         configs["num_entities"] = num_ent
         configs["num_relations"] = num_rel
@@ -139,7 +168,10 @@ def load_model(path_of_experiment_folder: str, model_name='model.pt',verbose=0)
     # (4) Select the model
     model, _ = intialize_model(configs,verbose)
     # (5) Put (1) into (4)
-    model.load_state_dict(weights)
+    if isinstance(weights,torch.jit._script.RecursiveScriptModule):
+        model.load_state_dict(weights.state_dict())
+    else:
+        model.load_state_dict(weights)
     # (6) Set it into eval model.
     for parameter in model.parameters():
         parameter.requires_grad = False
@@ -152,17 +184,20 @@ def load_model(path_of_experiment_folder: str, model_name='model.pt',verbose=0)
             print('Loading entity and relation indexes...', end=' ')
         try:
             # Maybe ? https://docs.python.org/3/library/mmap.html
+            # TODO:CD: Deprecate the pickle usage for data serialization.
+            # TODO: CD: We do not need to keep the mapping in memory
             with open(path_of_experiment_folder + '/entity_to_idx.p', 'rb') as f:
                 entity_to_idx = pickle.load(f)
         except FileNotFoundError:
-            print("entity_to_idx.p not found")
-            entity_to_idx = dict()
+            entity_to_idx = { v["entity"]:k for k,v in pd.read_csv(f"{path_of_experiment_folder}/entity_to_idx.csv",index_col=0).to_dict(orient='index').items()}
+
         try:
+            # TODO:CD: Deprecate the pickle usage for data serialization.
+            # TODO: CD: We do not need to keep the mapping in memory
             with open(path_of_experiment_folder + '/relation_to_idx.p', 'rb') as f:
                 relation_to_idx = pickle.load(f)
         except FileNotFoundError:
-            print("relation_to_idx.p not found")
-            relation_to_idx = dict()
+            relation_to_idx = { v["relation"]:k for k,v in pd.read_csv(f"{path_of_experiment_folder}/relation_to_idx.csv",index_col=0).to_dict(orient='index').items()}
         if verbose > 0:
             print(f'Done! It took {time.time() - start_time:.4f}')
         return model, (entity_to_idx, relation_to_idx)
@@ -218,10 +253,16 @@ def load_model_ensemble(path_of_experiment_folder: str) -> Tuple[BaseKGE, Tuple[
     model.eval()
     start_time = time.time()
     print('Loading entity and relation indexes...', end=' ')
-    with open(path_of_experiment_folder + '/entity_to_idx.p', 'rb') as f:
-        entity_to_idx = pickle.load(f)
-    with open(path_of_experiment_folder + '/relation_to_idx.p', 'rb') as f:
-        relation_to_idx = pickle.load(f)
+    # TODO: CD: We do not need to keep the mapping in memory
+    # TODO:CD: Deprecate the pickle usage for data serialization.
+
+    entity_to_idx = {v["entity"]: k for k, v in
+                     pd.read_csv(f"{path_of_experiment_folder}/entity_to_idx.csv", index_col=0).to_dict(
+                         orient='index').items()}
+    relation_to_idx = {v["relation"]: k for k, v in
+                     pd.read_csv(f"{path_of_experiment_folder}/relation_to_idx.csv", index_col=0).to_dict(
+                         orient='index').items()}
+
     assert isinstance(entity_to_idx, dict)
     assert isinstance(relation_to_idx, dict)
     print(f'Done! It took {time.time() - start_time:.4f}')
@@ -258,54 +299,38 @@ def numpy_data_type_changer(train_set: np.ndarray, num: int) -> np.ndarray:
 def save_checkpoint_model(model, path: str) -> None:
     """ Store Pytorch model into disk"""
     if isinstance(model, BaseKGE):
-        try:
-            torch.save(model.state_dict(), path)
-        except ReferenceError as e:
-            print(e)
-            print(model.name)
-            print('Could not save the model correctly')
+        torch.save(model.state_dict(), path)
+    elif isinstance(model, EnsembleKGE):
+        # path comes with ../model_...
+        for i, partial_model in enumerate(model):
+            new_path=path.replace("model.pt",f"model_partial_{i}.pt")
+            torch.save(partial_model.state_dict(), new_path)
     else:
         torch.save(model.model.state_dict(), path)
 
 
-def store(trainer,
-          trained_model, model_name: str = 'model', full_storage_path: str = None,
+def store(trained_model, model_name: str = 'model', full_storage_path: str = None,
           save_embeddings_as_csv=False) -> None:
-    """
-    Store trained_model model and save embeddings into csv file.
-    :param trainer: an instance of trainer class
-    :param full_storage_path: path to save parameters.
-    :param model_name: string representation of the name of the model.
-    :param trained_model: an instance of BaseKGE see core.models.base_model .
-    :param save_embeddings_as_csv: for easy access of embeddings.
-    :return:
-    """
     assert full_storage_path is not None
     assert isinstance(model_name, str)
     assert len(model_name) > 1
-
-    # (1) Save pytorch model in trained_model .
     save_checkpoint_model(model=trained_model, path=full_storage_path + f'/{model_name}.pt')
+
     if save_embeddings_as_csv:
         entity_emb, relation_ebm = trained_model.get_embeddings()
-        entity_to_idx = pickle.load(open(full_storage_path + '/entity_to_idx.p', 'rb'))
-        entity_str = entity_to_idx.keys()
-        # Ensure that the ordering is correct.
-        assert list(range(0, len(entity_str))) == list(entity_to_idx.values())
-        save_embeddings(entity_emb.numpy(), indexes=entity_str,
-                        path=full_storage_path + '/' + trained_model.name + '_entity_embeddings.csv')
-        del entity_to_idx, entity_str, entity_emb
+        print("Saving entity embeddings...")
+        entity=pd.read_csv(f"{full_storage_path}/entity_to_idx.csv",index_col=0)["entity"]
+        assert entity.index.is_monotonic_increasing
+        save_embeddings(entity_emb.numpy(), indexes=entity.to_list(), path=full_storage_path + '/' + trained_model.name + '_entity_embeddings.csv')
+        del entity, entity_emb
         if relation_ebm is not None:
-            relation_to_idx = pickle.load(open(full_storage_path + '/relation_to_idx.p', 'rb'))
-            relations_str = relation_to_idx.keys()
-
-            save_embeddings(relation_ebm.numpy(), indexes=relations_str,
-                            path=full_storage_path + '/' + trained_model.name + '_relation_embeddings.csv')
-            del relation_ebm, relations_str, relation_to_idx
+            print("Saving relation embeddings...")
+            relations = pd.read_csv(f"{full_storage_path}/relation_to_idx.csv", index_col=0)["relation"]
+            assert relations.index.is_monotonic_increasing
+            save_embeddings(relation_ebm.numpy(), indexes=relations, path=full_storage_path + '/' + trained_model.name + '_relation_embeddings.csv')
         else:
             pass
 
-
 def add_noisy_triples(train_set: pd.DataFrame, add_noise_rate: float) -> pd.DataFrame:
     """
     Add randomly constructed triples
@@ -335,22 +360,24 @@ def add_noisy_triples(train_set: pd.DataFrame, add_noise_rate: float) -> pd.Data
     assert num_triples + num_noisy_triples == len(train_set)
     return train_set
 
-
 def read_or_load_kg(args, cls):
     print('*** Read or Load Knowledge Graph  ***')
     start_time = time.time()
     kg = cls(dataset_dir=args.dataset_dir,
              byte_pair_encoding=args.byte_pair_encoding,
+             padding=True if args.byte_pair_encoding and args.model != "BytE" else False,
              add_noise_rate=args.add_noise_rate,
              sparql_endpoint=args.sparql_endpoint,
              path_single_kg=args.path_single_kg,
-             add_reciprical=args.apply_reciprical_or_noise,
+             add_reciprocal=args.apply_reciprical_or_noise,
              eval_model=args.eval_model,
              read_only_few=args.read_only_few,
              sample_triples_ratio=args.sample_triples_ratio,
              path_for_serialization=args.full_storage_path,
              path_for_deserialization=args.path_experiment_folder if hasattr(args, 'path_experiment_folder') else None,
-             backend=args.backend)
+             backend=args.backend,
+             training_technique=args.scoring_technique,
+             separator=args.separator)
     print(f'Preprocessing took: {time.time() - start_time:.3f} seconds')
     # (2) Share some info about data for easy access.
     print(kg.description_of_input)
@@ -409,9 +436,6 @@ def intialize_model(args: dict,verbose=0) -> Tuple[object, str]:
     elif model_name == 'KeciBase':
         model = KeciBase(args=args)
         form_of_labelling = 'EntityPrediction'
-    elif model_name == 'CMult':
-        model = CMult(args=args)
-        form_of_labelling = 'EntityPrediction'
     elif model_name == 'BytE':
         model = BytE(args=args)
         form_of_labelling = 'EntityPrediction'
@@ -511,19 +535,14 @@ def create_experiment_folder(folder_name='Experiments'):
 
 
 def continual_training_setup_executor(executor) -> None:
-    """
-    storage_path:str A path leading to a parent directory, where a subdirectory containing KGE related data
-
-    full_storage_path:str A path leading to a subdirectory containing KGE related data
-
-    """
+    # TODO:CD:Deprecate it
     if executor.is_continual_training:
         # (4.1) If it is continual, then store new models on previous path.
         executor.storage_path = executor.args.full_storage_path
     else:
         # Create a single directory containing KGE and all related data
         if executor.args.path_to_store_single_run:
-            os.makedirs(executor.args.path_to_store_single_run, exist_ok=True)
+            os.makedirs(executor.args.path_to_store_single_run, exist_ok=False)
             executor.args.full_storage_path = executor.args.path_to_store_single_run
         else:
             # Create a parent and subdirectory.
@@ -641,8 +660,11 @@ def download_files_from_url(base_url:str, destination_folder=".")->None:
 
     """
     # lazy import
-    from bs4 import BeautifulSoup
-
+    try:
+        from bs4 import BeautifulSoup
+    except ModuleNotFoundError:
+        print("Please install the 'beautifulsoup4' package by running: pip install beautifulsoup4")
+        raise
     response = requests.get(base_url)
     if response.status_code == 200:
         soup = BeautifulSoup(response.text, 'html.parser')
@@ -667,3 +689,83 @@ def download_pretrained_model(url: str) -> str:
         os.mkdir(dir_name)
         download_files_from_url(url_to_download_from, destination_folder=dir_name)
     return dir_name
+
+def write_csv_from_model_parallel(path: str) :
+    """Create"""
+    assert os.path.exists(path), "Path does not exist"
+    # Detect files that start with model_ and end with .pt
+    model_files = [f for f in os.listdir(path) if f.startswith("model_") and f.endswith(".pt")]
+    model_files.sort()  # Sort to maintain order if necessary (e.g., model_0.pt, model_1.pt)
+
+    entity_embeddings=[]
+    relation_embeddings=[]
+
+    # Process each model file
+    for model_file in model_files:
+        model_path = os.path.join(path, model_file)
+        # Load model
+        model = torch.load(model_path)
+        # Assuming model has a get_embeddings method
+        entity_emb, relation_emb = model["_orig_mod.entity_embeddings.weight"], model["_orig_mod.relation_embeddings.weight"]
+        entity_embeddings.append(entity_emb)
+        relation_embeddings.append(relation_emb)
+
+    return torch.cat(entity_embeddings, dim=1), torch.cat(relation_embeddings, dim=1)
+
+
+def from_pretrained_model_write_embeddings_into_csv(path: str) -> None:
+    """ """
+    assert os.path.exists(path), "Path does not exist"
+    config = load_json(path + '/configuration.json')
+    entity_csv_path = os.path.join(path, f"{config['model']}_entity_embeddings.csv")
+    relation_csv_path = os.path.join(path, f"{config['model']}_relation_embeddings.csv")
+
+    if config["trainer"]=="TP":
+        entity_emb, relation_emb = write_csv_from_model_parallel(path)
+    else:
+        # Load model
+        model = torch.load(os.path.join(path, "model.pt"))
+        # Assuming model has a get_embeddings method
+        entity_emb, relation_emb = model["entity_embeddings.weight"], model["relation_embeddings.weight"]
+    str_entity = pd.read_csv(f"{path}/entity_to_idx.csv", index_col=0)["entity"]
+    assert str_entity.index.is_monotonic_increasing
+    str_entity=str_entity.to_list()
+    # Write entity embeddings with headers and indices
+    with open(entity_csv_path, "w", newline="") as f:
+        writer = csv.writer(f)
+        # Add header (e.g., "", "0", "1", ..., "N")
+        headers = [""] + [f"{i}" for i in range(entity_emb.size(1))]
+        writer.writerow(headers)
+        # Add rows with index
+        for i_row, (name,row) in enumerate(zip(str_entity,entity_emb)):
+            writer.writerow([name] + row.tolist())
+    str_relations = pd.read_csv(f"{path}/relation_to_idx.csv", index_col=0)["relation"]
+    assert str_relations.index.is_monotonic_increasing
+
+    # Write relation embeddings with headers and indices
+    with open(relation_csv_path, "w", newline="") as f:
+        writer = csv.writer(f)
+        # Add header (e.g., "", "0", "1", ..., "N")
+        headers = [""] + [f"{i}" for i in range(relation_emb.size(1))]
+        writer.writerow(headers)
+        # Add rows with index
+        for i_row, (name, row) in enumerate(zip(str_relations,relation_emb)):
+            writer.writerow([name]+ row.tolist())
+
+    """
+    
+    # Write entity embeddings directly to CSV
+    with open(entity_csv_path, "w") as f:
+        for row in entity_emb:
+            f.write(",".join(map(str, row.tolist())) + "\n")
+
+    # Write relation embeddings directly to CSV
+    with open(relation_csv_path, "w") as f:
+        for row in relation_emb:
+            f.write(",".join(map(str, row.tolist())) + "\n")
+
+    # Convert to numpy
+    pd.DataFrame(entity_emb.numpy()).to_csv(entity_csv_path, index=True, header=False)
+    # If CSV files do not exist, create them
+    pd.DataFrame(relation_emb.numpy()).to_csv(relation_csv_path, index=True, header=False)
+    """
\ No newline at end of file
diff --git a/dicee/static_funcs_training.py b/dicee/static_funcs_training.py
index 7f4992a7..6b594e06 100644
--- a/dicee/static_funcs_training.py
+++ b/dicee/static_funcs_training.py
@@ -1,91 +1,195 @@
 import torch
-from typing import Dict, Tuple, List
+from typing import Dict, Tuple, List, Iterable
 import numpy as np
 from tqdm import tqdm
 
-def evaluate_lp(model, triple_idx, num_entities, er_vocab: Dict[Tuple, List], re_vocab: Dict[Tuple, List],
-                info='Eval Starts'):
-    """
-    Evaluate model in a standard link prediction task
-
-    for each triple
-    the rank is computed by taking the mean of the filtered missing head entity rank and
-    the filtered missing tail entity rank
-    :param model:
-    :param triple_idx:
-    :param info:
-    :return:
-    """
+def make_iterable_verbose(iterable_object, verbose, desc="Default", position=None, leave=True) -> Iterable:
+    if verbose:
+        return tqdm(iterable_object, desc=desc, position=position, leave=leave)
+    else:
+        return iterable_object
+
+@torch.no_grad()
+def evaluate_lp(model=None, triple_idx=None, num_entities=None, er_vocab: Dict[Tuple, List]=None,
+                re_vocab: Dict[Tuple, List]=None,
+                info='Eval Starts', batch_size=128, chunk_size=1000):
+    assert model is not None, "Model must be provided"
+    assert triple_idx is not None, "triple_idx must be provided"
+    assert num_entities is not None, "num_entities must be provided"
+    assert er_vocab is not None, "er_vocab must be provided"
+    assert re_vocab is not None, "re_vocab must be provided"
+
     model.eval()
     print(info)
     print(f'Num of triples {len(triple_idx)}')
-    print('** Evaluation without batching')
     hits = dict()
     reciprocal_ranks = []
     # Iterate over test triples
     all_entities = torch.arange(0, num_entities).long()
     all_entities = all_entities.reshape(len(all_entities), )
-    # Iterating one by one is not good when you are using batch norm
-    for i in tqdm(range(0, len(triple_idx))):
-        # (1) Get a triple (head entity, relation, tail entity
-        data_point = triple_idx[i]
-        h, r, t = data_point[0], data_point[1], data_point[2]
-
-        # (2) Predict missing heads and tails
-        x = torch.stack((torch.tensor(h).repeat(num_entities, ),
-                         torch.tensor(r).repeat(num_entities, ),
-                         all_entities), dim=1)
-
-        predictions_tails = model.forward_triples(x)
-        x = torch.stack((all_entities,
-                         torch.tensor(r).repeat(num_entities, ),
-                         torch.tensor(t).repeat(num_entities)
-                         ), dim=1)
-
-        predictions_heads = model.forward_triples(x)
-        del x
+    
+    # Evaluation without Batching
+    # for i in tqdm(range(0, len(triple_idx))):
+    #     # (1) Get a triple (head entity, relation, tail entity
+    #     data_point = triple_idx[i]
+    #     h, r, t = data_point[0], data_point[1], data_point[2]
 
-        # 3. Computed filtered ranks for missing tail entities.
-        # 3.1. Compute filtered tail entity rankings
-        filt_tails = er_vocab[(h, r)]
-        # 3.2 Get the predicted target's score
-        target_value = predictions_tails[t].item()
-        # 3.3 Filter scores of all triples containing filtered tail entities
-        predictions_tails[filt_tails] = -np.Inf
-        # 3.4 Reset the target's score
-        predictions_tails[t] = target_value
-        # 3.5. Sort the score
-        _, sort_idxs = torch.sort(predictions_tails, descending=True)
-        sort_idxs = sort_idxs.detach()
-        filt_tail_entity_rank = np.where(sort_idxs == t)[0][0]
+    #     # (2) Predict missing heads and tails
+    #     x = torch.stack((torch.tensor(h).repeat(num_entities, ),
+    #                      torch.tensor(r).repeat(num_entities, ),
+    #                      all_entities), dim=1)
 
-        # 4. Computed filtered ranks for missing head entities.
-        # 4.1. Retrieve head entities to be filtered
-        filt_heads = re_vocab[(r, t)]
-        # 4.2 Get the predicted target's score
-        target_value = predictions_heads[h].item()
-        # 4.3 Filter scores of all triples containing filtered head entities.
-        predictions_heads[filt_heads] = -np.Inf
-        predictions_heads[h] = target_value
-        _, sort_idxs = torch.sort(predictions_heads, descending=True)
-        sort_idxs = sort_idxs.detach()
-        filt_head_entity_rank = np.where(sort_idxs == h)[0][0]
+    #     predictions_tails = model(x)
+    #     x = torch.stack((all_entities,
+    #                      torch.tensor(r).repeat(num_entities, ),
+    #                      torch.tensor(t).repeat(num_entities)
+    #                      ), dim=1)
 
-        # 4. Add 1 to ranks as numpy array first item has the index of 0.
-        filt_head_entity_rank += 1
-        filt_tail_entity_rank += 1
+    #     predictions_heads = model(x)
+    #     del x
 
-        rr = 1.0 / filt_head_entity_rank + (1.0 / filt_tail_entity_rank)
-        # 5. Store reciprocal ranks.
-        reciprocal_ranks.append(rr)
-        # print(f'{i}.th triple: mean reciprical rank:{rr}')
+    #     # 3. Computed filtered ranks for missing tail entities.
+    #     # 3.1. Compute filtered tail entity rankings
+    #     filt_tails = er_vocab[(h, r)]
+    #     # 3.2 Get the predicted target's score
+    #     target_value = predictions_tails[t].item()
+    #     # 3.3 Filter scores of all triples containing filtered tail entities
+    #     predictions_tails[filt_tails] = -np.Inf
+    #     # 3.4 Reset the target's score
+    #     predictions_tails[t] = target_value
+    #     # 3.5. Sort the score
+    #     _, sort_idxs = torch.sort(predictions_tails, descending=True)
+    #     sort_idxs = sort_idxs.detach()
+    #     filt_tail_entity_rank = np.where(sort_idxs == t)[0][0]
 
-        # 4. Compute Hit@N
-        for hits_level in range(1, 11):
-            res = 1 if filt_head_entity_rank <= hits_level else 0
-            res += 1 if filt_tail_entity_rank <= hits_level else 0
-            if res > 0:
-                hits.setdefault(hits_level, []).append(res)
+    #     # 4. Computed filtered ranks for missing head entities.
+    #     # 4.1. Retrieve head entities to be filtered
+    #     filt_heads = re_vocab[(r, t)]
+    #     # 4.2 Get the predicted target's score
+    #     target_value = predictions_heads[h].item()
+    #     # 4.3 Filter scores of all triples containing filtered head entities.
+    #     predictions_heads[filt_heads] = -np.Inf
+    #     predictions_heads[h] = target_value
+    #     _, sort_idxs = torch.sort(predictions_heads, descending=True)
+    #     sort_idxs = sort_idxs.detach()
+    #     filt_head_entity_rank = np.where(sort_idxs == h)[0][0]
+
+    #     # 4. Add 1 to ranks as numpy array first item has the index of 0.
+    #     filt_head_entity_rank += 1
+    #     filt_tail_entity_rank += 1
+
+    #     rr = 1.0 / filt_head_entity_rank + (1.0 / filt_tail_entity_rank)
+    #     # 5. Store reciprocal ranks.
+    #     reciprocal_ranks.append(rr)
+    #     # print(f'{i}.th triple: mean reciprical rank:{rr}')
+
+    #     # 4. Compute Hit@N
+    #     for hits_level in range(1, 11):
+    #         res = 1 if filt_head_entity_rank <= hits_level else 0
+    #         res += 1 if filt_tail_entity_rank <= hits_level else 0
+    #         if res > 0:
+    #             hits.setdefault(hits_level, []).append(res)
+    
+    # Evaluation with Batching
+    for batch_start in tqdm(range(0, len(triple_idx), batch_size), desc="Evaluating Batches"):
+        batch_end = min(batch_start + batch_size, len(triple_idx))
+        batch_triples = triple_idx[batch_start:batch_end]
+        batch_size_current = len(batch_triples)
+        
+        # (1) Extract heads, relations, and tails for the batch
+        h_batch = torch.tensor([data_point[0] for data_point in batch_triples])
+        r_batch = torch.tensor([data_point[1] for data_point in batch_triples])
+        t_batch = torch.tensor([data_point[2] for data_point in batch_triples])
+
+        # Initialize score tensors
+        predictions_tails = torch.zeros(batch_size_current, num_entities)
+        predictions_heads = torch.zeros(batch_size_current, num_entities)
+
+        # Process entities in chunks to manage memory usage
+        for chunk_start in range(0, num_entities, chunk_size):
+            chunk_end = min(chunk_start + chunk_size, num_entities)
+            entities_chunk = all_entities[chunk_start:chunk_end]
+            chunk_size_current = entities_chunk.size(0)
+            
+            # (2) Predict missing heads and tails
+            # Prepare input tensors for tail prediction
+            x_tails = torch.stack((
+                h_batch.repeat_interleave(chunk_size_current),
+                r_batch.repeat_interleave(chunk_size_current),
+                entities_chunk.repeat(batch_size_current)
+            ), dim=1)
+
+            # Predict scores for missing tails
+            preds_tails = model(x_tails)
+            preds_tails = preds_tails.view(batch_size_current, chunk_size_current)
+            predictions_tails[:, chunk_start:chunk_end] = preds_tails
+            del x_tails
+
+            # Prepare input tensors for head prediction
+            x_heads = torch.stack((
+                entities_chunk.repeat(batch_size_current),
+                r_batch.repeat_interleave(chunk_size_current),
+                t_batch.repeat_interleave(chunk_size_current)
+            ), dim=1)
+
+            # Predict scores for missing heads
+            preds_heads = model(x_heads)
+            preds_heads = preds_heads.view(batch_size_current, chunk_size_current)
+            predictions_heads[:, chunk_start:chunk_end] = preds_heads
+            del x_heads
+    
+        # Iterating one by one is not good when you are using batch norm
+        for i in range(0, batch_size_current):
+        # (1) Get a triple (head entity, relation, tail entity)
+            h = h_batch[i].item()
+            r = r_batch[i].item()
+            t = t_batch[i].item()
+
+            # 3. Computed filtered ranks for missing tail entities.
+            # 3.1. Compute filtered tail entity rankings
+            filt_tails = er_vocab[(h, r)]
+            filt_tails_set = set(filt_tails) - {t}
+            filt_tails_indices = list(filt_tails_set)
+            # 3.2 Get the predicted target's score
+            target_value = predictions_tails[i, t].item()
+            # 3.3 Filter scores of all triples containing filtered tail entities
+            predictions_tails[i, filt_tails_indices] = -np.Inf
+            # 3.4 Reset the target's score
+            predictions_tails[i, t] = target_value
+            # 3.5. Sort the score
+            _, sort_idxs = torch.sort(predictions_tails[i], descending=True)
+            sort_idxs = sort_idxs.detach()
+            filt_tail_entity_rank = np.where(sort_idxs == t)[0][0]
+
+            # 4. Computed filtered ranks for missing head entities.
+            # 4.1. Retrieve head entities to be filtered
+            filt_heads = re_vocab[(r, t)]
+            filt_heads_set = set(filt_heads) - {h}
+            filt_heads_indices = list(filt_heads_set)
+            # 4.2 Get the predicted target's score
+            target_value = predictions_heads[i, h].item()
+            # 4.3 Filter scores of all triples containing filtered head entities.
+            predictions_heads[i, filt_heads_indices] = -np.Inf
+            predictions_heads[i, h] = target_value
+            _, sort_idxs = torch.sort(predictions_heads[i], descending=True)
+            sort_idxs = sort_idxs.detach()
+            filt_head_entity_rank = np.where(sort_idxs == h)[0][0]
+
+            # 4. Add 1 to ranks as numpy array first item has the index of 0.
+            filt_head_entity_rank += 1
+            filt_tail_entity_rank += 1
+
+            rr = 1.0 / filt_head_entity_rank + (1.0 / filt_tail_entity_rank)
+            # 5. Store reciprocal ranks.
+            reciprocal_ranks.append(rr)
+            # print(f'{i}.th triple: mean reciprical rank:{rr}')
+
+            # 4. Compute Hit@N
+            for hits_level in range(1, 11):
+                res = 1 if filt_head_entity_rank <= hits_level else 0
+                res += 1 if filt_tail_entity_rank <= hits_level else 0
+                if res > 0:
+                    hits.setdefault(hits_level, []).append(res)
 
     mean_reciprocal_rank = sum(reciprocal_ranks) / (float(len(triple_idx) * 2))
 
diff --git a/dicee/static_preprocess_funcs.py b/dicee/static_preprocess_funcs.py
index 4ef2ce84..bfdb1f0a 100644
--- a/dicee/static_preprocess_funcs.py
+++ b/dicee/static_preprocess_funcs.py
@@ -39,37 +39,29 @@ def preprocesses_input_args(args):
     assert args.weight_decay >= 0.0
     args.learning_rate = args.lr
     args.deterministic = True
-
     assert args.init_param in ['xavier_normal', None]
-
     # No need to eval. Investigate runtime performance
     args.check_val_every_n_epoch = 10 ** 6  # ,i.e., no eval
     assert args.add_noise_rate is None or isinstance(args.add_noise_rate, float)
     args.logger = False
-    try:
-        assert args.eval_model in [None, 'None', 'train', 'val', 'test', 'train_val', 'train_test', 'val_test',
-                                   'train_val_test']
-    except AssertionError:
-        raise AssertionError(f'Unexpected input for eval_model ***\t{args.eval_model}\t***')
-
+    assert args.eval_model in [None, 'None', 'train', 'val', 'test', 'train_val', 'train_test', 'val_test',
+                                   'train_val_test'], f'Unexpected input for eval_model ***\t{args.eval_model}\t***'
     if args.eval_model == 'None':
         args.eval_model = None
-
     # reciprocal checking
-    if args.scoring_technique in ["AllvsAll", "KvsSample", "KvsAll", "1vsAll"]:
+    if args.scoring_technique in ["AllvsAll", "1vsSample", "KvsAll", "1vsAll", "KvsSample"]:
         args.apply_reciprical_or_noise = True
     elif args.scoring_technique in ["NegSample", "Sentence"]:
         args.apply_reciprical_or_noise = False
     else:
         raise KeyError(f'Unexpected input for scoring_technique \t{args.scoring_technique}')
-
     if args.sample_triples_ratio is not None:
         assert 1.0 >= args.sample_triples_ratio >= 0.0
     assert args.backend in ["pandas", "polars", "rdflib"]
     sanity_checking_with_arguments(args)
     if args.model == 'Shallom':
         args.scoring_technique = 'KvsAll'
-    # TODO: we need need to define as "NONE ?
+
     if args.normalization == 'None':
         args.normalization = None
     assert args.normalization in [None, 'LayerNorm', 'BatchNorm1d']
diff --git a/dicee/trainer/dice_trainer.py b/dicee/trainer/dice_trainer.py
index 4ab3376d..2808e48a 100644
--- a/dicee/trainer/dice_trainer.py
+++ b/dicee/trainer/dice_trainer.py
@@ -1,15 +1,14 @@
 import lightning as pl
-
-import gc
-
+import polars
 from typing import Union
-
 from dicee.models.base_model import BaseKGE
 from dicee.static_funcs import select_model
 from dicee.callbacks import ASWA, Eval, KronE, PrintCallback, AccumulateEpochLossCallback, Perturb
-from dicee.dataset_classes import construct_dataset, reload_dataset
+from dicee.dataset_classes import construct_dataset
 from .torch_trainer import TorchTrainer
 from .torch_trainer_ddp import TorchDDPTrainer
+from .model_parallelism import TensorParallel
+from ..models.ensemble import EnsembleKGE
 from ..static_funcs import timeit
 import os
 import torch
@@ -17,19 +16,23 @@
 import copy
 from typing import List, Tuple
 from ..knowledge_graph import KG
+import numpy as np
 
+def load_term_mapping(file_path=str):
+    return polars.read_csv(file_path + ".csv")
 
-def initialize_trainer(args, callbacks):
+
+def initialize_trainer(args, callbacks)->TorchTrainer | TensorParallel | TorchDDPTrainer | pl.Trainer:
     if args.trainer == 'torchCPUTrainer':
         print('Initializing TorchTrainer CPU Trainer...', end='\t')
-        return TorchTrainer(args, callbacks=callbacks)
+        trainer = TorchTrainer(args, callbacks=callbacks)
+    elif args.trainer == 'TP':
+        print('Initializing TensorParallel...', end='\t')
+        trainer= TensorParallel(args, callbacks=callbacks)
     elif args.trainer == 'torchDDP':
-        if torch.cuda.is_available():
-            print('Initializing TorchDDPTrainer GPU', end='\t')
-            return TorchDDPTrainer(args, callbacks=callbacks)
-        else:
-            print('Initializing TorchTrainer CPU Trainer', end='\t')
-            return TorchTrainer(args, callbacks=callbacks)
+        assert torch.cuda.is_available()
+        print('Initializing TorchDDPTrainer GPU', end='\t')
+        trainer = TorchDDPTrainer(args, callbacks=callbacks)
     elif args.trainer == 'PL':
         print('Initializing Pytorch-lightning Trainer', end='\t')
         kwargs = vars(args)
@@ -63,7 +66,7 @@ def initialize_trainer(args, callbacks):
         default_root_dir: Optional[_PATH] = None,)
         """
         # @TODO: callbacks need to be ad
-        return pl.Trainer(accelerator=kwargs.get("accelerator", "auto"),
+        trainer= pl.Trainer(accelerator=kwargs.get("accelerator", "auto"),
                           strategy=kwargs.get("strategy", "auto"),
                           num_nodes=kwargs.get("num_nodes", 1),
                           precision=kwargs.get("precision", None),
@@ -77,8 +80,11 @@ def initialize_trainer(args, callbacks):
                           detect_anomaly=False,
                           barebones=False)
     else:
-        print('Initialize TorchTrainer CPU Trainer', end='\t')
-        return TorchTrainer(args, callbacks=callbacks)
+        print('Initializing TorchTrainer CPU Trainer...', end='\t')
+        trainer = TorchTrainer(args, callbacks=callbacks)
+    assert trainer is not None
+    return trainer
+
 
 
 def get_callbacks(args):
@@ -108,7 +114,6 @@ def get_callbacks(args):
             raise RuntimeError(f'Incorrect callback:{k}')
     return callbacks
 
-
 class DICE_Trainer:
     """
    DICE_Trainer implement
@@ -131,7 +136,7 @@ class DICE_Trainer:
     report:dict
     """
 
-    def __init__(self, args, is_continual_training, storage_path, evaluator=None):
+    def __init__(self, args, is_continual_training:bool, storage_path, evaluator=None):
         self.report = dict()
         self.args = args
         self.trainer = None
@@ -140,13 +145,13 @@ def __init__(self, args, is_continual_training, storage_path, evaluator=None):
         # Required for CV.
         self.evaluator = evaluator
         self.form_of_labelling = None
-        print(
-            f'# of CPUs:{os.cpu_count()} | # of GPUs:{torch.cuda.device_count()} | # of CPUs for dataloader:{self.args.num_core}')
-
+        print(f'# of CPUs:{os.cpu_count()} |'
+              f' # of GPUs:{torch.cuda.device_count()} |'
+              f' # of CPUs for dataloader:{self.args.num_core}')
         for i in range(torch.cuda.device_count()):
             print(torch.cuda.get_device_name(i))
 
-    def continual_start(self):
+    def continual_start(self,knowledge_graph):
         """
         (1) Initialize training.
         (2) Load model
@@ -164,18 +169,15 @@ def continual_start(self):
 
         self.trainer = self.initialize_trainer(callbacks=get_callbacks(self.args))
         model, form_of_labelling = self.initialize_or_load_model()
-        assert form_of_labelling in ['EntityPrediction', 'RelationPrediction', 'Pyke']
-        assert self.args.scoring_technique in ["AllvsAll",'KvsSample', '1vsAll', 'KvsAll', 'NegSample']
-        train_loader = self.initialize_dataloader(
-            reload_dataset(path=self.storage_path, form_of_labelling=form_of_labelling,
-                           scoring_technique=self.args.scoring_technique,
-                           neg_ratio=self.args.neg_ratio,
-                           label_smoothing_rate=self.args.label_smoothing_rate))
-        self.trainer.fit(model, train_dataloaders=train_loader)
+        # TODO: Here we need to load memory pag
+        self.trainer.evaluator = self.evaluator
+        self.trainer.dataset = knowledge_graph
+        self.trainer.form_of_labelling = form_of_labelling
+        self.trainer.fit(model, train_dataloaders=self.init_dataloader(self.init_dataset()))
         return model, form_of_labelling
 
     @timeit
-    def initialize_trainer(self, callbacks: List) -> pl.Trainer:
+    def initialize_trainer(self, callbacks: List) -> pl.Trainer | TensorParallel | TorchTrainer | TorchDDPTrainer:
         """ Initialize Trainer from input arguments """
         return initialize_trainer(self.args, callbacks)
 
@@ -188,7 +190,7 @@ def initialize_or_load_model(self):
         return model, form_of_labelling
 
     @timeit
-    def initialize_dataloader(self, dataset: torch.utils.data.Dataset) -> torch.utils.data.DataLoader:
+    def init_dataloader(self, dataset: torch.utils.data.Dataset) -> torch.utils.data.DataLoader:
         print('Initializing Dataloader...', end='\t')
         # https://pytorch.org/docs/stable/data.html#multi-process-data-loading
         # https://github.com/pytorch/pytorch/issues/13246#issuecomment-905703662
@@ -197,42 +199,92 @@ def initialize_dataloader(self, dataset: torch.utils.data.Dataset) -> torch.util
                                            num_workers=self.args.num_core, persistent_workers=False)
 
     @timeit
-    def initialize_dataset(self, dataset: KG, form_of_labelling) -> torch.utils.data.Dataset:
+    def init_dataset(self) -> torch.utils.data.Dataset:
         print('Initializing Dataset...', end='\t')
-        train_dataset = construct_dataset(train_set=dataset.train_set,
-                                          valid_set=dataset.valid_set,
-                                          test_set=dataset.test_set,
-                                          train_target_indices=dataset.train_target_indices,
-                                          target_dim=dataset.target_dim,
-                                          ordered_bpe_entities=dataset.ordered_bpe_entities,
-                                          entity_to_idx=dataset.entity_to_idx,
-                                          relation_to_idx=dataset.relation_to_idx,
-                                          form_of_labelling=form_of_labelling,
-                                          scoring_technique=self.args.scoring_technique,
-                                          neg_ratio=self.args.neg_ratio,
-                                          label_smoothing_rate=self.args.label_smoothing_rate,
-                                          byte_pair_encoding=self.args.byte_pair_encoding,
-                                          block_size=self.args.block_size)
-        if self.args.eval_model is None:
-            del dataset.train_set
-            gc.collect()
+        if isinstance(self.trainer.dataset,KG):
+            # Create a memory map of training dataset to reduce the memory usage
+            train_set_shape=self.trainer.dataset.train_set.shape
+            train_set_dtype=self.trainer.dataset.train_set.dtype
+            path_memory_map=self.trainer.dataset.path_for_serialization + '/memory_map_train_set.npy'
+            memmap_kg = np.memmap(path_memory_map, dtype=train_set_dtype, mode='w+', shape=train_set_shape)
+            memmap_kg[:] = self.trainer.dataset.train_set[:]
+            memmap_kg[:].flush()
+            del memmap_kg
+            self.trainer.dataset.train_se = np.memmap(path_memory_map,
+                                             mode='r',
+                                             dtype=train_set_dtype,
+                                             shape=train_set_shape)
+            train_dataset = construct_dataset(train_set=self.trainer.dataset.train_set,
+                                              valid_set=self.trainer.dataset.valid_set,
+                                              test_set=self.trainer.dataset.test_set,
+                                              train_target_indices=self.trainer.dataset.train_target_indices,
+                                              target_dim=self.trainer.dataset.target_dim,
+                                              ordered_bpe_entities=self.trainer.dataset.ordered_bpe_entities,
+                                              entity_to_idx=self.trainer.dataset.entity_to_idx,
+                                              relation_to_idx=self.trainer.dataset.relation_to_idx,
+                                              form_of_labelling=self.trainer.form_of_labelling,
+                                              scoring_technique=self.args.scoring_technique,
+                                              neg_ratio=self.args.neg_ratio,
+                                              label_smoothing_rate=self.args.label_smoothing_rate,
+                                              byte_pair_encoding=self.args.byte_pair_encoding,
+                                              block_size=self.args.block_size)
+        else:
+            assert isinstance(self.trainer.dataset, np.memmap), ("Train dataset must be an instance of memmap. "
+                                                                 f"Currently, {type(np.memmap)}!")
+            if self.args.continual_learning:
+                path = self.args.continual_learning
+            else:
+                path = self.args.path_to_store_single_run
+
+            train_dataset = construct_dataset(train_set=self.trainer.dataset,
+                                              valid_set=None,
+                                              test_set=None,
+                                              train_target_indices=None,
+                                              target_dim=None,
+                                              ordered_bpe_entities=None,
+                                              entity_to_idx={v["entity"]:k for k,v in pd.read_csv(f"{path}/entity_to_idx.csv",index_col=0).to_dict(orient='index').items()},
+                                              relation_to_idx={v["relation"]:k for k,v in pd.read_csv(f"{path}/relation_to_idx.csv",index_col=0).to_dict(orient='index').items()},
+                                              form_of_labelling=self.trainer.form_of_labelling,
+                                              scoring_technique=self.args.scoring_technique,
+                                              neg_ratio=self.args.neg_ratio,
+                                              label_smoothing_rate=self.args.label_smoothing_rate,
+                                              byte_pair_encoding=self.args.byte_pair_encoding,
+                                              block_size=self.args.block_size)
+
+
         return train_dataset
 
-    def start(self, knowledge_graph: KG) -> Tuple[BaseKGE, str]:
+    def start(self, knowledge_graph: Union[KG,np.memmap]) -> Tuple[BaseKGE, str]:
+        """
+        Start the training
+
+        (1) Initialize Trainer
+        (2) Initialize or load a pretrained KGE model
+
+        in DDP setup, we need to load the memory map of already read/index KG.
+        """
         """ Train selected model via the selected training strategy """
         print('------------------- Train -------------------')
-
+        assert isinstance(knowledge_graph, np.memmap) or isinstance(knowledge_graph, KG), \
+            f"knowledge_graph must be an instance of KG or np.memmap. Currently {type(knowledge_graph)}"
         if self.args.num_folds_for_cv == 0:
-            # Initialize Trainer
-            self.trainer: Union[TorchTrainer, TorchDDPTrainer, pl.Trainer]
+            self.trainer: Union[TensorParallel, TorchTrainer, TorchDDPTrainer, pl.Trainer]
             self.trainer = self.initialize_trainer(callbacks=get_callbacks(self.args))
-            # Initialize or load model
             model, form_of_labelling = self.initialize_or_load_model()
             self.trainer.evaluator = self.evaluator
             self.trainer.dataset = knowledge_graph
             self.trainer.form_of_labelling = form_of_labelling
-            self.trainer.fit(model, train_dataloaders=self.initialize_dataloader(
-                self.initialize_dataset(knowledge_graph, form_of_labelling)))
+            # TODO: Later, maybe we should write a callback to save the models in disk
+
+            if isinstance(self.trainer, TensorParallel):
+                assert isinstance(model, EnsembleKGE), type(model)
+
+                model = self.trainer.fit(model, train_dataloaders=self.init_dataloader(self.init_dataset()))
+                assert isinstance(model,EnsembleKGE)
+            else:
+                self.trainer.fit(model, train_dataloaders=self.init_dataloader(self.init_dataset()))
+
+
             return model, form_of_labelling
         else:
             return self.k_fold_cross_validation(knowledge_graph)
@@ -270,7 +322,7 @@ def k_fold_cross_validation(self, dataset) -> Tuple[BaseKGE, str]:
             train_set_for_i_th_fold, test_set_for_i_th_fold = dataset.train_set[train_index], dataset.train_set[
                 test_index]
 
-            trainer.fit(model, train_dataloaders=self.initialize_dataloader(
+            trainer.fit(model, train_dataloaders=self.init_dataloader(
                 construct_dataset(train_set=train_set_for_i_th_fold,
                                   entity_to_idx=dataset.entity_to_idx,
                                   relation_to_idx=dataset.relation_to_idx,
diff --git a/dicee/trainer/model_parallelism.py b/dicee/trainer/model_parallelism.py
new file mode 100644
index 00000000..a900bfd7
--- /dev/null
+++ b/dicee/trainer/model_parallelism.py
@@ -0,0 +1,299 @@
+import torch
+from ..abstracts import AbstractTrainer
+from ..static_funcs_training import make_iterable_verbose
+from ..models.ensemble import EnsembleKGE
+from typing import Tuple
+import time
+
+def extract_input_outputs(z: list, device=None):
+    # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
+    if len(z) == 2:
+        x_batch, y_batch = z
+        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
+        if device:
+            x_batch, y_batch = x_batch.to(device, non_blocking=True), y_batch.pin_memory().to(device,
+                                                                                              non_blocking=True)
+        return x_batch, y_batch
+    elif len(z) == 3:
+        x_batch, y_idx_batch, y_batch, = z
+        if device:
+            x_batch, y_batch, y_idx_batch = x_batch.pin_memory().to(device,
+                                                                    non_blocking=True), y_batch.pin_memory().to(
+                device, non_blocking=True), y_idx_batch.pin_memory().to(device, non_blocking=True)
+        return (x_batch, y_idx_batch), y_batch
+    else:
+        raise ValueError('Unexpected batch shape..')
+
+
+def find_good_batch_size(train_loader,tp_ensemble_model):
+    # () Initial batch size.
+    initial_batch_size=train_loader.batch_size
+    # () # of training data points.
+    training_dataset_size=len(train_loader.dataset)
+    # () Batch is large enough.
+    if initial_batch_size >= training_dataset_size:
+        return training_dataset_size, None
+    # () Log the number of training data points.
+    print("Number of training data points:",training_dataset_size)
+
+    def increase_batch_size_until_cuda_out_of_memory(ensemble_model, train_loader, batch_size,delta: int = None):
+        assert delta is not None, "delta cannot be None."
+        assert isinstance(delta, int), "delta must be a positive integer."
+        # () Store the batch sizes and GPU memory usages in a tuple.
+        batch_sizes_and_mem_usages = []
+        # () Increase the batch size until a stopping criterion is reached.
+        try:
+            while True:
+                start_time=time.time()
+                # () Initialize a dataloader with a current batch_size
+                train_dataloaders = torch.utils.data.DataLoader(train_loader.dataset,
+                                                                batch_size=batch_size,
+                                                                shuffle=True,
+                                                                sampler=None,
+                                                                batch_sampler=None,
+                                                                num_workers=train_loader.num_workers,
+                                                                collate_fn=train_loader.dataset.collate_fn,
+                                                                pin_memory=False,
+                                                                drop_last=False,
+                                                                timeout=0,
+                                                                worker_init_fn=None,
+                                                                persistent_workers=False)
+                
+                batch_loss = None
+                for i, batch_of_training_data in enumerate(train_dataloaders):
+                    batch_loss = forward_backward_update_loss(batch_of_training_data, ensemble_model)
+                    break
+                
+                global_free_memory, total_memory = torch.cuda.mem_get_info(device="cuda:0")
+                percentage_used_gpu_memory = (total_memory - global_free_memory) / total_memory
+                rt=time.time()-start_time
+
+                print(f"Random Batch Loss: {batch_loss:0.4}\tGPU Usage: {percentage_used_gpu_memory:0.3}\tRuntime: {rt:.3f}\tBatch Size: {batch_size}")
+
+                # Store the batch size and the runtime
+                batch_sizes_and_mem_usages.append((batch_size, rt))
+                
+                # ()
+                # https://github.com/pytorch/pytorch/issues/21819
+                # CD: as we reach close to 1.0 GPU memory usage, we observe RuntimeError: CUDA error: an illegal memory access was encountered.
+                # CD: To avoid this problem, we add the following condition as a temp solution.
+                if percentage_used_gpu_memory > 0.9:
+                    # Mimik out of memory error
+                    return batch_sizes_and_mem_usages, False
+                if batch_size < training_dataset_size:
+                    # Increase the batch size.
+                    batch_size += int(batch_size / delta)
+                else:
+                    return batch_sizes_and_mem_usages,True
+                        
+        except torch.OutOfMemoryError as e:
+            print(f"torch.OutOfMemoryError caught! {e}\n\n")
+            return batch_sizes_and_mem_usages, False
+
+    history_batch_sizes_and_mem_usages=[]
+    batch_size=initial_batch_size
+
+    for delta in range(1,5,1):
+        result,flag= increase_batch_size_until_cuda_out_of_memory(tp_ensemble_model, train_loader, batch_size,delta=delta)
+        
+        history_batch_sizes_and_mem_usages.extend(result)
+
+        if flag:
+            batch_size, batch_rt = history_batch_sizes_and_mem_usages[-1]
+        else:
+            assert len(history_batch_sizes_and_mem_usages)>2, "GPU memory errorin the first try"
+            # CUDA ERROR Observed 
+            batch_size, batch_rt=history_batch_sizes_and_mem_usages[-2]
+            # https://github.com/pytorch/pytorch/issues/21819
+            break
+
+        if batch_size>=training_dataset_size:
+            batch_size=training_dataset_size
+            break
+        else:
+            continue
+
+    return batch_size, batch_rt
+
+
+def forward_backward_update_loss(z:Tuple, ensemble_model)->float:
+    # () Get a random batch of data points (z).
+    x_batch, y_batch = extract_input_outputs(z)
+    # () Move the batch of labels into the master GPU : GPU-0.
+    y_batch = y_batch.to("cuda:0")
+    # () Forward pas on the batch of input data points (yhat on the master GPU).
+    yhat = ensemble_model(x_batch)
+    # () Compute the loss.
+    loss = torch.nn.functional.binary_cross_entropy_with_logits(yhat, y_batch)
+    # () Compute the gradient of the loss w.r.t. parameters.
+    loss.backward()
+    # () Parameter update.
+    ensemble_model.step()
+    return loss.item()
+
+class TensorParallel(AbstractTrainer):
+    def __init__(self, args, callbacks):
+        super().__init__(args, callbacks)
+
+
+    def fit(self, *args, **kwargs):
+        """ Train model        """
+        assert len(args) == 1
+        ensemble_model, = args
+        assert isinstance(ensemble_model,EnsembleKGE), (f"Selected model must "
+                                                        f"be an instance of EnsembleKGE{type(ensemble_model)}")
+        # () Run on_fit_start callbacks.
+        self.on_fit_start(self, ensemble_model)
+        # () Sanity checking
+        assert torch.cuda.device_count()== len(ensemble_model)
+        # () Get DataLoader
+        train_dataloader = kwargs['train_dataloaders']
+        # () Find a batch size so that available GPU memory is *almost* fully used.
+        if self.attributes.auto_batch_finding:
+            batch_size, batch_rt=find_good_batch_size(train_dataloader, ensemble_model)
+
+            train_dataloader = torch.utils.data.DataLoader(train_dataloader.dataset,
+                                                            batch_size=batch_size,
+                                                            shuffle=True,
+                                                            sampler=None,
+                                                            batch_sampler=None,
+                                                            num_workers=self.attributes.num_core,
+                                                            collate_fn=train_dataloader.dataset.collate_fn,
+                                                            pin_memory=False,
+                                                            drop_last=False,
+                                                            timeout=0,
+                                                            worker_init_fn=None,
+                                                            persistent_workers=False)
+            #if batch_rt is not None:
+            #    expected_training_time=batch_rt * len(train_dataloader) * self.attributes.num_epochs
+            # print(f"Exp.Training Runtime: {expected_training_time/60 :.3f} in mins\t|\tBatch Size:{batch_size}\t|\tBatch RT:{batch_rt:.3f}\t|\t # of batches:{len(train_dataloader)}\t|\t# of epochs:{self.attributes.num_epochs}")
+
+        # () Number of batches to reach a single epoch.
+        num_of_batches = len(train_dataloader)
+        # () Start training.
+        for epoch in (tqdm_bar := make_iterable_verbose(range(self.attributes.num_epochs),
+                                                        verbose=True, position=0, leave=True)):
+            # () Accumulate the batch losses.
+            epoch_loss = 0
+            # () Iterate over batches.
+            for i, z in enumerate(train_dataloader):
+                # () Forward, Loss, Backward, and Update on a given batch of data points.
+                batch_loss = forward_backward_update_loss(z,ensemble_model)
+                # () Accumulate the batch losses to compute the epoch loss.
+                epoch_loss += batch_loss
+                # if verbose=TRue, show info.
+                if hasattr(tqdm_bar, 'set_description_str'):
+                    tqdm_bar.set_description_str(f"Epoch:{epoch + 1}")
+                    if i > 0:
+                        tqdm_bar.set_postfix_str(
+                            f"batch={i} | {num_of_batches}, loss_step={batch_loss:.5f}, loss_epoch={epoch_loss / i:.5f}")
+                    else:
+                        tqdm_bar.set_postfix_str(f"loss_step={batch_loss:.5f}, loss_epoch={batch_loss:.5f}")
+            # Store the epoch loss
+            ensemble_model.loss_history.append(epoch_loss)
+        # Run on_fit_end callbacks after the training is done.
+        self.on_fit_end(self, ensemble_model)
+        # TODO: Later, maybe we should write a callback to save the models in disk
+        return ensemble_model
+    
+    """
+    
+    def batchwisefit(self, *args, **kwargs):
+        assert len(args) == 1
+        model, = args
+        # (1) Run the fit the start callback.
+        self.on_fit_start(self, model)
+        # (2) Setup DDP.
+        optimizer = model.configure_optimizers()
+        num_gpus = torch.cuda.device_count()
+        for epoch in (tqdm_bar := make_iterable_verbose(range(self.attributes.num_epochs),
+                                                        verbose=True, position=0, leave=True)):
+            epoch_loss = 0
+            num_of_batches = len(kwargs['train_dataloaders'])
+            for i, (x_batch, y_batch) in enumerate(kwargs['train_dataloaders']):
+                # Define a large batch into small batches
+                x_splits = torch.chunk(x_batch, num_gpus)
+                y_splits = torch.chunk(y_batch, num_gpus)
+
+                # Forward pass. We need to paralelize it
+                gpu_losses = []
+                for gpu_id, (x_split, y_split) in enumerate(zip(x_splits, y_splits)):
+                    y_split = y_split.to(f"cuda:{gpu_id}")
+                    h_emb, r_emb, t_emb = model.get_triple_representation(x_split)
+                    h_emb, r_emb, t_emb = h_emb.pin_memory().to(f"cuda:{gpu_id}",
+                                                                non_blocking=True), r_emb.pin_memory().to(f"cuda:{gpu_id}", non_blocking=True), t_emb.pin_memory().to(f"cuda:{gpu_id}", non_blocking=True)
+                    yhat = model.score(h_emb, r_emb, t_emb)
+                    gpu_losses.append(torch.nn.functional.binary_cross_entropy_with_logits(yhat, y_split).to("cuda:0"))
+
+                loss = sum(gpu_losses) / len(gpu_losses)
+
+                loss.backward()
+                batch_loss = loss.item()
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                epoch_loss += batch_loss
+
+                if hasattr(tqdm_bar, 'set_description_str'):
+                    tqdm_bar.set_description_str(f"Epoch:{epoch + 1}")
+                    if i > 0:
+                        tqdm_bar.set_postfix_str(
+                            f"batch={i} | {num_of_batches}, loss_step={batch_loss:.5f}, loss_epoch={epoch_loss / i:.5f}")
+                    else:
+                        tqdm_bar.set_postfix_str(f"loss_step={batch_loss:.5f}, loss_epoch={batch_loss:.5f}")
+
+    def torch_buggy_fit(self, *args, **kwargs):
+        assert len(args) == 1
+        model, = args
+        # () Run the fit the start callback.
+        self.on_fit_start(self, model)
+        # () Init Process Group with NCCL.
+        torch.distributed.init_process_group(backend="nccl")
+        # () Get Rank and World Size.
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+        # () Reinitialize Rank based on manuel seed rank.
+        torch.manual_seed(rank)
+        model.param_init(model.entity_embeddings.weight.data) 
+        model.param_init(model.relation_embeddings.weight.data)
+        # () .
+        device = torch.device(f'cuda:{rank}')
+        model.to(device)
+        # () .
+        optimizer = model.configure_optimizers()
+        # () .
+        for epoch in (tqdm_bar := make_iterable_verbose(range(self.attributes.num_epochs),
+                                                        verbose=True, position=0, leave=True)):
+            epoch_loss = 0
+            num_of_batches = len(kwargs['train_dataloaders'])
+            # () .
+            for i, z in enumerate(kwargs['train_dataloaders']):
+                optimizer.zero_grad()
+                # () Get batch and move it on GPUs .
+                inputs,targets = extract_input_outputs(z,device)
+                # () Predict .
+                yhats = model(inputs)   
+                # () TODO: Pytorch Bug https://github.com/pytorch/pytorch/issues/58005 .
+                dist.all_reduce(yhats,op=dist.ReduceOp.SUM)
+                # () Compute loss .
+                loss = torch.nn.functional.binary_cross_entropy_with_logits(yhats, targets)
+                # () Backward .
+                loss.backward()
+                # () .
+                batch_loss = loss.item()
+                # () .
+                optimizer.step()
+                # () .
+                epoch_loss +=batch_loss
+                # () .
+                if rank==0 and hasattr(tqdm_bar, 'set_description_str'):
+                    tqdm_bar.set_description_str(f"Epoch:{epoch + 1}")
+                    if i > 0:
+                        tqdm_bar.set_postfix_str(f"batch={i} | {num_of_batches}, loss_step={batch_loss:.5f}, loss_epoch={epoch_loss / i:.5f}")
+                    else:
+                        tqdm_bar.set_postfix_str(f"loss_step={batch_loss:.5f}, loss_epoch={batch_loss:.5f}")
+        # () .
+        torch.distributed.destroy_process_group()
+        # () .
+        self.on_fit_end(self, model)
+    """
diff --git a/dicee/trainer/torch_trainer.py b/dicee/trainer/torch_trainer.py
index a35948c8..6bda30da 100644
--- a/dicee/trainer/torch_trainer.py
+++ b/dicee/trainer/torch_trainer.py
@@ -4,7 +4,7 @@
 import time
 import os
 import psutil
-
+from tqdm import tqdm
 
 class TorchTrainer(AbstractTrainer):
     """
@@ -58,48 +58,6 @@ def _run_batch(self, i: int, x_batch, y_batch) -> float:
         # (3) Loss Forward and Backward w.r.t the batch.
         return self.forward_backward_update(x_batch, y_batch)
 
-    def _run_epoch(self, epoch: int) -> float:
-        """
-            Iterate over the training dataset
-
-            Arguments
-           ----------
-           epoch:int
-           -------
-           average loss over the dataset
-       """
-        epoch_loss = 0
-        i = 0
-        construct_mini_batch_time = None
-        batch: list
-        for i, batch in enumerate(self.train_dataloaders):
-            # (1) Extract Input and Outputs and set them on the dice
-            x_batch, y_batch = self.extract_input_outputs_set_device(batch)
-            start_time = time.time()
-            if construct_mini_batch_time:
-                construct_mini_batch_time = start_time - construct_mini_batch_time
-            # (2) Forward-Backward-Update.
-            batch_loss = self._run_batch(i, x_batch, y_batch)
-            epoch_loss += batch_loss
-            if construct_mini_batch_time:
-                print(
-                    f"Epoch:{epoch + 1} "
-                    f"| Batch:{i + 1} "
-                    f"| Loss:{batch_loss:.10f} "
-                    f"| ForwardBackwardUpdate:{(time.time() - start_time):.2f}sec "
-                    f"| BatchConst.:{construct_mini_batch_time:.2f}sec "
-                    f"| Mem. Usage {self.process.memory_info().rss / 1_000_000: .5}MB "
-                    f" ({psutil.virtual_memory().percent} %)")
-            else:
-                print(
-                    f"Epoch:{epoch + 1} "
-                    f"| Batch:{i + 1} "
-                    f"| Loss:{batch_loss} "
-                    f"| ForwardBackwardUpdate:{(time.time() - start_time):.2f}secs "
-                    f"| Mem. Usage {self.process.memory_info().rss / 1_000_000: .5}MB ")
-            construct_mini_batch_time = time.time()
-        return epoch_loss / (i + 1)
-
     def fit(self, *args, train_dataloaders, **kwargs) -> None:
         """
             Training starts
@@ -130,13 +88,27 @@ def fit(self, *args, train_dataloaders, **kwargs) -> None:
               f'| LearningRate:{self.model.learning_rate} '
               f'| BatchSize:{self.train_dataloaders.batch_size} '
               f'| EpochBatchsize:{len(train_dataloaders)}')
-        for epoch in range(self.attributes.max_epochs):
-            start_time = time.time()
 
-            avg_epoch_loss = self._run_epoch(epoch)
-            print(f"Epoch:{epoch + 1} "
-                  f"| Loss:{avg_epoch_loss:.8f} "
-                  f"| Runtime:{(time.time() - start_time) / 60:.3f} mins")
+        for epoch in (tqdm_bar := tqdm(range(self.attributes.max_epochs))):
+            epoch_loss = 0
+            i = 0
+            construct_mini_batch_time = None
+            batch: list
+            for i, batch in enumerate(self.train_dataloaders):
+                # (1) Extract Input and Outputs and set them on the dice
+                x_batch, y_batch = self.extract_input_outputs_set_device(batch)
+                start_time = time.time()
+                if construct_mini_batch_time:
+                    construct_mini_batch_time = start_time - construct_mini_batch_time
+                # (2) Forward-Backward-Update.
+                batch_loss = self._run_batch(i, x_batch, y_batch)
+                epoch_loss += batch_loss
+                tqdm_bar.set_description_str(f"Epoch:{epoch + 1}")
+                if i>0:
+                    tqdm_bar.set_postfix_str(f"loss_step={batch_loss:.5f}, loss_epoch={epoch_loss/i:.5f}")
+                else:
+                    tqdm_bar.set_postfix_str(f"loss_step={batch_loss:.5f}, loss_epoch={batch_loss:.5f}")
+            avg_epoch_loss = epoch_loss / len(self.train_dataloaders)
             """
             # Autobatch Finder: Double the current batch size if memory allows and repeat this process at mast 5 times.
             if self.attributes.auto_batch_finder and psutil.virtual_memory().percent < 30.0 and counter < 5:
@@ -205,4 +177,5 @@ def extract_input_outputs_set_device(self, batch: list) -> Tuple:
             return (x_batch, y_idx_batch), y_batch
         else:
             print(len(batch))
-            raise ValueError('Unexpected batch shape..')
+            print("Unexpected batch shape..")
+            raise RuntimeError
\ No newline at end of file
diff --git a/dicee/trainer/torch_trainer_ddp.py b/dicee/trainer/torch_trainer_ddp.py
index 368f7e97..61afbc34 100644
--- a/dicee/trainer/torch_trainer_ddp.py
+++ b/dicee/trainer/torch_trainer_ddp.py
@@ -1,17 +1,17 @@
 import os
 import torch
-import time
-from torch.nn.parallel import DistributedDataParallel as DDP
-
+from typing import Iterable
 from dicee.abstracts import AbstractTrainer
-from dicee.static_funcs_training import efficient_zero_grad
 from torch.utils.data import DataLoader
+from tqdm import tqdm
 
+torch.set_float32_matmul_precision('high')
 
-# DDP with gradiant accumulation https://gist.github.com/mcarilli/bf013d2d2f4b4dd21ade30c9b52d5e2e
-def print_peak_memory(prefix, device):
-    if device == 0:
-        print(f"{prefix}: {torch.cuda.max_memory_allocated(device) // 1e6}MB ")
+def make_iterable_verbose(iterable_object, verbose, desc="Default", position=None, leave=True) -> Iterable:
+    if verbose:
+        return tqdm(iterable_object, desc=desc, position=position, leave=leave)
+    else:
+        return iterable_object
 
 
 class TorchDDPTrainer(AbstractTrainer):
@@ -48,9 +48,7 @@ def fit(self, *args, **kwargs):
         model, = args
         # (1) Run the fit the start callback.
         self.on_fit_start(self, model)
-        print("DDP starts")
         # (2) Setup DDP.
-        print("NCCL is being initialized....")
         torch.distributed.init_process_group(backend="nccl")
         train_dataset_loader = kwargs['train_dataloaders']
         # (1) Create DATA LOADER.
@@ -63,11 +61,8 @@ def fit(self, *args, **kwargs):
                                           collate_fn=kwargs['train_dataloaders'].dataset.collate_fn,
                                           sampler=torch.utils.data.distributed.DistributedSampler(
                                               train_dataset_loader.dataset))
-
-        # (2) Initialize OPTIMIZER.
-        optimizer = model.configure_optimizers()
         # (3) Start NodeTrainer.
-        NodeTrainer(self, model, train_dataset_loader, optimizer, self.callbacks, self.attributes.num_epochs).train()
+        NodeTrainer(self, model, train_dataset_loader, self.callbacks, self.attributes.num_epochs).train()
         torch.distributed.destroy_process_group()
         self.on_fit_end(self, model)
 
@@ -77,39 +72,26 @@ def __init__(self,
                  trainer,
                  model: torch.nn.Module,
                  train_dataset_loader: DataLoader,
-                 optimizer: torch.optim.Optimizer,
                  callbacks,
                  num_epochs: int) -> None:
-        print("Initializing Node Trainer....")
         # (1) Trainer.
         self.trainer = trainer
         # (2) Local and Global Ranks.
         self.local_rank = int(os.environ["LOCAL_RANK"])
         self.global_rank = int(os.environ["RANK"])
+        self.optimizer = model.configure_optimizers()
         # (3) Send model to local trainer.
-        self.model = model.to(self.local_rank)
         self.train_dataset_loader = train_dataset_loader
-        self.loss_func = self.model.loss
-        self.optimizer = optimizer
+        self.loss_func = model.loss
         self.callbacks = callbacks
-        # (3) Wrap the model with DDP() along with GPU ID that model lives on.
-        print("Initializing device on",self.local_rank)
-        self.model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[self.local_rank],output_device=self.local_rank)
+        self.model = torch.compile(model).to(self.local_rank)
+        self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[self.local_rank])#, output_device=self.local_rank)
         self.num_epochs = num_epochs
-        print("HEREEE")
-        print_peak_memory("Max memory allocated after creating DDP local local_rank:", self.local_rank)
-        print(f'Global Rank {self.global_rank}\t Local Rank:{self.local_rank}')
-        print(self.model)
-        print(self.optimizer)
-        print(f'Global:{self.global_rank}'
-              f' | Local:{self.local_rank}'
-              f' | NumOfDataPoints:{len(self.train_dataset_loader.dataset)}'
-              f' | NumOfEpochs:{self.num_epochs}'
-              f' | LearningRate:{self.model.module.learning_rate}'
-              f' | BatchSize:{self.train_dataset_loader.batch_size}'
-              f' | EpochBatchsize:{len(self.train_dataset_loader)}')
-
         self.loss_history = []
+        # TODO: CD: This should be given as an input param
+        ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}["float16"]
+        self.ctx = torch.amp.autocast(device_type="cuda",dtype=ptdtype)
+        self.scaler = torch.amp.GradScaler("cuda",enabled=True)
 
     def _load_snapshot(self, snapshot_path):
         raise NotImplementedError
@@ -128,22 +110,27 @@ def _run_batch(self, source: torch.LongTensor, targets: torch.FloatTensor):
         batch loss
 
         """
-        self.optimizer.zero_grad()
-        output = self.model(source)
-        loss = self.loss_func(output, targets)
-        batch_loss = loss.item()
-        loss.backward()
-        self.optimizer.step()
+        with self.ctx:
+            output = self.model(source)
+            loss = self.loss_func(output, targets)
+            batch_loss = loss.item()
+        self.scaler.scale(loss).backward()
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        # flush the gradients as soon as we can, no need for this memory anymore
+        self.optimizer.zero_grad(set_to_none=True)
         return batch_loss
 
     def extract_input_outputs(self, z: list):
+        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
         if len(z) == 2:
             x_batch, y_batch = z
-            return x_batch.to(self.local_rank), y_batch.to(self.local_rank)
+            # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
+            x_batch, y_batch = x_batch.pin_memory().to(self.local_rank, non_blocking=True), y_batch.pin_memory().to(self.local_rank, non_blocking=True)
+            return x_batch, y_batch
         elif len(z) == 3:
             x_batch, y_idx_batch, y_batch, = z
-            x_batch, y_idx_batch, y_batch = x_batch.to(self.local_rank), y_idx_batch.to(self.local_rank), y_batch.to(
-                self.local_rank)
+            x_batch, y_batch,y_idx_batch = x_batch.pin_memory().to(self.local_rank, non_blocking=True), y_batch.pin_memory().to(self.local_rank, non_blocking=True),y_idx_batch.pin_memory().to(self.local_rank, non_blocking=True)
             return (x_batch, y_idx_batch), y_batch
         else:
             raise ValueError('Unexpected batch shape..')
@@ -164,33 +151,10 @@ def _run_epoch(self, epoch: int) -> float:
         self.train_dataset_loader.sampler.set_epoch(epoch)
         epoch_loss = 0
         i = 0
-        construct_mini_batch_time = None
         for i, z in enumerate(self.train_dataset_loader):
             source, targets = self.extract_input_outputs(z)
-            start_time = time.time()
-            if construct_mini_batch_time:
-                construct_mini_batch_time = start_time - construct_mini_batch_time
             batch_loss = self._run_batch(source, targets)
             epoch_loss += batch_loss
-            if True:  # self.local_rank == self.global_rank==0:
-                if construct_mini_batch_time:
-                    print(
-                        f"Global:{self.global_rank}"
-                        f" | Local:{self.local_rank}"
-                        f" | Epoch:{epoch + 1}"
-                        f" | Batch:{i + 1}"
-                        f" | Loss:{batch_loss}"
-                        f" | ForwardBackwardUpdate:{(time.time() - start_time):.2f}sec"
-                        f" | BatchConst.:{construct_mini_batch_time:.2f}sec")
-                else:
-                    print(
-                        f"Global:{self.global_rank}"
-                        f" | Local:{self.local_rank}"
-                        f" | Epoch:{epoch + 1}"
-                        f" | Batch:{i + 1}"
-                        f" | Loss:{batch_loss}"
-                        f" | ForwardBackwardUpdate:{(time.time() - start_time):.2f}secs")
-            construct_mini_batch_time = time.time()
         return epoch_loss / (i + 1)
 
     def train(self):
@@ -201,109 +165,27 @@ def train(self):
         -------
 
         """
-        for epoch in range(self.num_epochs):
-            start_time = time.time()
-            epoch_loss = self._run_epoch(epoch)
-
-            print(f"Global:{self.global_rank}"
-                  f" | Local:{self.local_rank}"
-                  f" | Epoch:{epoch + 1}"
-                  f" | Loss:{epoch_loss:.8f}"
-                  f" | Runtime:{(time.time() - start_time) / 60:.3f}mins")
-
-            if True:  # self.local_rank == self.global_rank == 0:
-                self.model.module.loss_history.append(epoch_loss)
+        num_of_batches=len(self.train_dataset_loader)
+        for epoch in (tqdm_bar := make_iterable_verbose(range(self.num_epochs),
+                                                      verbose=self.local_rank == self.global_rank == 0,
+                                                      position=0,
+                                                        leave=True)):
+            self.train_dataset_loader.sampler.set_epoch(epoch)
+            epoch_loss = 0
+            for i, z in enumerate(self.train_dataset_loader):
+                source, targets = self.extract_input_outputs(z)
+                batch_loss = self._run_batch(source, targets)
+                epoch_loss += batch_loss
+                if hasattr(tqdm_bar, 'set_description_str'):
+                    tqdm_bar.set_description_str(f"Epoch:{epoch + 1}")
+                    if i > 0:
+                        tqdm_bar.set_postfix_str(f"batch={i} | {num_of_batches}, loss_step={batch_loss:.5f}, loss_epoch={epoch_loss / i:.5f}")
+                    else:
+                        tqdm_bar.set_postfix_str(f"loss_step={batch_loss:.5f}, loss_epoch={batch_loss:.5f}")
+
+            avg_epoch_loss = epoch_loss / num_of_batches
+
+            if self.local_rank == self.global_rank == 0:
+                self.model.module.loss_history.append(avg_epoch_loss)
                 for c in self.callbacks:
                     c.on_train_epoch_end(self.trainer, self.model.module)
-
-
-class DDPTrainer:
-    def __init__(self,
-                 model: torch.nn.Module,
-                 train_dataset_loader: DataLoader,
-                 optimizer: torch.optim.Optimizer,
-                 gpu_id: int, callbacks, num_epochs) -> None:
-        self.gpu_id = gpu_id
-        self.model = model.to(gpu_id)
-        self.train_dataset_loader = train_dataset_loader
-        self.loss_func = self.model.loss
-        self.optimizer = optimizer
-        self.callbacks = callbacks
-        # (1) Wrap the model with DDP() along with GPU ID that model lives on.
-        self.model = DDP(model, device_ids=[gpu_id])
-        self.num_epochs = num_epochs
-        print_peak_memory("Max memory allocated after creating DDP:", gpu_id)
-        print('GPU:{self.gpu_id}')
-        print(self.model)
-        print(self.optimizer)
-        print(
-            f'NumOfDataPoints:{len(self.train_dataset_loader.dataset)}'
-            f'|NumOfEpochs:{self.num_epochs}'
-            f'|LearningRate:{self.model.module.learning_rate}'
-            f'|BatchSize:{self.train_dataset_loader.batch_size}'
-            f'|EpochBatchsize:{len(self.train_dataset_loader)}')
-
-        self.loss_history = []
-
-    def _run_batch(self, source, targets):
-        # (1) Zero the gradients.
-        # self.optimizer.zero_grad()
-        efficient_zero_grad(self.model)
-        output = self.model(source)
-        loss = self.loss_func(output, targets)
-        batch_loss = loss.item()
-        loss.backward()
-        self.optimizer.step()
-        # @TODO: Tips to decrease mem usage
-        #  https://github.com/pytorch/pytorch/issues/13246#issuecomment-905703662
-        #  torch.cuda.empty_cache()
-        return batch_loss
-
-    def extract_input_outputs(self, z: list):
-        if len(z) == 2:
-            x_batch, y_batch = z
-            return x_batch.to(self.gpu_id), y_batch.to(self.gpu_id)
-        elif len(z) == 3:
-            x_batch, y_idx_batch, y_batch, = z
-            x_batch, y_idx_batch, y_batch = x_batch.to(self.gpu_id), y_idx_batch.to(self.gpu_id), y_batch.to(
-                self.gpu_id)
-            return (x_batch, y_idx_batch), y_batch
-        else:
-            raise ValueError('Unexpected batch shape..')
-
-    def _run_epoch(self, epoch):
-        self.train_dataset_loader.sampler.set_epoch(epoch)
-        epoch_loss = 0
-        i = 0
-        construct_mini_batch_time = None
-        for i, z in enumerate(self.train_dataset_loader):
-            source, targets = self.extract_input_outputs(z)
-            start_time = time.time()
-            if construct_mini_batch_time:
-                construct_mini_batch_time = start_time - construct_mini_batch_time
-            batch_loss = self._run_batch(source, targets)
-            epoch_loss += batch_loss
-            if self.gpu_id == 0:
-                if construct_mini_batch_time:
-                    print(
-                        f"Epoch:{epoch + 1}|Batch:{i + 1}"
-                        f"|Loss:{batch_loss}"
-                        f"|ForwardBackwardUpdate:{(time.time() - start_time):.2f}sec"
-                        f"|BatchConst.:{construct_mini_batch_time:.2f}sec")
-                else:
-                    print(
-                        f"Epoch:{epoch + 1}|Batch:{i + 1}"
-                        f"|Loss:{batch_loss}"
-                        f"|ForwardBackwardUpdate:{(time.time() - start_time):.2f}secs")
-            construct_mini_batch_time = time.time()
-        return epoch_loss / (i + 1)
-
-    def train(self):
-        for epoch in range(self.num_epochs):
-            start_time = time.time()
-            epoch_loss = self._run_epoch(epoch)
-            if self.gpu_id == 0:
-                print(f"Epoch:{epoch + 1} | Loss:{epoch_loss:.8f} | Runtime:{(time.time() - start_time) / 60:.3f}mins")
-                self.model.module.loss_history.append(epoch_loss)
-                for c in self.callbacks:
-                    c.on_train_epoch_end(None, self.model.module)
diff --git a/docs/_static/images/dicee_logo.png b/docs/_static/images/dicee_logo.png
new file mode 100644
index 00000000..5df44b96
Binary files /dev/null and b/docs/_static/images/dicee_logo.png differ
diff --git a/docs/_static/images/favicon.ico b/docs/_static/images/favicon.ico
new file mode 100644
index 00000000..42e8969e
Binary files /dev/null and b/docs/_static/images/favicon.ico differ
diff --git a/docs/conf.py b/docs/conf.py
index 86bd1d69..7e29d0ce 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -67,6 +67,10 @@
 
 html_static_path = ['_static']
 
+html_logo = '_static/images/dicee_logo.png'
+
+html_favicon = '_static/images/favicon.ico'
+
 if stanford_theme_mod:
     html_theme = 'sphinx_rtd_theme'
 
diff --git a/docs/usage/main.md b/docs/usage/main.md
index 9f89c45d..c8af55bc 100644
--- a/docs/usage/main.md
+++ b/docs/usage/main.md
@@ -324,6 +324,52 @@ To test the Docker image:
 docker run --rm -v ~/.local/share/dicee/KGs:/dicee/KGs dice-embeddings ./main.py --model AConEx --embedding_dim 16
 ```
 
+## Coverage Report
+
+The coverage report is generated using [coverage.py](https://coverage.readthedocs.io/en/7.6.0/):
+
+```
+Name                                                   Stmts   Miss  Cover   Missing
+------------------------------------------------------------------------------------
+dicee/__init__.py                                          7      0   100%
+dicee/abstracts.py                                       201     82    59%   104-105, 123, 146-147, 152, 165, 197, 240-254, 257-260, 263-266, 301, 314-317, 320-324, 364-375, 390-398, 413, 424-428, 555-575, 581-585, 589-591
+dicee/callbacks.py                                       245    102    58%   50-55, 67-73, 76, 88-93, 98-103, 106-109, 116-133, 138-142, 146-147, 276-280, 286-287, 305-311, 314, 319-320, 332-338, 344-353, 358-360, 405, 416-429, 433-468, 480-486
+dicee/config.py                                           93      2    98%   141-142
+dicee/dataset_classes.py                                 299     74    75%   41, 54, 87, 93, 99-106, 109, 112, 115-139, 195-201, 204, 207-209, 314, 325-328, 344, 410-411, 429, 528-536, 539, 543-557, 700-707, 710-714
+dicee/eval_static_funcs.py                               227     95    58%   101, 106, 111, 258-353, 360-411
+dicee/evaluator.py                                       262     51    81%   46, 51, 56, 84, 89-90, 93, 109, 126, 137, 141, 146, 177-188, 195-206, 314, 344-367, 455, 465, 482-487
+dicee/executer.py                                        113      4    96%   116, 258-259, 291
+dicee/knowledge_graph.py                                  65      3    95%   79, 110, 114
+dicee/knowledge_graph_embeddings.py                      636    443    30%   27, 30-31, 39-52, 57-90, 93-127, 131-139, 170-184, 215-228, 254-274, 324-327, 330-333, 346, 381-426, 484-486, 502-503, 509-517, 522-525, 528-533, 538, 547, 592-598, 630, 688-1053, 1084-1145, 1149-1177, 1200, 1227-1265
+dicee/models/__init__.py                                   9      0   100%
+dicee/models/base_model.py                               234     31    87%   54, 56, 82, 88-103, 157, 190, 230, 236, 245, 248, 252, 259, 263, 265, 280, 288-289, 296-297, 351, 354, 427, 439
+dicee/models/clifford.py                                 556    357    36%   31-42, 68-117, 122-133, 156-168, 190-220, 235, 237, 241, 248-249, 276-280, 303-311, 325-327, 332-333, 364-384, 406, 413, 417-478, 495-499, 511, 514, 519, 524, 571-607, 625-631, 644, 647, 652, 657, 686-692, 705, 708, 713, 718, 728-737, 753-754, 774-845, 856-859, 884-909, 933-966, 1002-1006, 1019, 1029, 1032, 1037, 1042, 1047, 1051, 1055, 1064-1065, 1095, 1102, 1107, 1135-1139, 1167-1176, 1186-1194, 1212-1214, 1232-1234, 1250-1252
+dicee/models/complex.py                                  151     15    90%   86-109
+dicee/models/dualE.py                                     59     10    83%   93-102, 142-156
+dicee/models/function_space.py                           262    221    16%   10-24, 28-37, 40-49, 53-70, 77-86, 89-98, 101-110, 114-126, 134-156, 159-165, 168-185, 188-194, 197-205, 208, 213-234, 243-246, 250-254, 258-267, 271-292, 301-307, 311-328, 332-335, 344-352, 355, 366-372, 392-406, 424-438, 443-453, 461-465, 474-478
+dicee/models/octonion.py                                 227     83    63%   21-44, 320-329, 334-345, 348-370, 374-416, 426-474
+dicee/models/pykeen_models.py                             50      5    90%   60-63, 118
+dicee/models/quaternion.py                               192     69    64%   7-21, 30-55, 68-72, 107, 185, 328-342, 345-364, 368-389, 399-426
+dicee/models/real.py                                      61     12    80%   36-39, 66-69, 87, 103-106
+dicee/models/static_funcs.py                              10      0   100%
+dicee/models/transformers.py                             236    189    20%   24-43, 46, 60-75, 84-102, 105-116, 123-125, 128, 134-151, 155-180, 186-190, 193-197, 203-207, 210-212, 229-256, 265-268, 271-276, 279-304, 310-315, 319-372, 376-398, 404-414
+dicee/query_generator.py                                 374    346     7%   18-52, 56, 62-65, 69-70, 78-92, 100-147, 155-188, 192-206, 212-269, 274-303, 307-443, 453-472, 480-501, 508-512, 517, 522-528
+dicee/read_preprocess_save_load_kg/__init__.py             3      0   100%
+dicee/read_preprocess_save_load_kg/preprocess.py         256     41    84%   34, 40, 78, 102-127, 133, 138-151, 184, 214, 388-389, 444
+dicee/read_preprocess_save_load_kg/read_from_disk.py      36     11    69%   33, 38-40, 47, 55, 58-72
+dicee/read_preprocess_save_load_kg/save_load_disk.py      45     18    60%   39-60
+dicee/read_preprocess_save_load_kg/util.py               219    126    42%   65-67, 72-73, 91-97, 100-102, 107-109, 121, 134, 140-143, 148-156, 161-167, 172-177, 182-187, 199-220, 226-282, 286-290, 294-295, 299, 303-304, 334, 351, 356, 363-364
+dicee/sanity_checkers.py                                  54     23    57%   8-12, 21-31, 46, 51, 58, 64-79, 85, 89, 96
+dicee/static_funcs.py                                    418    163    61%   40, 50, 56-61, 83, 105-106, 115, 138, 152, 157-159, 163-165, 167, 194-198, 246, 254, 263-268, 290-304, 316-336, 340-357, 362, 386-387, 392-393, 410-411, 413-414, 416-417, 419-420, 428, 446-450, 467-470, 474-479, 483-487, 491-492, 498-500, 526-527, 539-542, 547-550, 559-610, 615-627, 644-658, 661-669
+dicee/static_funcs_training.py                           123     63    49%   118-215, 223-224
+dicee/static_preprocess_funcs.py                         100     44    56%   17-25, 52, 56, 64, 67, 78, 91-115, 120-123, 128-131, 136-139
+dicee/trainer/__init__.py                                  1      0   100%
+dicee/trainer/dice_trainer.py                            126     13    90%   27-32, 91, 98, 103-108, 147
+dicee/trainer/torch_trainer.py                            79      4    95%   31, 196, 207-208
+dicee/trainer/torch_trainer_ddp.py                       152    128    16%   13-14, 43, 47-72, 83-112, 131-137, 140-149, 164-194, 204-217, 226-246, 251-260, 263-272, 275-299, 302-309
+------------------------------------------------------------------------------------
+TOTAL                                                   6181   2828    54%
+```
 
 ## How to cite
 Currently, we are working on our manuscript describing our framework. 
diff --git a/examples/LP_BenchmarkResults.md b/examples/LP_BenchmarkResults.md
deleted file mode 100644
index 2db8188b..00000000
--- a/examples/LP_BenchmarkResults.md
+++ /dev/null
@@ -1,225 +0,0 @@
-# Link Prediction on Benchmark Datasets
-
-Here, we show that generalization performance of knowledge graph embedding models do not differ much if they are trained
-well.
-
-# Hyperparameter Setting
-
-Hyperparameters play an important role in the successful applications of knowledge graph embedding models. In our
-experiments, we selected such hyperparameter configuration so that experiments can be done less than a minute on UMLS
-dataset
-
-# Link Prediction Performance Analysis on KINSHIP
-
-| model_name   |   train_mrr |   train_h1 |   train_h3 |   train_h10 |   val_mrr |   val_h1 |   val_h3 |   val_h10 |   test_mrr |   test_h1 |   test_h3 |   test_h10 |   runtime |
-|:-------------|------------:|-----------:|-----------:|------------:|----------:|---------:|---------:|----------:|-----------:|----------:|----------:|-----------:|----------:|
-| DistMult     |    0.781619 |   0.681706 |   0.853406 |    0.967228 |  0.656955 | 0.516386 | 0.744382 |  0.941948 |   0.640159 |  0.496741 |  0.724395 |   0.938547 |   43.4192 |
-| ComplEx      |    0.883142 |    0.81291 |   0.947507 |     0.99011 |  0.757905 | 0.636236 | 0.860019 |   0.96161 |   0.762234 |  0.646182 |  0.851955 |   0.972533 |   45.6701 |
-| QMult        |    0.864761 |   0.789033 |   0.930185 |    0.985194 |  0.742049 |  0.61985 |  0.83661 |  0.955993 |    0.75211 |  0.636406 |  0.841713 |   0.956238 |   51.2721 |
-| OMult        |     0.85052 |   0.771126 |   0.917427 |    0.981449 |  0.744363 | 0.628277 | 0.834738 |   0.94897 |   0.745078 |  0.625233 |  0.838454 |   0.963687 |   47.2538 |
-| ConEx        |    0.726664 |   0.603406 |   0.815485 |    0.952891 |   0.68413 | 0.549157 | 0.779494 |  0.941479 |   0.676157 |  0.544693 |  0.768156 |   0.934358 |   52.2601 |
-| ConvQ        |    0.715911 |   0.586142 |   0.807584 |    0.957104 |  0.669007 | 0.525749 | 0.768258 |  0.948502 |   0.667853 |  0.524209 |  0.770019 |   0.935754 |   60.0243 |
-| ConvO        |    0.589155 |   0.433462 |   0.677434 |    0.918305 |  0.558573 | 0.401685 | 0.639513 |  0.902154 |    0.55491 |   0.39851 |  0.634078 |   0.902235 |   51.1115 |
-| AConvO       |    0.855045 |   0.770365 |   0.931355 |    0.986189 |   0.73474 | 0.609551 | 0.831461 |  0.950843 |   0.741717 |   0.61825 |  0.836127 |   0.962291 |   50.5469 |
-| AConEx       |    0.888418 |     0.8204 |    0.95055 |    0.989993 |   0.75704 | 0.637172 | 0.858614 |  0.955993 |   0.753787 |  0.634078 |  0.847765 |   0.967877 |   48.5488 |
-| AConvQ       |    0.888232 |   0.822156 |   0.947507 |    0.989876 |   0.74992 | 0.625468 | 0.852528 |  0.961142 |   0.752967 |  0.625698 |  0.856145 |   0.969274 |   46.3772 |
-
-# Link Prediction Performance Analysis on UMLS
-
-1. Multiplicative models fit the training dataset split of UMLS better than convolutional neural network based models.
-2. Replacing the multiplicative connections of conv(h,r) with additive connections leads convolutional neural network
-   based models to fit better the training data.
-3. Additive connections decrease the runtimes of the neural network based models.
-
-| model_name   |   train_mrr |   train_h1 |   train_h3 |   train_h10 |   val_mrr |   val_h1 |   val_h3 |   val_h10 |   test_mrr |   test_h1 |   test_h3 |   test_h10 |   runtime |
-|:-------------|------------:|-----------:|-----------:|------------:|----------:|---------:|---------:|----------:|-----------:|----------:|----------:|-----------:|----------:|
-| DistMult     |    0.997396 |   0.994919 |          1 |           1 |     0.739 | 0.622699 | 0.821319 |  0.954755 |   0.738951 |  0.620272 |  0.824508 |   0.947806 |   32.1166 |
-| ComplEx      |    0.998107 |   0.996357 |   0.999904 |           1 |  0.725287 | 0.604294 | 0.809816 |  0.950153 |   0.748761 |  0.633888 |   0.83056 |   0.958396 |   39.2883 |
-| QMult        |    0.997148 |   0.994824 |   0.999712 |           1 |  0.733427 | 0.611963 |  0.82592 |  0.958589 |   0.732836 |  0.611195 |  0.826778 |   0.960666 |   37.2082 |
-| OMult        |    0.985969 |   0.972776 |   0.999233 |           1 |  0.754292 | 0.629601 | 0.855828 |  0.959356 |   0.762545 |  0.636914 |  0.857035 |    0.96823 |   38.2151 |
-| ConEx        |    0.963993 |   0.932707 |   0.995399 |    0.999041 |   0.80534 | 0.705521 | 0.883436 |  0.957822 |   0.810177 |   0.70121 |  0.906203 |   0.968986 |   46.6404 |
-| ConvO        |    0.912093 |   0.837615 |    0.98495 |    0.997316 |  0.832247 | 0.721626 | 0.934049 |  0.977761 |   0.830195 |  0.714826 |  0.933434 |   0.984871 |   58.9782 |
-| ConvQ        |    0.904769 |   0.831959 |   0.975748 |    0.993194 |  0.834325 | 0.730828 | 0.927147 |  0.976227 |   0.827083 |  0.711044 |  0.939486 |   0.977307 |   42.0177 |
-| AConEx       |    0.995367 |   0.991181 |   0.999617 |           1 |  0.719798 |  0.58589 | 0.824387 |  0.943252 |   0.734349 |  0.610439 |  0.824508 |    0.95764 |   40.6111 |
-| AConvQ       |    0.995162 |   0.990702 |   0.999808 |           1 |  0.709874 | 0.568252 | 0.819018 |  0.956288 |   0.737209 |  0.607413 |  0.850983 |   0.953858 |   42.8576 |
-| AConvO       |    0.986693 |   0.974214 |   0.999041 |           1 |  0.698491 | 0.548313 | 0.817485 |  0.946319 |   0.732535 |  0.594554 |   0.83888 |   0.961422 |   44.2366 |
-
-# Link Prediction Performance Analysis on FB15K-237
-
-| model_name | train_mrr | train_h1 | train_h3 | train_h10 |  val_mrr |   val_h1 |   val_h3 |  val_h10 | test_mrr |  test_h1 |  test_h3 | test_h10 | runtime |
-|:-----------|----------:|---------:|---------:|----------:|---------:|---------:|---------:|---------:|---------:|---------:|---------:|---------:|--------:|
-| DistMult   |  0.593914 | 0.499502 | 0.647212 |  0.773351 |  0.26749 | 0.189136 | 0.290077 | 0.425463 | 0.264973 |   0.1867 |   0.2854 | 0.426024 | 1006.55 |
-| ComplEx    |  0.574006 | 0.476723 | 0.627397 |  0.760223 | 0.261119 | 0.182948 | 0.282606 |  0.41996 | 0.259313 | 0.180592 | 0.281809 | 0.421602 | 1031.37 |
-| QMult      |  0.549874 | 0.448279 | 0.605929 |  0.743963 | 0.266945 | 0.188822 | 0.290248 | 0.423667 | 0.265085 | 0.186114 | 0.288161 | 0.424265 | 1038.14 |
-| OMult      |  0.497422 | 0.389729 |  0.55521 |   0.70461 | 0.276785 | 0.197149 | 0.302509 | 0.435329 |  0.27315 | 0.193809 |   0.2972 | 0.432791 | 1098.29 |
-| ConEx      |   0.16171 |  0.10981 | 0.171578 |  0.260664 | 0.185139 | 0.130026 | 0.197177 | 0.292159 | 0.184531 | 0.128286 | 0.197743 | 0.293633 | 1096.44 |
-| AConEx     |  0.559767 | 0.458587 | 0.615846 |   0.75322 | 0.271471 | 0.192273 | 0.294953 | 0.430539 | 0.270874 | 0.191977 | 0.293975 | 0.431227 | 1106.65 |
-| AConvQ     |   0.49182 | 0.383151 | 0.546962 |  0.702877 | 0.263918 | 0.184631 | 0.285458 | 0.424636 | 0.263824 | 0.184306 | 0.286988 | 0.424411 | 1118.75 |
-| AConvO     |  0.481607 | 0.374538 | 0.535467 |  0.690759 | 0.275071 | 0.197092 | 0.296664 | 0.435672 |  0.27353 | 0.194225 | 0.298568 | 0.433744 | 1203.71 |
-
-
-## Rerun Experiments
-
-**Datasets: KGs/UMLS, KGs/KINSHIP, and KGs/FB15k-237**
-## ConEx (Convolutional Complex Knowledge Graph Embeddings)
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "ConEx" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --kernel_size 3 --num_of_output_channels 3 --seed_for_computation 0
-```
-```bash
-Total computation time: 38.546 seconds
-Evaluate ConEx on Train set: Evaluate ConEx on Train set
-{'H@1': 0.932707055214724, 'H@3': 0.995398773006135, 'H@10': 0.9990414110429447, 'MRR': 0.9639930469284392}
-Evaluate ConEx on Validation set: Evaluate ConEx on Validation set
-{'H@1': 0.7055214723926381, 'H@3': 0.8834355828220859, 'H@10': 0.9578220858895705, 'MRR': 0.8053400891316801}
-Evaluate ConEx on Test set: Evaluate ConEx on Test set
-{'H@1': 0.7012102874432677, 'H@3': 0.9062027231467473, 'H@10': 0.9689863842662633, 'MRR': 0.8101772584084878}
-```
-
-## AConEx (ConEx with Additive Connections)
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "AConEx" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --kernel_size 3 --num_of_output_channels 3 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 38.555 seconds
-Evaluate AConEx on Train set: Evaluate AConEx on Train set
-{'H@1': 0.991180981595092, 'H@3': 0.9996165644171779, 'H@10': 1.0, 'MRR': 0.9953668200408997}
-Evaluate AConEx on Validation set: Evaluate AConEx on Validation set
-{'H@1': 0.5858895705521472, 'H@3': 0.8243865030674846, 'H@10': 0.9432515337423313, 'MRR': 0.7197978156406648}
-Evaluate AConEx on Test set: Evaluate AConEx on Test set
-{'H@1': 0.6104387291981845, 'H@3': 0.8245083207261724, 'H@10': 0.9576399394856279, 'MRR': 0.734348890030545}
-```
-
-## ConvQ
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "ConvQ" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --kernel_size 3 --num_of_output_channels 3 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 41.013 seconds
-Evaluate ConvQ on Train set: Evaluate ConvQ on Train set
-{'H@1': 0.8319593558282209, 'H@3': 0.9757476993865031, 'H@10': 0.993194018404908, 'MRR': 0.9047693553188044}
-Evaluate ConvQ on Validation set: Evaluate ConvQ on Validation set
-{'H@1': 0.7308282208588958, 'H@3': 0.9271472392638037, 'H@10': 0.9762269938650306, 'MRR': 0.834325039087529}
-Evaluate ConvQ on Test set: Evaluate ConvQ on Test set
-{'H@1': 0.7110438729198184, 'H@3': 0.9394856278366112, 'H@10': 0.9773071104387292, 'MRR': 0.8270827146655849}
-```
-
-## AConvQ (ConvQ with with Additive Connections)
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "AConvQ" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --kernel_size 3 --num_of_output_channels 3 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 41.238 seconds
-Evaluate AConvQ on Train set: Evaluate AConvQ on Train set
-{'H@1': 0.9907016871165644, 'H@3': 0.999808282208589, 'H@10': 1.0, 'MRR': 0.9951623210633946}
-Evaluate AConvQ on Validation set: Evaluate AConvQ on Validation set
-{'H@1': 0.5682515337423313, 'H@3': 0.8190184049079755, 'H@10': 0.9562883435582822, 'MRR': 0.7098740236669477}
-Evaluate AConvQ on Test set: Evaluate AConvQ on Test set
-{'H@1': 0.6074130105900152, 'H@3': 0.850983358547655, 'H@10': 0.953857791225416, 'MRR': 0.737209298931209}
-```
-
-## ConvO
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "ConvO" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --kernel_size 3 --num_of_output_channels 3 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 54.756 seconds
-Evaluate ConvO on Train set: Evaluate ConvO on Train set
-{'H@1': 0.8376150306748467, 'H@3': 0.9849501533742331, 'H@10': 0.9973159509202454, 'MRR': 0.9120930363683276}
-Evaluate ConvO on Validation set: Evaluate ConvO on Validation set
-{'H@1': 0.7216257668711656, 'H@3': 0.9340490797546013, 'H@10': 0.977760736196319, 'MRR': 0.8322465353766179}
-Evaluate ConvO on Test set: Evaluate ConvO on Test set
-{'H@1': 0.7148260211800302, 'H@3': 0.9334341906202723, 'H@10': 0.9848714069591528, 'MRR': 0.8301952668172322}
-
-```
-
-## AConvO
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "AConvO" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --kernel_size 3 --num_of_output_channels 3 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 42.910 seconds
-Evaluate AConvO on Train set: Evaluate AConvO on Train set
-{'H@1': 0.9742139570552147, 'H@3': 0.9990414110429447, 'H@10': 1.0, 'MRR': 0.986693187627812}
-Evaluate AConvO on Validation set: Evaluate AConvO on Validation set
-{'H@1': 0.5483128834355828, 'H@3': 0.8174846625766872, 'H@10': 0.946319018404908, 'MRR': 0.6984911756060336}
-Evaluate AConvO on Test set: Evaluate AConvO on Test set
-{'H@1': 0.594553706505295, 'H@3': 0.8388804841149773, 'H@10': 0.9614220877458396, 'MRR': 0.7325348232341585}
-```
-
-## DistMult
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "DistMult" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0  --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 43.789 seconds
-Evaluate DistMult on Train set: Evaluate DistMult on Train set
-{'H@1': 0.9949194785276073, 'H@3': 1.0, 'H@10': 1.0, 'MRR': 0.9973958333333333}
-Evaluate DistMult on Validation set: Evaluate DistMult on Validation set
-{'H@1': 0.6226993865030674, 'H@3': 0.821319018404908, 'H@10': 0.9547546012269938, 'MRR': 0.7389999406726034}
-Evaluate DistMult on Test set: Evaluate DistMult on Test set
-{'H@1': 0.6202723146747352, 'H@3': 0.8245083207261724, 'H@10': 0.9478063540090772, 'MRR': 0.7389508344813291}
-```
-
-## ComplEx
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "ComplEx" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0 --feature_map_dropout_rate 0.0 --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 37.448 seconds
-Evaluate ComplEx on Train set: Evaluate ComplEx on Train set
-{'H@1': 0.9963573619631901, 'H@3': 0.9999041411042945, 'H@10': 1.0, 'MRR': 0.9981067868098159}
-Evaluate ComplEx on Validation set: Evaluate ComplEx on Validation set
-{'H@1': 0.6042944785276073, 'H@3': 0.8098159509202454, 'H@10': 0.9501533742331288, 'MRR': 0.725286550123257}
-Evaluate ComplEx on Test set: Evaluate ComplEx on Test set
-{'H@1': 0.6338880484114977, 'H@3': 0.8305597579425114, 'H@10': 0.9583963691376702, 'MRR': 0.7487605914404327}
-
-```
-
-## QMult
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "QMult" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0  --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 35.082 seconds
-Evaluate QMult on Train set: Evaluate QMult on Train set
-{'H@1': 0.9948236196319018, 'H@3': 0.9997124233128835, 'H@10': 1.0, 'MRR': 0.9971481978527608}
-Evaluate QMult on Validation set: Evaluate QMult on Validation set
-{'H@1': 0.6119631901840491, 'H@3': 0.825920245398773, 'H@10': 0.9585889570552147, 'MRR': 0.7334272790319509}
-Evaluate QMult on Test set: Evaluate QMult on Test set
-{'H@1': 0.6111951588502269, 'H@3': 0.8267776096822995, 'H@10': 0.9606656580937972, 'MRR': 0.7328362412430357}
-
-```
-
-## OMult
-
-```bash
-python main.py --path_dataset_folder KGs/UMLS --model "OMult" --optim Adam --embedding_dim 32 --num_epochs 256 --batch_size 1024 --lr 0.1 --backend 'pandas' --trainer 'PL' --scoring_technique 'KvsAll' --weight_decay 0.0 --input_dropout_rate 0.0 --hidden_dropout_rate 0.0  --normalization LayerNorm --init_param xavier_normal --label_smoothing_rate 0.0 --seed_for_computation 0 --num_core 4
-```
-
-```bash
-Total computation time: 52.970 seconds
-Evaluate OMult on Train set: Evaluate OMult on Train set
-{'H@1': 0.9727760736196319, 'H@3': 0.9992331288343558, 'H@10': 1.0, 'MRR': 0.985969452965235}
-Evaluate OMult on Validation set: Evaluate OMult on Validation set
-{'H@1': 0.629601226993865, 'H@3': 0.8558282208588958, 'H@10': 0.9593558282208589, 'MRR': 0.754292154154622}
-Evaluate OMult on Test set: Evaluate OMult on Test set
-{'H@1': 0.6369137670196672, 'H@3': 0.857034795763994, 'H@10': 0.9682299546142209, 'MRR': 0.7625450316560369}
-```
-
diff --git a/examples/Pykeen.md b/examples/Pykeen.md
index 42a70fab..ebe99e21 100644
--- a/examples/Pykeen.md
+++ b/examples/Pykeen.md
@@ -1,18 +1,29 @@
 # Using Pykeen
 
-# Training a Pykeen Model with Pytorch-lightning
+# Training Few Pykeen Models with KvsAll
+
+```
+dicee --dataset_dir KGs/UMLS --model Pykeen_MuRE --num_epochs 100 --batch_size 256 --lr 0.1 --scoring_technique KvsAll
+dicee --dataset_dir KGs/UMLS --model Pykeen_HolE --num_epochs 100 --batch_size 256 --lr 0.1 --scoring_technique KvsAll
+dicee --dataset_dir KGs/UMLS --model Pykeen_DistMult --num_epochs 100 --batch_size 256 --lr 0.1 --scoring_technique KvsAll
+dicee --dataset_dir KGs/UMLS --model Pykeen_ComplEx --num_epochs 100 --batch_size 256 --lr 0.1 --scoring_technique KvsAll
+dicee --dataset_dir KGs/UMLS --model Pykeen_QuatE --num_epochs 100 --batch_size 256 --lr 0.1 --scoring_technique KvsAll
+```
+
 ```
-python main.py --model Pykeen_MuRE --num_epochs 10 --batch_size 256 --lr 0.1 --trainer "PL" --num_core 4 --scoring_technique KvsAll --pykeen_model_kwargs embedding_dim=64
-python main.py --model Pykeen_HolE --num_epochs 10 --batch_size 256 --lr 0.1 --trainer "PL" --num_core 4 --scoring_technique KvsAll --pykeen_model_kwargs embedding_dim=64
-python main.py --model Pykeen_DistMult --num_epochs 10 --batch_size 256 --lr 0.1 --trainer "PL" --num_core 4 --scoring_technique KvsAll --pykeen_model_kwargs embedding_dim=64
-python main.py --model Pykeen_ComplEx --num_epochs 10 --batch_size 256 --lr 0.1 --trainer "PL" --num_core 4 --scoring_technique KvsAll --pykeen_model_kwargs embedding_dim=32
-python main.py --model Pykeen_QuatE --num_epochs 10 --batch_size 256 --lr 0.1 --trainer "PL" --num_core 4 --scoring_technique KvsAll --pykeen_model_kwargs embedding_dim=16
-python analyse_experiments.py
-| model_name      |   train_mrr |   train_h1 |   train_h3 |   train_h10 |   val_mrr |   val_h1 |   val_h3 |   val_h10 |   test_mrr |   test_h1 |   test_h3 |   test_h10 |   runtime |   params |
-|:----------------|------------:|-----------:|-----------:|------------:|----------:|---------:|---------:|----------:|-----------:|----------:|----------:|-----------:|----------:|---------:|
-| Pykeen_DistMult |    0.827423 |   0.707918 |   0.938267 |    0.992331 |  0.794849 | 0.661043 | 0.918712 |  0.980061 |   0.776859 |  0.639183 |  0.901664 |   0.982602 |   3.9479  |    14528 |
-| Pykeen_MuRE     |    0.96874  |   0.94306  |   0.995878 |    1        |  0.79331  | 0.685583 | 0.882669 |  0.969325 |   0.804933 |  0.705749 |  0.887292 |   0.967474 |   4.89179 |    20686 |
-| Pykeen_ComplEx  |    0.771764 |   0.673792 |   0.846817 |    0.925422 |  0.709708 | 0.598926 | 0.784509 |  0.895706 |   0.701002 |  0.58472  |  0.7882   |   0.8941   |   3.13179 |    14528 |
-| Pykeen_HolE     |    0.945462 |   0.901169 |   0.990798 |    0.999425 |  0.778686 | 0.667945 | 0.868098 |  0.960123 |   0.760102 |  0.641452 |  0.854766 |   0.962935 |   5.46916 |    14528 |
-| Pykeen_QuatE    |    0.92043  |   0.870207 |   0.96434  |    0.988976 |  0.817052 | 0.72546  | 0.898006 |  0.95092  |   0.819355 |  0.726929 |  0.894856 |   0.955371 |   4.24787 |    14528 |
+python dicee/analyse_experiments.py --dir Experiments --features "model" "trainMRR" "testMRR"
+
+\begin{tabular}{lrrr}
+\toprule
+model & trainMRR & testMRR & NumParam \\
+\midrule
+Pykeen_MuRE & 0.879 & 0.836 & 10478 \\
+Pykeen_HolE & 0.830 & 0.689 & 7264 \\
+Pykeen_QuatE & 1.000 & 0.683 & 29056 \\
+Pykeen_ComplEx & 1.000 & 0.648 & 14528 \\
+Pykeen_DistMult & 0.818 & 0.588 & 7264 \\
+\bottomrule
+\end{tabular}
+
+
 ```
\ No newline at end of file
diff --git a/examples/continual_training.py b/examples/continual_training.py
deleted file mode 100644
index 88038333..00000000
--- a/examples/continual_training.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from dicee.executer import ContinuousExecute
-import argparse
-def argparse_default(description=None):
-    parser = argparse.ArgumentParser(add_help=False)
-    # Dataset and storage related
-    parser.add_argument("--path_experiment_folder", type=str, default="Experiments/2023-01-07 18:44:47.703307",
-                        help="The path of a folder containing pretrained model")
-    parser.add_argument("--num_epochs", type=int, default=1, help='Number of epochs for training.')
-    parser.add_argument("--lr", type=float, default=None)
-    parser.add_argument("--num_core", type=int, default=None, help='Number of cores to be used.')
-    parser.add_argument('--scoring_technique', default=None, help="KvsSample, 1vsAll, KvsAll, NegSample")
-    parser.add_argument('--neg_ratio', type=int, default=None,
-                        help='The number of negative triples generated per positive triple.')
-    parser.add_argument('--optim', type=str, default=None, help='[NAdam, Adam, SGD]')
-    parser.add_argument('--batch_size', type=int, default=None, help='Mini batch size')
-    parser.add_argument("--seed_for_computation", type=int, default=0, help='Seed for all, see pl seed_everything().')
-    parser.add_argument("--trainer", type=str, default=None,
-                        help='PL (pytorch lightning trainer), torchDDP (custom ddp), torchCPUTrainer (custom cpu only)')
-    if description is None:
-        return parser.parse_args()
-    return parser.parse_args(description)
-
-
-if __name__ == '__main__':
-    args = argparse_default()
-    ContinuousExecute(args).continual_start()
diff --git a/examples/download_pretrained_model.py b/examples/download_pretrained_model.py
index a6dbc9e5..44bce06b 100644
--- a/examples/download_pretrained_model.py
+++ b/examples/download_pretrained_model.py
@@ -1,4 +1,9 @@
-# pip install dicee
+"""
+Download a pretrained KGE model
+Evaluate its performance
+python examples/download_pretrained_model.py
+"""
+
 from dicee import KGE
 import pandas as pd
 from dicee.static_funcs import get_er_vocab
diff --git a/examples/large_scale_training.py b/examples/large_scale_training.py
deleted file mode 100644
index 7b1dd9a4..00000000
--- a/examples/large_scale_training.py
+++ /dev/null
@@ -1,321 +0,0 @@
-"""
-Example script to train KGE models via model parallel with GPU-off-loading
-"""
-from typing import Tuple
-import os
-import torch
-import dicee
-from dicee import Keci
-from dicee import NegSampleDataset
-import polars as pl
-import time
-import pandas as pd
-import numpy as np
-from tqdm import tqdm
-import argparse
-from argparse import ArgumentParser
-import pickle
-
-def input_arguments():
-    parser = ArgumentParser()
-    parser.add_argument("--path_kg", type=str, default="dbpedia-2022-12-nt-wo-lit-polars.parquet.snappy",
-                        help="path parquet formatted polars dataframe")
-    parser.add_argument("--path_idx_kg", type=str, default="data.npy",
-                        help="path to numpy ndarray")
-    parser.add_argument("--path_checkpoint", type=str, default="Keci_1_14.torch"
-                        )
-    parser.add_argument("--path_checkpoint2", type=str, default="Keci_1_14.torch")
-
-    parser.add_argument("--batch_size", type=int, default=10_000_000)
-    parser.add_argument("--neg_sample_ratio", type=float, default=1.0)
-    parser.add_argument("--embedding_dim", type=int, default=20)
-    parser.add_argument("--num_epochs", type=int, default=1)
-    parser.add_argument("--read_only", default=None)
-    parser.add_argument("--lr", type=float, default=0.1)
-
-    return parser.parse_args()
-
-class MultiEpochsDataLoader(torch.utils.data.DataLoader):
-    """ To avoid the excessive time spent to fetch the first batch at each new epoch
-    See https://discuss.pytorch.org/t/enumerate-dataloader-slow/87778/2
-    """
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._DataLoader__initialized = False
-        self.batch_sampler = _RepeatSampler(self.batch_sampler)
-        self._DataLoader__initialized = True
-        self.iterator = super().__iter__()
-
-    def __len__(self):
-        return len(self.batch_sampler.sampler)
-
-    def __iter__(self):
-        for i in range(len(self)):
-            yield next(self.iterator)
-
-
-class _RepeatSampler(object):
-    """ Sampler that repeats forever.
-
-    Args:
-        sampler (Sampler)
-    """
-
-    def __init__(self, sampler):
-        self.sampler = sampler
-
-    def __iter__(self):
-        while True:
-            yield from iter(self.sampler)
-
-def get_data(args) -> Tuple[np.ndarray, int, int]:
-    if args.path_kg:
-        """ do this """
-        print("Reading KG...\n")
-        start_time = time.time()
-        data = pl.read_parquet(source=args.path_kg,n_rows=args.read_only)
-        print(f"took {time.time() - start_time}")
-        print("Unique entities...")
-        start_time = time.time()
-        unique_entities = pl.concat((data.get_column('subject'), data.get_column('object'))).unique().rename(
-            'entity')
-        print(f"Number of unique entities:{len(unique_entities)}")
-        unique_entities = unique_entities.to_list()
-        print(f"took {time.time() - start_time}")
-
-        print("Unique relations...")
-        start_time = time.time()
-        unique_relations = data.unique(subset=["relation"]).select("relation").to_series()
-        print(f"Number of unique relations:{len(unique_relations)}")
-        unique_relations = unique_relations.to_list()
-        print(f"took {time.time() - start_time}")
-
-        print("Entity index mapping...")
-        start_time = time.time()
-        entity_to_idx = {ent: idx for idx, ent in enumerate(unique_entities)}
-        pickle.dump(entity_to_idx, open("entity_to_idx.p", "wb"))
-
-        print(f"took {time.time() - start_time}")
-
-        print("Relation index mapping...")
-        start_time = time.time()
-        rel_to_idx = {rel: idx for idx, rel in enumerate(unique_relations)}
-        pickle.dump(rel_to_idx, open("relation_to_idx.p","wb")) 
-        print(f"took {time.time() - start_time}")
-        print("Constructing training data...")
-        start_time = time.time()
-        data = data.with_columns(pl.col("subject").map_dict(entity_to_idx).alias("subject"),
-                                 pl.col("relation").map_dict(rel_to_idx).alias("relation"),
-                                 pl.col("object").map_dict(entity_to_idx).alias("object")).to_numpy()
-        print(f"took {time.time() - start_time}")
-
-        num_entities = len(unique_entities)
-        num_relations = len(unique_relations)
-        # TODO: maybe save the data into some designated folter
-        with open("data.npy", 'wb') as f:
-            np.save(f, data)
-
-        return data, num_entities, num_relations
-
-    elif args.path_idx_kg:
-        print("Loading the index numpy KG..\n")
-        #data=np.load(args.path_idx_kg, mmap_mode='r')
-        with open(args.path_idx_kg, 'rb') as f:
-            data = np.load(f)
-        num_entities = 1 + max(max(data[:, 0]), max(data[:, 2]))
-        num_relations = 1 + max(data[:, 1])
-        return data, num_entities, num_relations
-    else:
-        raise RuntimeError
-
-
-def init_model(args, num_entities, num_relations):
-    start_time = time.time()
-    print('Initializing models...')
-    model1 = Keci(
-        args={"optim": "Adam", "p": 0, "q": 1, "num_entities": num_entities, "num_relations": num_relations,
-              "embedding_dim": args.embedding_dim, 'learning_rate': args.lr})
-    model2 = Keci(
-        args={"optim": "Adam", "p": 0, "q": 1, "num_entities": num_entities, "num_relations": num_relations,
-            "embedding_dim": args.embedding_dim, 'learning_rate': args.lr})
-    print(f"took {time.time() - start_time}")
-    return (model1, model2), (model1.configure_optimizers(), model2.configure_optimizers())
-
-
-def get_model(args, num_entities: int, num_relations: int):
-    # Initialize |GPUs| models on a single node
-    models, optimizers = init_model(args, num_entities, num_relations)
-    if args.path_checkpoint:
-        """ Load the checkpoint"""
-        # update models
-        model1, model2 = models
-        opt1, opt2 = optimizers
-
-        model1.load_state_dict(torch.load(args.path_checkpoint,map_location='cpu'))
-        model2.load_state_dict(torch.load(args.path_checkpoint,map_location='cpu'))
-        models = (model1, model2)
-        optimizers = (opt1, opt2)
-    return models, optimizers
-
-
-def get_train_loader(args):
-    data: np.ndarray
-    data, num_ent, num_rel = get_data(args)
-    data: torch.utils.data.DataLoader
-    print('Creating dataset...')
-    data: NegSampleDataset
-    # TODO: neg_sample_ratio is not used at the moment
-    data = NegSampleDataset(train_set=data,
-                            num_entities=num_ent, num_relations=num_rel,
-                            neg_sample_ratio=1.0)
-    data: torch.utils.data.DataLoader
-    data = MultiEpochsDataLoader(dataset=data,
-            batch_size=args.batch_size, shuffle=True,
-                                       num_workers=32)
-    print('Number of triples', len(data.dataset))
-    return data, num_ent, num_rel
-
-def run_epoch(loss_function,dataloader,model1,model2,opt1,opt2):
-    device1 = "cuda:0"
-    device2 = "cuda:1"
-    epoch_loss = 0.0
-    for ith, (x, y) in enumerate(tqdm(dataloader)):
-        # (1) Shape the batch
-        x = x.flatten(start_dim=0, end_dim=1)
-        y = y.flatten(start_dim=0, end_dim=1)
-            
-        # (2) Empty the gradients
-        opt1.zero_grad(set_to_none=True)
-        opt2.zero_grad(set_to_none=True)
-            
-        # (3) Forward Backward and Parameter Update
-        start_time = time.time()
-        # (3.1) Select embeddings of triples
-        h1, r1, t1 = model1.get_triple_representation(x)
-        # (3.2) Move (3.1) into a single GPU
-        h1, r1, t1, y = h1.pin_memory().to(device1, non_blocking=True), r1.pin_memory().to(device1,non_blocking=True), t1.pin_memory().to(device1, non_blocking=True), y.pin_memory().to(device1, non_blocking=True)
-        # (3.3) Compute triple score (Forward Pass)
-        yhat1 = model1.score(h1, r1, t1)
-
-        # (3.4) Select second part of the embeddings of triples
-        h2, r2, t2 = model2.get_triple_representation(x)
-        # (3.5) Move (3.4) into a single GPU
-        h2, r2, t2 = h2.pin_memory().to(device2, non_blocking=True), r2.pin_memory().to(device2, non_blocking=True), t2.pin_memory().to(device2, non_blocking=True)
-        # 3.6 Forward Pass
-        yhat2 = model2.score(h2, r2, t2).to(device1)
-        # (3.7) Composite Prediction
-        yhat = yhat1 + yhat2
-        # (3.8) Compute Loss
-        batch_loss = loss_function(yhat, y)
-        # (3.9) Compute gradients (Backward Pass)
-        batch_loss.backward()
-        # (3.10) Update parameters
-        opt1.step()
-        opt2.step()
-        # (4) Update epoch loss
-        numpy_batch_loss = batch_loss.item()
-        epoch_loss += numpy_batch_loss
-        print(f"\tBatch Loss:{numpy_batch_loss}\tForward-Backward-Update: {time.time() - start_time}")
-    print(f"Epoch Loss:{epoch_loss}")
-
-def run(args):
-    # (1) Get training data
-    dataloader: torch.utils.data.DataLoader
-    dataloader, num_ent, num_rel = get_train_loader(args)
-    # (2) Get model
-    models, optimizers = get_model(args, num_ent, num_rel)
-    print("Compiling...")
-    # (3) Compile models
-    model1, model2 = models
-    print('####### Model 1 #######')
-    print(model1)
-    print(model1.summarize())
-    model1 = torch.compile(model1)
-    print(model1)
-    print('######## Model2 #######')
-    print(model2)
-    print(model2.summarize())
-    model2 = torch.compile(model2)
-    print(model2)
-    # (4) Get optim
-    opt1, opt2 = optimizers
-    print(opt1)
-    print(opt2)
-    # (5) Get loss func
-    loss_function = model1.loss_function
-    print("Training...")
-
-    device1 = "cuda:0"
-    device2 = "cuda:1"
-    # @TODO: Ensure the multi-node training
-    for e in range(args.num_epochs):
-        epoch_loss = 0
-        if e==-1:
-            args.batch_size+=args.batch_size
-            print(f"Increase Batch size to {args.batch_size}")
-            args.batch_size+=args.batch_size
-            dataloader = MultiEpochsDataLoader(dataset=dataloader.dataset,batch_size=args.batch_size, shuffle=True,num_workers=32)
-
-
-        run_epoch(loss_function,dataloader,model1,model2,opt1,opt2)
-        """
-        for ith, (x, y) in enumerate(tqdm(dataloader)):
-            # (1) Shape the batch
-            x = x.flatten(start_dim=0, end_dim=1)
-            y = y.flatten(start_dim=0, end_dim=1)
-            
-            # (2) Empty the gradients
-            opt1.zero_grad(set_to_none=True)
-            opt2.zero_grad(set_to_none=True)
-            
-            # (3) Forward Backward and Parameter Update
-            start_time = time.time()
-            # (3.1) Select embeddings of triples
-            h1, r1, t1 = model1.get_triple_representation(x)
-            # (3.2) Move (3.1) into a single GPU
-            h1, r1, t1, y = h1.pin_memory().to(device1, non_blocking=True), r1.pin_memory().to(device1,non_blocking=True), t1.pin_memory().to(device1, non_blocking=True), y.pin_memory().to(device1, non_blocking=True)
-            # (3.3) Compute triple score (Forward Pass)
-            yhat1 = model1.score(h1, r1, t1)
-
-            # (3.4) Select second part of the embeddings of triples
-            h2, r2, t2 = model2.get_triple_representation(x)
-            # (3.5) Move (3.4) into a single GPU
-            h2, r2, t2 = h2.pin_memory().to(device2, non_blocking=True), r2.pin_memory().to(device2, non_blocking=True), t2.pin_memory().to(device2, non_blocking=True)
-            # 3.6 Forward Pass
-            yhat2 = model2.score(h2, r2, t2).to(device1)
-            # (3.7) Composite Prediction
-            yhat = yhat1 + yhat2
-            # (3.8) Compute Loss
-            batch_loss = loss_function(yhat, y)
-            # (3.9) Compute gradients (Backward Pass)
-            batch_loss.backward()
-            # (3.10) Update parameters
-            opt1.step()
-            opt2.step()
-            # (4) Update epoch loss
-            numpy_batch_loss = batch_loss.item()
-            epoch_loss += numpy_batch_loss
-            if ith % 1 == 0: # init an argument
-                print(f"\tBatch Loss:{numpy_batch_loss}\tForward-Backward-Update: {time.time() - start_time}")
-        print(f"Epoch:{e}\tEpoch Loss:{epoch_loss}")
-        """
-    print("Saving....")
-    start_time=time.time()
-    model1.to("cpu")
-    model2.to("cpu")
-    print(model1._orig_mod.state_dict())
-    torch.save(model1._orig_mod.state_dict(),f"{model1._orig_mod.name}_1_{e}.torch")
-    print(model2._orig_mod.state_dict())
-    torch.save(model2._orig_mod.state_dict(),f"{model1._orig_mod.name}_2_{e}.torch")
-    print('DONE')
-    print(f"took {time.time() - start_time}")
-
-
-if __name__ == '__main__':
-    run(input_arguments())
-
-# @TODO Post Processing
-# Note mode1 and model2 keci with p=0, q=1
-# model1 real_m1:[] complex_m1[]
-# model2 real_m2:[] complex_m2[]
-# y1 y2 => Final model = real_m1[],real_m2[], complex_m1[] complex_m2[]
diff --git a/examples/multi_hop_query_answering/benchmarking.py b/examples/multi_hop_query_answering/benchmarking.py
index 25286986..9987c0cb 100644
--- a/examples/multi_hop_query_answering/benchmarking.py
+++ b/examples/multi_hop_query_answering/benchmarking.py
@@ -31,7 +31,7 @@
     args = Namespace()
     args.model = kge_name
     args.scoring_technique = "KvsAll"
-    args.path_dataset_folder = "KGs/UMLS"
+    args.dataset_dir = "KGs/UMLS"
     args.num_epochs = 20
     args.batch_size = 1024
     args.lr = 0.1
diff --git a/examples/ppe.py b/examples/ppe.py
deleted file mode 100644
index a39a6c5d..00000000
--- a/examples/ppe.py
+++ /dev/null
@@ -1,44 +0,0 @@
-from dicee.executer import Execute
-from dicee.config import Namespace
-
-args = Namespace()
-args.model = 'Keci'
-args.p = 0
-args.q = 1
-args.scoring_technique = "KvsAll"
-args.dataset_dir = "KGs/UMLS"
-args.num_epochs = 200
-args.lr = 0.1
-args.embedding_dim = 32
-args.batch_size = 1024
-reports = Execute(args).start()
-"""
-Evaluate Keci on Train set: Evaluate Keci on Train set
-{'H@1': 0.9966449386503068, 'H@3': 1.0, 'H@10': 1.0, 'MRR': 0.9983064928425357}
-Evaluate Keci on Validation set: Evaluate Keci on Validation set
-{'H@1': 0.6134969325153374, 'H@3': 0.8098159509202454, 'H@10': 0.9424846625766872, 'MRR': 0.7293869361804316}
-Evaluate Keci on Test set: Evaluate Keci on Test set
-{'H@1': 0.6437216338880484, 'H@3': 0.8275340393343419, 'H@10': 0.959909228441755, 'MRR': 0.751216359363361}
-Total Runtime: 13.259 seconds
-"""
-args = Namespace()
-args.model = 'Keci'
-args.p = 0
-args.q = 1
-args.scoring_technique = "KvsAll"
-args.dataset_dir = "KGs/UMLS"
-args.num_epochs = 200
-args.lr = 0.1
-args.embedding_dim = 32
-args.batch_size = 1024
-args.callbacks = {"PPE": {"epoch_to_start": 100}}
-reports = Execute(args).start()
-"""
-Evaluate Keci on Train set: Evaluate Keci on Train set
-{'H@1': 0.9934815950920245, 'H@3': 1.0, 'H@10': 1.0, 'MRR': 0.9966609151329243}
-Evaluate Keci on Validation set: Evaluate Keci on Validation set
-{'H@1': 0.7001533742331288, 'H@3': 0.8696319018404908, 'H@10': 0.9585889570552147, 'MRR': 0.7946759330503159}
-Evaluate Keci on Test set: Evaluate Keci on Test set
-{'H@1': 0.710287443267776, 'H@3': 0.8789712556732224, 'H@10': 0.9780635400907716, 'MRR': 0.8082179592109334}
-Total Runtime: 12.497 seconds
-"""
\ No newline at end of file
diff --git a/examples/train_kge_on_kg_larger_than_memory.py b/examples/train_kge_on_kg_larger_than_memory.py
new file mode 100644
index 00000000..fa491f09
--- /dev/null
+++ b/examples/train_kge_on_kg_larger_than_memory.py
@@ -0,0 +1,340 @@
+"""
+Training a large knowledge graph embedding model on a knowledge graph that does not fit in memory.
+Computation
+(1) Read/Index/Store a knowledge graph
+(1.1) Extract unique entities and assign unique integers them
+(1.2) Extract unique relations and assign unique integers them
+(1.4) Index an input knowledge graph batch wise
+(1.4.1) Read a batch in memory, replace each subject, predicate, and object with their indices
+(1.4.2) Store a batch into disk as CSV.
+
+(2) Train a KGE model
+(2.1) Initialize a KGE model
+(2.2) Construct a memory map for the indexed knowledge graph that is a concatenation of index batch csv files.
+(2.3) Initialize a dataset based on (2.2) that generates negative examples on the fly.
+
+        SETUP
++ A machine with 62.2 RAM and Intel Core U7
++ dbpedia-generic-snapshot-2022-12 without literals having a size of 590,550,727 (wc -l dbpedia.nt)
+
+        RUN
+(dice) python examples/train_kge_on_kg_larger_than_memory.py --path_dataset dbpedia.nt --path_csv_index_entities dbpedia_index_entities.csv --path_csv_index_relations dbpedia_index_relations.csv --path_indexed_dataset indexed_dbpedia
+Collecting subject entities...
+Unique number of subjects:40890146
+Collecting object entities...
+Unique number of objects:117367061
+Batch Processing...
+Step:0  Seen Triples: 59228337  Total Runtime: 28.149 secs
+Step:1  Seen Triples: 118617771 Total Runtime: 58.033 secs
+Step:2  Seen Triples: 177962380 Total Runtime: 86.136 secs
+Step:3  Seen Triples: 237353864 Total Runtime: 116.173 secs
+Step:4  Seen Triples: 296632299 Total Runtime: 143.573 secs
+Step:5  Seen Triples: 355872505 Total Runtime: 172.268 secs
+Step:6  Seen Triples: 415059906 Total Runtime: 200.339 secs
+Step:7  Seen Triples: 474349288 Total Runtime: 231.305 secs
+Step:8  Seen Triples: 533641626 Total Runtime: 264.210 secs
+Step:9  Seen Triples: 578714291 Total Runtime: 290.475 secs
+Step:10 Seen Triples: 590550727 Total Runtime: 304.888 secs
+Total Runtime:304.8876941204071
+Memory usage of entity_to_idx dataframe is 7788.37 MB
+Memory usage of relation_to_idx dataframe is 0.54 MB
+Parsing CSV files...
+Creating a memory-map to a graph stored in a binary file on disk.
+Adding indexed_dbpedia/Start_row_0_End_row_59228337.np
+Adding indexed_dbpedia/Start_row_59228337_End_row_118617771.np
+Adding indexed_dbpedia/Start_row_118617771_End_row_177962380.np
+Adding indexed_dbpedia/Start_row_177962380_End_row_237353864.np
+Adding indexed_dbpedia/Start_row_237353864_End_row_296632299.np
+Adding indexed_dbpedia/Start_row_296632299_End_row_355872505.np
+Adding indexed_dbpedia/Start_row_355872505_End_row_415059906.np
+Adding indexed_dbpedia/Start_row_415059906_End_row_474349288.np
+Adding indexed_dbpedia/Start_row_474349288_End_row_533641626.np
+Adding indexed_dbpedia/Start_row_533641626_End_row_578714291.np
+Adding indexed_dbpedia/Start_row_578714291_End_row_590550727.np
+Concatenated data saved to indexed_dbpedia/indexed_knowledge_graph.npy
+Total number of parameters: 4088141280
+Total memory size: 16352565120 bytes
+Number of batch updates per epoch 11812
+Batch [0 | 11812]       Loss [2.3336198329925537]       Forward/Backward/Update [8.225] BatchFetch [1.447]      RT [9.672]
+Batch [1 | 11812]       Loss [2.3810811042785645]       Forward/Backward/Update [4.082] BatchFetch [1.500]      RT [5.582]
+Batch [2 | 11812]       Loss [2.3562967777252197]       Forward/Backward/Update [3.949] BatchFetch [1.339]      RT [5.287]
+Batch [3 | 11812]       Loss [2.368497371673584]        Forward/Backward/Update [4.032] BatchFetch [1.278]      RT [5.309]
+Batch [4 | 11812]       Loss [2.3351900577545166]       Forward/Backward/Update [4.021] BatchFetch [1.404]      RT [5.424]
+Batch [5 | 11812]       Loss [2.4003183841705322]       Forward/Backward/Update [3.938] BatchFetch [1.254]      RT [5.192]
+...
+
+
+"""
+import argparse
+import polars
+import polars as pl
+from torch.nn import functional as F
+from dicee import Keci
+import os
+import torch
+from torch.utils.data import Dataset, DataLoader
+import time
+import numpy as np
+from typing import Tuple
+
+class OnevsSampleDataset(Dataset):
+    def __init__(self, mmap_kg:np.memmap, num_entities: int = None, neg_sample_ratio: int = 2):
+        assert num_entities >= 1
+        self.memmap_g = mmap_kg
+        self.num_points = len(self.memmap_g)
+        self.num_entities = 30_000 if num_entities > 2 ** 20 else num_entities
+        self.neg_sample_ratio = neg_sample_ratio
+
+    def __len__(self):
+        return self.num_points
+
+    def __getitem__(self, idx):
+        # pytorch triple from the memory map of numpy.
+        triple = torch.from_numpy(self.memmap_g[idx].copy())
+        x = triple[:2]
+        y = triple[-1].unsqueeze(-1)
+        # Initialize weights for negative sampling. This corresponds to sampling with replacement
+        negative_idx = torch.randint(low=0,
+                                     high=self.num_entities,
+                                     size=(self.neg_sample_ratio,),
+                                     device="cpu")
+        # Concatenate the true tail entity with the negative samples
+        y_idx = torch.cat((y, negative_idx), 0).long()
+        y_vec = torch.zeros(self.neg_sample_ratio + 1, device='cpu')
+        y_vec[0] = 1
+        return x, y_idx, y_vec
+
+def get_default_arguments():
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--path_dataset", type=str, default=None,
+                        help="Path to a knowledge graph in the n-triple format without containing literals")
+    parser.add_argument("--path_csv_index_entities", type=str, default=None, required=True)
+    parser.add_argument("--path_csv_index_relations", type=str, default=None, required=True)
+    parser.add_argument("--path_indexed_dataset", type=str, default=None, required=True)
+    parser.add_argument("--preprocessing_batch_size", type=int, default=50_000_000)
+    parser.add_argument("--batch_size", type=int, default=50_000, help="Batch size for the SGD training.")
+    parser.add_argument("--neg_sample_ratio", type=int, default=100, help="Number of negative examples per positive example.")
+    parser.add_argument("--num_epochs", type=int, default=100, help="Number of training epochs.")
+    parser.add_argument("--embedding_dim", type=int, default=32, help="Number of training epochs.")
+    return parser.parse_args()
+
+def create_indexing(args)->Tuple[polars.DataFrame,polars.DataFrame]:
+    # () Collect unique entities.
+    # Note: Setting maintain_order=True in unique() increases the memory usage.
+    lazy_df = pl.scan_csv(args.path_dataset,
+                          has_header=False,
+                          separator=" ",
+                          new_columns=['subject', 'relation', 'object', "end"]).drop('end')
+    # () Select unique subject entities.
+    print("Collecting subject entities...")
+    subjects = lazy_df.select(pl.col("subject").unique().alias("entity")).collect()
+    print(f"Unique number of subjects:{len(subjects)}")
+    # () Select unique object entities.
+    print("Collecting object entities...")
+    objects = lazy_df.select(pl.col("object").unique().alias("entity")).collect()
+    print(f"Unique number of objects:{len(objects)}")
+    # () Select unique entities.
+    entity_to_idx = pl.concat([subjects, objects], how="vertical").unique()
+    entity_to_idx = entity_to_idx.with_row_index("index").select(["index", "entity"])
+    # () Write unique entities with indices.
+    entity_to_idx.write_csv(file=args.path_csv_index_entities, include_header=True)
+    del subjects, objects
+    # () Collect relations.
+    relation_to_idx = lazy_df.select(pl.col("relation").unique()).collect(streaming=True).with_row_index(
+        "index").select(["index", "relation"])
+    relation_to_idx.write_csv(file=args.path_csv_index_relations, include_header=True)
+    return entity_to_idx, relation_to_idx
+
+def index_knowledge_graph_index(path_dataset:str=None,
+                                path_indexed_dataset:str=None,
+                                preprocessing_batch_size:int=None,
+                                entity_to_idx:polars.DataFrame=None,
+                                relation_to_idx:polars.DataFrame=None)->None:
+    assert path_dataset is not None, f"path_dataset ({path_dataset}) must be specified)"
+    assert path_indexed_dataset is not None, f"path_indexed_dataset ({path_indexed_dataset}) must be specified)"
+    # Create directory if it does not exist, without raising an error
+    os.makedirs(path_indexed_dataset, exist_ok=True)
+    start_time=time.time()
+    start_row = 0
+    end_row = 0
+    reader = pl.read_csv_batched(source=path_dataset,
+                                 batch_size=preprocessing_batch_size,
+                                 has_header=False,
+                                 separator=" ",
+                                 schema_overrides=pl.Schema({"subject": pl.String,
+                                                             "relation": pl.String(),
+                                                             "object": pl.String(),
+                                                             "end": pl.String()}))
+    print("Batch Processing...")
+
+    for i, batches in enumerate(iter(lambda: reader.next_batches(1), None)):
+        # if next_batches(n>2), we need to apply the concat opt.
+        pl_batch = pl.concat(batches).drop("end")
+        end_row += len(pl_batch)
+
+        # Step : Join on 'relation' to replace relation with its index
+        df_merged = pl_batch.join(relation_to_idx, on="relation", how="left")
+        df_merged = df_merged.select([pl.col("subject"), pl.col("index").alias("relation"), pl.col("object")])
+        # Step :  Consider Left Table on subject and Right Table on entity with the left join
+        # Returns all rows from the left table, and the matched rows from the right table
+        df_merged = df_merged.join(entity_to_idx, left_on="subject", right_on="entity", how="left")
+        df_merged = df_merged.drop("subject").rename({"index": "subject"})
+
+        # Step 3: Join on 'object' to replace object with its index
+        df_final = df_merged.join(entity_to_idx, left_on="object", right_on="entity", how="left")
+        df_final = df_final.drop("object").rename({"index": "object"})
+
+        # Step 4: Select the desired columns
+        df_final = df_final.select([pl.col("subject"), pl.col("relation"), pl.col("object")])
+        # why numpy?
+        df_final.write_csv(f"{path_indexed_dataset}/Start_row_{start_row}_End_row_{end_row}.csv", include_header=False)
+
+        start_row = end_row
+        print(f"Step:{i}\tSeen Triples: {end_row}\tTotal Runtime: {time.time() - start_time:.3f} secs")
+
+    print(f"Total Runtime:{time.time() - start_time}")
+
+def prepare_dataset(args)->Tuple[int,int,np.memmap]:
+    if (os.path.exists(args.path_indexed_dataset) and os.path.exists(args.path_csv_index_entities)
+            and os.path.exists(args.path_csv_index_relations)):
+            # () Read entities with indices.
+            print("Reading entities...", end="\t")
+            entity_to_idx = pl.read_csv(args.path_csv_index_entities)
+            print(f"Unique number of entities:\t{len(entity_to_idx)}")
+            print("Reading relations...", end="\t")
+            relation_to_idx = pl.read_csv(args.path_csv_index_relations)
+            print(f"Unique number of relations:\t{len(relation_to_idx)}")
+
+    else:
+        # Ensure that the folder exists.
+        assert os.path.exists(args.path_dataset), f"Path ({args.path_dataset}) does not exist!"
+        entity_to_idx, relation_to_idx = create_indexing(args)
+        assert isinstance(args.preprocessing_batch_size,int),f"--preprocessing_batch_size must be int!"
+        assert args.preprocessing_batch_size > 1_000_000, f"--preprocessing_batch_size must be greater than 10⁶ !"
+        index_knowledge_graph_index(path_dataset=args.path_dataset,
+                                    path_indexed_dataset=args.path_indexed_dataset,
+                                    preprocessing_batch_size=args.preprocessing_batch_size,
+                                    entity_to_idx=entity_to_idx,
+                                    relation_to_idx=relation_to_idx)
+
+    print(f"Memory usage of entity_to_idx dataframe is {round(entity_to_idx.estimated_size('mb'), 2)} MB")
+    print(f"Memory usage of relation_to_idx dataframe is {round(relation_to_idx.estimated_size('mb'), 2)} MB")
+
+    num_relations = len(relation_to_idx)
+    num_entities = len(entity_to_idx)
+
+    # Create a memory-mapped file for the concatenated array
+    output_filename = f"{args.path_indexed_dataset}/indexed_knowledge_graph.npy"
+
+    if os.path.exists(output_filename) is False:
+        file_ranges = []
+        print("Parsing CSV files...")
+        for filename in os.listdir(args.path_indexed_dataset):
+            if filename.endswith('.csv'):
+                # Extract start and end row numbers from the filename
+                parts = filename.replace('.csv', '').split('_')
+                start_row = int(parts[2])
+                end_row = int(parts[-1])
+                file_path = os.path.join(args.path_indexed_dataset, filename)
+                file_ranges.append((file_path, start_row, end_row))
+        # Sort by start_row
+        file_ranges = sorted(file_ranges, key=lambda x: x[1])
+        file_shapes = []  # To keep track of the shapes of each numpy array
+        print("Creating a memory-map to a graph stored in a binary file on disk.")
+        created_files = []
+        num_triples = 0
+        for f in file_ranges:
+            csv_path, start_row, end_row = f
+            # csv to pt file name change
+            new_name = csv_path.replace(".csv", ".np")
+            np_triple = pl.read_csv(csv_path, has_header=False).to_numpy()
+            # https://numpy.org/doc/stable/reference/generated/numpy.ndarray.tofile.html#numpy-ndarray-tofile
+            np_triple.tofile(new_name)
+            num_triples += np_triple.size // 3
+            created_files.append(new_name)
+            file_shapes.append(np_triple.shape)
+        mmap_kg = np.memmap(output_filename, dtype=np.int64, mode='w+', shape=(num_triples, 3))
+        # Concatenate all the data into the new memory map
+        current_position = 0
+        for file_path, shape in zip(created_files, file_shapes):
+            g = np.memmap(file_path, mode="r", dtype=np.int64, shape=shape)
+            print("Adding", file_path)
+            mmap_kg[current_position:current_position + len(g):, :] = g
+            current_position += len(g)
+        # Flush the memory-mapped file to disk
+        mmap_kg.flush()
+        print(f"Concatenated data saved to {output_filename}")
+    else:
+        n = len(np.memmap(output_filename, dtype=np.int64, mode='r'))#, shape=(num_triples, 3))
+        mmap_kg = np.memmap(output_filename, dtype=np.int64, mode='r', shape=(n//3, 3))
+
+    return num_entities, num_relations, mmap_kg
+
+def run(args):
+    num_entities, num_relations, mmap_kg = prepare_dataset(args)
+    model = Keci(args={"embedding_dim": args.embedding_dim, "p": 0, "q": 1,
+                       "num_entities": num_entities, "num_relations": num_relations})
+
+    # Number of parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Total number of parameters: {total_params}")
+
+    # Memory size in bytes
+    total_memory = sum(p.numel() * p.element_size() for p in model.parameters())
+    print(f"Total memory size: {total_memory} bytes")
+
+    optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01)
+    rows_to_copy = 100
+    sub_mmap_kg = np.memmap('destination_file.dat', dtype=mmap_kg.dtype, mode='w+', shape=(rows_to_copy, 3))
+    # Copy the first few rows from the source to the destination
+    sub_mmap_kg[:] = mmap_kg[:rows_to_copy]
+
+    # Multi-class classification per triple.
+    dataset = OnevsSampleDataset(mmap_kg=sub_mmap_kg,
+                                 num_entities=num_entities,
+                                 neg_sample_ratio=args.neg_sample_ratio)
+    dataloader = DataLoader(dataset,
+                            batch_size=args.batch_size,
+                            shuffle=False)
+    num_of_batches = len(dataloader)
+    print(f"Number of batch updates per epoch {num_of_batches}")
+    for epoch_id in range(args.num_epochs):
+        iter_dataloader = iter(dataloader)
+        for batch_id in range(num_of_batches):
+            batch_time = time.time()
+            batch_fetch_time = time.time()
+            # () Get the next batch.
+            try:
+                x, y_idx, y_label = next(iter_dataloader)
+            except StopIteration:
+                continue
+            batch_fetch_time = time.time() - batch_fetch_time
+            rt_forward_backward_update = time.time()
+            # () Clean gradients
+            optimizer.zero_grad()
+            # () - log(1/N) => -math.log(1/127741846) = 18.83
+            logits = model.forward_k_vs_sample(x=x, target_entity_idx=y_idx)
+            # ()
+            batch_loss = F.binary_cross_entropy_with_logits(logits, y_label)
+            # ()
+            batch_loss_float = batch_loss.item()
+            # ()
+            batch_loss.backward()
+            # ()
+            optimizer.step()
+            # ()
+            rt_forward_backward_update = time.time() - rt_forward_backward_update
+            # ()
+            print(f"Epoch [{epoch_id} | {args.num_epochs}]\t"
+                  f"Batch [{batch_id} | {num_of_batches}]\t"
+                  f"Loss [{batch_loss_float}]\t"
+                  f"Forward/Backward/Update [{rt_forward_backward_update:.3f}]\t"
+                  f"BatchFetch [{batch_fetch_time:.3f}]\t"
+                  f"RT [{time.time() - batch_time:.3f}]")
+        # TODO: Epochs End Save the model
+if __name__ == '__main__':
+    run(get_default_arguments())
+
+
+
+
diff --git a/examples/training.py b/examples/training.py
new file mode 100644
index 00000000..a2ed7535
--- /dev/null
+++ b/examples/training.py
@@ -0,0 +1,34 @@
+from dicee.models import Keci
+import numpy as np
+import torch
+# Indexed Triples
+X=torch.from_numpy(np.array([[0,0,1],
+                               [0,1,1],
+                               [1,1,1]])).long()
+
+# Labels
+y=torch.from_numpy(np.array([1,1,0])).float()
+
+# Model
+model=Keci(args={"num_entities":2,"num_relations":2,"embedding_dim":8,"optim":"Adopt"})
+# Optim
+optim=torch.optim.Adam(model.parameters(),lr=0.1)
+#
+model.train()
+for i in range(10):
+    optim.zero_grad()
+
+    yhat=model(X)
+
+    loss=torch.nn.functional.binary_cross_entropy_with_logits(yhat,y)
+
+    print(loss.item())
+
+    loss.backward()
+
+    optim.step()
+
+model.eval()
+
+with torch.no_grad():
+    print(model(X).mean())
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 8c925a22..f0f8291b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 numpy==1.26.4
-torch>=2.2.0
+torch>=2.5.1
 lightning>=2.1.3
 pandas>=2.1.0
 polars>=0.16.14
diff --git a/setup.py b/setup.py
index 63a70d0a..5b9b97d3 100644
--- a/setup.py
+++ b/setup.py
@@ -5,11 +5,11 @@
 # Dev version :pip3 install -e .["dev"]
 
 _deps = [
-    "torch>=2.2.0",
+    "torch>=2.5.1",
     "lightning>=2.1.3",
     "pandas>=2.1.0",
     "numpy==1.26.4",
-    "polars==0.16.14",
+    "polars==1.9.0",
     "scikit-learn>=1.2.2",
     "pyarrow>=11.0.0",
     "pykeen>=1.10.2",


Test patch:
diff --git a/tests/test_answer_multi_hop_query.py b/tests/test_answer_multi_hop_query.py
index 091e9d8c..30970b06 100644
--- a/tests/test_answer_multi_hop_query.py
+++ b/tests/test_answer_multi_hop_query.py
@@ -1,11 +1,6 @@
 from dicee.executer import Execute
 from dicee.config import Namespace
 from dicee.knowledge_graph_embeddings import KGE
-from dicee import QueryGenerator
-from dicee.static_funcs import evaluate
-from dicee.static_funcs import load_pickle
-from dicee.static_funcs import load_json
-
 import pytest
 
 class TestQueryAnswering:
@@ -65,6 +60,6 @@ def test_query_answering_family(self):
                                                                         "http://www.benchmark.org/family#hasSibling",
                                                                         "http://www.benchmark.org/family#married",
                                                                         "http://www.w3.org/1999/02/22-rdf-syntax-ns#type")),
-                                                             tnorm="min", k=5)
+                                                             tnorm="min", k=10)
         top_entities = [topk_entity for topk_entity, query_score in predictions]
         assert "http://www.benchmark.org/family#Person" in top_entities
\ No newline at end of file
diff --git a/tests/test_continual_training.py b/tests/test_continual_training.py
index c8c63fd5..5aca93e0 100644
--- a/tests/test_continual_training.py
+++ b/tests/test_continual_training.py
@@ -29,7 +29,7 @@ def test_negative_sampling(self):
         result = Execute(args).start()
         assert os.path.isdir(result['path_experiment_folder'])
         pre_trained_kge = KGE(path=result['path_experiment_folder'])
-        kg = KG(dataset_dir=args.dataset_dir)
+        kg = KG(dataset_dir=args.dataset_dir,separator="\s+")
         pre_trained_kge.train(kg, epoch=1, batch_size=args.batch_size)
 
     @pytest.mark.filterwarnings('ignore::UserWarning')
@@ -54,5 +54,5 @@ def test_negative_sampling_Family(self):
         assert os.path.isdir(result['path_experiment_folder'])
         pre_trained_kge = KGE(path=result['path_experiment_folder'])
         kg = KG(args.dataset_dir, entity_to_idx=pre_trained_kge.entity_to_idx,
-                relation_to_idx=pre_trained_kge.relation_to_idx)
+                relation_to_idx=pre_trained_kge.relation_to_idx,separator="\s+")
         pre_trained_kge.train(kg, epoch=1, batch_size=args.batch_size)
diff --git a/tests/test_different_backends.py b/tests/test_different_backends.py
index 8b40bcc3..a68210ee 100644
--- a/tests/test_different_backends.py
+++ b/tests/test_different_backends.py
@@ -28,4 +28,5 @@ def test_pandas_as_backend(self):
         args = Namespace()
         args.dataset_dir = 'KGs/UMLS'
         args.backend = 'polars'
+        args.separator="	"
         Execute(args).start()
diff --git a/tests/test_kvssample.py b/tests/test_onevssample.py
similarity index 95%
rename from tests/test_kvssample.py
rename to tests/test_onevssample.py
index c3811d2e..d3910dfc 100644
--- a/tests/test_kvssample.py
+++ b/tests/test_onevssample.py
@@ -16,7 +16,7 @@ def test_k_vs_sample(self):
         args.input_dropout_rate = 0.0
         args.hidden_dropout_rate = 0.0
         args.feature_map_dropout_rate = 0.0
-        args.scoring_technique = 'KvsSample'
+        args.scoring_technique = '1vsSample'
         args.neg_ratio = 10
         args.num_folds_for_cv = 0
         args.eval_model = 'train_val_test'
@@ -43,7 +43,7 @@ def test_k_vs_sample_regression(self):
         args.input_dropout_rate = 0.0
         args.hidden_dropout_rate = 0.0
         args.feature_map_dropout_rate = 0.0
-        args.scoring_technique = 'KvsSample'
+        args.scoring_technique = '1vsSample'
         # size of entity vocabulary
         args.neg_ratio = 10
         args.weight_decay = 0.0
diff --git a/tests/test_online_learning.py b/tests/test_online_learning.py
index 7a980918..f993599b 100644
--- a/tests/test_online_learning.py
+++ b/tests/test_online_learning.py
@@ -10,11 +10,12 @@ class TestRegressionOnlineLearning:
     def test_umls(self):
         args = Namespace()
         args.model = 'AConEx'
-        args.scoring_technique = 'KvsSample'
+        args.scoring_technique = '1vsSample'
         args.optim = 'Adam'
         args.dataset_dir = 'KGs/UMLS'
         args.num_epochs = 0
         args.batch_size = 1024
+        args.neg_ratio=1
         args.lr = 0.1
         args.embedding_dim = 32
         args.input_dropout_rate = 0.0
@@ -23,7 +24,7 @@ def test_umls(self):
         args.read_only_few = None
         args.sample_triples_ratio = None
         args.num_folds_for_cv = None
-        args.backend = 'polars'
+        args.backend = 'pandas'
         args.trainer = 'torchCPUTrainer'
         result = Execute(args).start()
         assert os.path.isdir(result['path_experiment_folder'])
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index 355ad364..13842906 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -3,7 +3,7 @@
 import pickle
 import pytest
 from dicee.config import Namespace
-
+# TODO:CD: Using pickle is quite inefficient timewise. Why do we need the pickle models ?!
 class TestPickle:
     @pytest.mark.filterwarnings('ignore::UserWarning')
     def test_dismult_pickle(self):
diff --git a/tests/test_regression_cl.py b/tests/test_regression_cl.py
index 3494db1c..cf562dbf 100644
--- a/tests/test_regression_cl.py
+++ b/tests/test_regression_cl.py
@@ -6,11 +6,12 @@ class TestRegressionAConEx:
     def test_k_vs_all(self):
         args = Namespace()
         args.model = 'AConEx'
-        args.scoring_technique = 'KvsSample'
+        args.scoring_technique = '1vsSample'
         args.optim = 'Adam'
         args.dataset_dir = 'KGs/UMLS'
-        args.num_epochs = 5
-        args.batch_size = 4096
+        args.num_epochs = 10
+        args.neg_ratio = 1
+        args.batch_size = 256
         args.lr = 0.1
         args.embedding_dim = 32
         args.input_dropout_rate = 0.0
diff --git a/tests/test_regression_model_paralelisim.py b/tests/test_regression_model_paralelisim.py
new file mode 100644
index 00000000..82c0a247
--- /dev/null
+++ b/tests/test_regression_model_paralelisim.py
@@ -0,0 +1,29 @@
+import pytest
+from dicee.static_funcs import write_csv_from_model_parallel
+from dicee.executer import Execute
+from dicee.config import Namespace
+import os
+import torch
+class TestRegressionTensorParallel:
+    @pytest.mark.filterwarnings('ignore::UserWarning')
+    def test_k_vs_all(self):
+        if torch.cuda.is_available():
+            os.system(f'rm -rf Keci_UMLS')
+            args = Namespace()
+            args.model = 'Keci'
+            args.trainer = "TP"
+            args.scoring_technique = "KvsAll"  # 1vsAll, or AllvsAll, or NegSample
+            args.dataset_dir = "KGs/UMLS"
+            args.path_to_store_single_run = "Keci_UMLS"
+            args.save_embeddings_as_csv=True
+            args.optim="Adopt"
+            args.num_epochs = 100
+            args.embedding_dim = 32
+            args.batch_size = 32
+            args.lr=0.1
+            reports = Execute(args).start()
+            assert reports["Train"]["MRR"] >= 0.60
+            assert reports["Test"]["MRR"] >= 0.58
+            assert os.path.exists("Keci_UMLS/Keci_entity_embeddings.csv")
+            assert os.path.exists("Keci_UMLS/Keci_relation_embeddings.csv")
+            os.system(f'rm -rf Keci_UMLS')
diff --git a/tests/test_saving_embeddings.py b/tests/test_saving_embeddings.py
new file mode 100644
index 00000000..f83a03e0
--- /dev/null
+++ b/tests/test_saving_embeddings.py
@@ -0,0 +1,47 @@
+import pytest
+from dicee.static_funcs import from_pretrained_model_write_embeddings_into_csv
+from dicee.executer import Execute
+from dicee.config import Namespace
+import torch
+
+class TestSavingEmbeddings:
+    def test_saving_embeddings(self):
+        # (1) Train a KGE model
+        args = Namespace()
+        args.model = 'Keci'
+        args.p = 0
+        args.q = 1
+        args.optim = 'Adam'
+        args.scoring_technique = "KvsAll"
+        args.path_single_kg = "KGs/Family/family-benchmark_rich_background.owl"
+        args.backend = "rdflib"
+        args.num_epochs = 0
+        args.batch_size = 1024
+        args.lr = 0.1
+        args.embedding_dim = 512
+        result = Execute(args).start()
+        from_pretrained_model_write_embeddings_into_csv(result["path_experiment_folder"])
+
+    def test_model_parallel_saving_embeddings(self):
+        # (1) Train a KGE model
+        if torch.cuda.is_available():
+            from dicee.config import Namespace
+            from dicee.executer import Execute
+            import os
+            args = Namespace()
+            args.model = 'Keci'
+            args.p = 0
+            args.q = 1
+            args.optim = 'Adam'
+            args.scoring_technique = "KvsAll"
+            args.path_single_kg = "KGs/Family/family-benchmark_rich_background.owl"
+            args.backend = "rdflib"
+            args.trainer = "TP"
+            args.num_epochs = 1
+            args.batch_size = 1024
+            args.lr = 0.1
+            args.embedding_dim = 32
+            args.save_embeddings_as_csv = True
+            result = Execute(args).start()
+            assert os.path.exists(result["path_experiment_folder"] + "/Keci_entity_embeddings.csv")
+            assert os.path.exists(result["path_experiment_folder"] + "/Keci_relation_embeddings.csv")