Repository: arangoml/fastgraphml (language: python)

Issue:
Bandit B101 - assert used
Creating an issue to preserve conversation with @sachinsharma9780 regarding the Bandit security issue with the use of `assert` in non-test code.

Here is the error produced:
```
>> Issue: [B101:assert_used] Use of assert detected. The enclosed code will be removed when compiling to optimised byte code.
   Severity: Low   Confidence: High
   CWE: CWE-703 (https://cwe.mitre.org/data/definitions/703.html)
   Location: fastgraphml/graph_embeddings/downstream_tasks/similarity_search.py:64:8
   More Info: https://bandit.readthedocs.io/en/0.0.0/plugins/b101_assert_used.html
63          else:
64              assert (
65                  search_type == "exact" or search_type == "approx"
66              ), "pass search type either exact or approx"
```
Upon reading further into https://bandit.readthedocs.io/en/0.0.0/plugins/b101_assert_used.html as assert can be removed from code while compiling in some instances it seems reasonable to abide by the rule and replace assert with something else, in this case, perhaps an exception.

I can make the change and will close the issue unless we have any other reason to keep assert? 


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index b0c5fa3..2f31853 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -2,21 +2,65 @@ name: build
 on:
   workflow_dispatch:
   push:
-    branches: [ master ]
+    branches: [ main ]
   pull_request:
-    branches: [ master ]
+    branches: [ main ]
+env:
+  PACKAGE_DIR: fastgraphml
+  TESTS_DIR: tests
+  CONDA_ENV: fastgraphml
 jobs:
-  build:
+  lint:
     runs-on: self-hosted
     defaults:
       run:
         shell: bash -l {0}
-    name: gpu
+    strategy:
+      matrix:
+        python: ["3.8"]
+    name: Lint - Python ${{ matrix.python }}
     steps:
       - uses: actions/checkout@v2
       - name: Activating conda env
         run: |
           source ~/miniconda3/etc/profile.d/conda.sh 
-          conda activate fastgraphml
-      - name: Run pytest in conda env
-        run: conda run -n fastgraphml pytest 
+          conda activate ${{env.CONDA_ENV}}
+      - name: Setup Python ${{ matrix.python }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python }}
+      - name: Install packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install .[dev]
+      - name: Run black
+        run: conda run -n ${{env.CONDA_ENV}} black --check --verbose --diff --color ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run flake8
+        run: conda run -n ${{env.CONDA_ENV}} flake8 ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run isort
+        run: conda run -n ${{env.CONDA_ENV}} isort --check --profile=black ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run mypy
+        run: conda run -n ${{env.CONDA_ENV}} mypy ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run bandit
+        run: conda run -n ${{env.CONDA_ENV}} bandit --exclude "./tests/*" --recursive ./
+  test:
+    runs-on: self-hosted
+    defaults:
+      run:
+        shell: bash -l {0}
+    strategy:
+      matrix:
+        python: ["3.8"]
+    name: Test - Python ${{ matrix.python }}
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup Python ${{ matrix.python }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python }}
+      - name: Activating conda env
+        run: |
+          source ~/miniconda3/etc/profile.d/conda.sh 
+          conda activate ${{env.CONDA_ENV}}
+      - name: Install packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install .[dev]
+      - name: Run pytest
+        run: conda run -n ${{env.CONDA_ENV}} pytest --cov=${{env.PACKAGE_DIR}} --cov-report xml --cov-report term-missing -v --color=yes --no-cov-on-fail --code-highlight=yes --cov-fail-under=75
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
new file mode 100644
index 0000000..6302818
--- /dev/null
+++ b/.github/workflows/release.yml
@@ -0,0 +1,155 @@
+name: release
+on:
+  workflow_dispatch:
+  release:
+    types: [published]
+env:
+  PACKAGE_DIR: fastgraphml
+  TESTS_DIR: tests
+  CONDA_ENV: fastgraphml
+jobs:
+  lint:
+    runs-on: self-hosted
+    defaults:
+      run:
+        shell: bash -l {0}
+    strategy:
+      matrix:
+        python: ["3.8"]
+    name: Lint - Python ${{ matrix.python }}
+    steps:
+      - uses: actions/checkout@v2
+      - name: Activating conda env
+        run: |
+          source ~/miniconda3/etc/profile.d/conda.sh 
+          conda activate ${{env.CONDA_ENV}}
+      - name: Setup Python ${{ matrix.python }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python }}
+      - name: Install packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install .[dev]
+      - name: Run black
+        run: conda run -n ${{env.CONDA_ENV}} black --check --verbose --diff --color ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run flake8
+        run: conda run -n ${{env.CONDA_ENV}} flake8 ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run isort
+        run: conda run -n ${{env.CONDA_ENV}} isort --check --profile=black ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run mypy
+        run: conda run -n ${{env.CONDA_ENV}} mypy ${{env.PACKAGE_DIR}} ${{env.TESTS_DIR}}
+      - name: Run bandit
+        run: conda run -n ${{env.CONDA_ENV}} bandit --exclude "./tests/*" --recursive ./
+  test:
+    needs: lint
+    runs-on: self-hosted
+    defaults:
+      run:
+        shell: bash -l {0}
+    strategy:
+      matrix:
+        python: ["3.8"]
+    name: Test - Python ${{ matrix.python }}
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup Python ${{ matrix.python }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python }}
+      - name: Activating conda env
+        run: |
+          source ~/miniconda3/etc/profile.d/conda.sh 
+          conda activate ${{env.CONDA_ENV}}
+      - name: Install packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install .[dev]
+      - name: Run pytest
+        run: conda run -n ${{env.CONDA_ENV}} pytest --cov=${{env.PACKAGE_DIR}} --cov-report xml --cov-report term-missing -v --color=yes --no-cov-on-fail --code-highlight=yes --cov-fail-under=75
+  release:
+    needs: test
+    runs-on: ubuntu-latest
+    name: Release package
+    steps:
+      - name: Activating conda env
+        run: |
+          source ~/miniconda3/etc/profile.d/conda.sh 
+          conda activate ${{env.CONDA_ENV}}
+      - uses: actions/checkout@v2
+
+      - name: Fetch complete history for all tags and branches
+        run: git fetch --prune --unshallow
+
+      - name: Setup python
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.8"
+
+      - name: Install release packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install setuptools wheel twine setuptools-scm[toml]
+
+      - name: Install packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install .[dev]
+
+      - name: Build distribution
+        run: conda run -n ${{env.CONDA_ENV}} python setup.py sdist bdist_wheel
+
+      - name: Publish to PyPI Test
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD_TEST }}
+        run: conda run -n ${{env.CONDA_ENV}} twine upload --repository testpypi dist/* #--skip-existing
+      - name: Publish to PyPI
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}
+        run: conda run -n ${{env.CONDA_ENV}} twine upload --repository pypi dist/* #--skip-existing
+
+  changelog:
+    needs: release
+    runs-on: ubuntu-latest
+    name: Update Changelog
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Create new branch
+        run: git checkout -b actions/changelog
+
+      - name: Set branch upstream
+        run: git push -u origin actions/changelog
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Setup python
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.8"
+
+      - name: Install release packages
+        run: conda run -n ${{env.CONDA_ENV}} pip install wheel gitchangelog pystache
+
+      - name: Set variables
+        run: echo "VERSION=$(curl ${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/releases/latest | python -c "import sys; import json; print(json.load(sys.stdin)['tag_name'])")" >> $GITHUB_ENV
+
+      - name: Generate newest changelog
+        run: gitchangelog ${{env.VERSION}} > CHANGELOG.md
+
+      - name: Make commit for auto-generated changelog
+        uses: EndBug/add-and-commit@v7
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          add: "CHANGELOG.md"
+          branch: actions/changelog
+          message: "!gitchangelog"
+
+      - name: Create pull request for the auto generated changelog
+        run: |
+          echo "PR_URL=$(gh pr create \
+            --title "changelog: release ${{env.VERSION}}" \
+            --body "beep boop, i am a robot" \
+            --label documentation)" >> $GITHUB_ENV
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Alert developer of open PR
+        run: echo "Changelog $PR_URL is ready to be merged by developer."
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 372c13e..9182c1b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,6 @@
 __pycache__/
+build/
+fastgraphml.egg-info/
+fastgraphml/.DS_Store
+dist/
 
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..c87955d
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,33 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.2.0
+    hooks:
+    -   id: check-merge-conflict
+-   repo: https://github.com/PyCQA/isort
+    rev: 5.10.1
+    hooks:
+    -   id: isort
+-   repo: https://github.com/psf/black
+    rev: 22.8.0
+    hooks:
+    -   id: black
+        args:
+            - -l 88
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: "v0.982"
+    hooks:
+    -   id: mypy
+        additional_dependencies: [types-requests]
+        exclude: ^tests/
+-   repo: https://github.com/PyCQA/flake8
+    rev: 5.0.4
+    hooks:
+    -   id: flake8
+        args: 
+            - "--max-line-length=88"
+            - "--ignore=E203, W503, E251"
+-   repo: https://github.com/PyCQA/bandit
+    rev: 1.7.4
+    hooks:
+    -   id: bandit
+        exclude: ^tests/
\ No newline at end of file
diff --git a/README.md b/README.md
index 669ca93..b445d26 100644
--- a/README.md
+++ b/README.md
@@ -2,12 +2,13 @@
 Given an input graph it generates Graph Embeddings using Low-Code framework built on top of [PyG](https://pytorch-geometric.readthedocs.io/en/latest/). The package supports training on both GPU and CPU enabled machines. Training jobs on GPUs results in much faster execution and increased in performance when it comes to handling large graphs as compared to CPUs. In addition, the framework provides tight integration with  [ArangoDB](https://www.arangodb.com/) which is a scalable, fully managed graph database, document store and search engine in one place. Once Graph Embeddings are generated, they can be used for various downstream machine learning tasks like Node Classification, Link Prediction, Visualisation, Community Detection, Similartiy Search, Recommendation, etc. 
 
 ## Installation
-#### Additional Dependencies
-1. [pytorch](https://pytorch.org/) 
-2. [pyg](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) 
-3. [FAISS](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md) 
-
-Note: For FAISS-CPU one needs numba==0.53.0
+#### Required Dependencies
+1. PyTorch `1.12.*` is required.
+    * Install using previous version that matches your CUDA version: [pytorch](https://pytorch.org/get-started/previous-versions/)
+        * To find your installed CUDA version run `nvidia-smi` in your terminal.
+2. [pyg](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html)
+3. [FAISS](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md)
+    * Note: For FAISS-CPU one needs `numba==0.53.0`
 
 #### Latest Release
 ```
@@ -28,29 +29,34 @@ git clone https://github.com/arangoml/fastgraphml.git
 
 ```python
 from fastgraphml.graph_embeddings import SAGE, GAT
-from fastgraphml.graph_embeddings import downstream_tasks 
+from fastgraphml.graph_embeddings import downstream_tasks
+from fastgraphml import Datasets 
 from arango import ArangoClient
 
 # Initialize the ArangoDB client.
 client = ArangoClient("http://127.0.0.1:8529")
-db = client.db('_system', username='root')
+db = client.db('_system', username='root', password='openSesame')
+
+# Loading Amazon Computer Products dataset into ArangoDB
+Datasets(db).load("AMAZON_COMPUTER_PRODUCTS")
+
+# Optionally use arangodb graph
+# arango_graph = db.graph('product_graph')
 
-# arangodb graph name
-arango_graph = db.graph('cora_graph')
 # metadata information of arango_graph
 metagraph = {
     "vertexCollections": {
-        "Paper": {"x": "features", "y": "label"},
+        "Computer_Products": {"x": "features", "y": "label"},
     },
     "edgeCollections": {
-        "Cites": {},
+        "bought_together": {},
     },
 }
 
 # generating graph embeddings with 3 lines of code
-model = SAGE(db, arango_graph, metagraph, embedding_size=64) # define graph embedding model
-model._train(model, epochs=10) # train
-embeddings = model.get_embeddings(model=model) # get embeddings
+model = SAGE(db,'product_graph', metagraph, embedding_size=64) # define graph embedding model
+model._train(epochs=10) # train
+embeddings = model.get_embeddings() # get embeddings
 ```
 
 #### Example Heterogeneous Graphs
@@ -58,13 +64,20 @@ embeddings = model.get_embeddings(model=model) # get embeddings
 ```python
 from fastgraphml.graph_embeddings import METAPATH2VEC, DMGI
 from fastgraphml.graph_embeddings import downstream_tasks 
+from fastgraphml import Datasets 
+
 from arango import ArangoClient
 
 # Initialize the ArangoDB client.
 client = ArangoClient("http://127.0.0.1:8529")
 db = client.db('_system', username='root')
 
-arango_graph = db.graph("IMDB")
+# Loading IMDB Dataset into ArangoDB
+Datasets(db).load("IMDB_X")
+
+# Optionally use ArangoDB Graph
+# arango_graph = db.graph("IMDB")
+
 metagraph = {
     "vertexCollections": {
     
@@ -80,7 +93,7 @@ metapaths = [('movie', 'to','actor'),
              ('actor', 'to', 'movie'), ] # MAM # co-actor relationship
 
 # generating graph embeddings with 3 lines of code
-model = METAPATH2VEC(db, arango_graph, metagraph, metapaths, key_node='movie', embedding_size=128,
+model = METAPATH2VEC(db, "IMDB_X", metagraph, metapaths, key_node='movie', embedding_size=128,
                      walk_length=5, context_size=6, walks_per_node=5, num_negative_samples=5,
                      sparse=True) # define model
 model._train(epochs=10, lr=0.03) # train
@@ -100,8 +113,8 @@ data = dataset[0]
 
 # generating graph embeddings with 3 lines of code
 model = SAGE(pyg_graph=data, embedding_size=64) # define graph embedding model
-model._train(model, epochs=10) # train
-embeddings = model.get_embeddings(model=model) # get embeddings
+model._train(epochs=10) # train
+embeddings = model.get_embeddings() # get embeddings
 ```
 ## Models Supported
 
diff --git a/examples/sage_amazon_computer.ipynb b/examples/sage_amazon_computer.ipynb
index f1eab22..83c13d8 100644
--- a/examples/sage_amazon_computer.ipynb
+++ b/examples/sage_amazon_computer.ipynb
@@ -84,8 +84,8 @@
    "source": [
     "# generating graph embeddings with 3 lines of code\n",
     "model = SAGE(db, arango_graph, metagraph, embedding_size=256) # define graph embedding model\n",
-    "model._train(model, epochs=6, lr=0.0001) # train\n",
-    "embeddings = model.get_embeddings(model=model) # get embeddings"
+    "model._train(epochs=6, lr=0.0001) # train\n",
+    "embeddings = model.get_embeddings() # get embeddings"
    ]
   },
   {
diff --git a/fastgraphml/__init__.py b/fastgraphml/__init__.py
index d565cc7..a26e2a4 100644
--- a/fastgraphml/__init__.py
+++ b/fastgraphml/__init__.py
@@ -1,4 +1,8 @@
+from arango_datasets.datasets import Datasets
+
+from fastgraphml.graph_embeddings.models.dmgi import DMGI
+from fastgraphml.graph_embeddings.models.gat import GAT
 from fastgraphml.graph_embeddings.models.graph_sage import SAGE
-from fastgraphml.graph_embeddings.models.gat import GAT    
 from fastgraphml.graph_embeddings.models.metapath2vec import METAPATH2VEC
-from fastgraphml.graph_embeddings.models.dmgi import DMGI
+
+__all__ = ["DMGI", "GAT", "SAGE", "METAPATH2VEC", "Datasets"]
diff --git a/fastgraphml/graph_embeddings/__init__.py b/fastgraphml/graph_embeddings/__init__.py
index 4575acd..a9279a3 100644
--- a/fastgraphml/graph_embeddings/__init__.py
+++ b/fastgraphml/graph_embeddings/__init__.py
@@ -1,4 +1,6 @@
-from fastgraphml.graph_embeddings.models.gat import GAT  
-from fastgraphml.graph_embeddings.models.graph_sage import SAGE  
-from fastgraphml.graph_embeddings.models.metapath2vec import METAPATH2VEC
 from fastgraphml.graph_embeddings.models.dmgi import DMGI
+from fastgraphml.graph_embeddings.models.gat import GAT
+from fastgraphml.graph_embeddings.models.graph_sage import SAGE
+from fastgraphml.graph_embeddings.models.metapath2vec import METAPATH2VEC
+
+__all__ = ["DMGI", "GAT", "SAGE", "METAPATH2VEC"]
diff --git a/fastgraphml/graph_embeddings/downstream_tasks/__init__.py b/fastgraphml/graph_embeddings/downstream_tasks/__init__.py
index 2279279..b636b28 100644
--- a/fastgraphml/graph_embeddings/downstream_tasks/__init__.py
+++ b/fastgraphml/graph_embeddings/downstream_tasks/__init__.py
@@ -1,2 +1,8 @@
-from fastgraphml.graph_embeddings.downstream_tasks.visualization import visualize_embeddings
-from fastgraphml.graph_embeddings.downstream_tasks.similarity_search import similarity_search
\ No newline at end of file
+from fastgraphml.graph_embeddings.downstream_tasks.similarity_search import (  # noqa: E501
+    similarity_search,
+)
+from fastgraphml.graph_embeddings.downstream_tasks.visualization import (  # noqa: E501
+    visualize_embeddings,
+)
+
+__all__ = ["similarity_search", "visualize_embeddings"]
diff --git a/fastgraphml/graph_embeddings/downstream_tasks/similarity_search.py b/fastgraphml/graph_embeddings/downstream_tasks/similarity_search.py
index fc6c3da..ab72869 100644
--- a/fastgraphml/graph_embeddings/downstream_tasks/similarity_search.py
+++ b/fastgraphml/graph_embeddings/downstream_tasks/similarity_search.py
@@ -1,54 +1,66 @@
+from typing import Any, Tuple, Union
+
 import faiss
+import numpy as np
+import numpy.typing as npt
+
 
+def similarity_search(
+    graph_emb: Union[npt.NDArray[np.float64], Any],
+    top_k_nbors: int = 10,
+    nlist: int = 10,
+    search_type: str = "exact",
+) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.int64]]:
+    """Performs similarity search in sets of vectors of any size, up to ones that
+    possibly do not fit in RAM using FAISS.
 
-def similarity_search(graph_emb, top_k_nbors=10, nlist=10, search_type='exact'):
-    """Performs similarity search in sets of vectors of any size, up to ones that possibly do not fit in RAM using FAISS 
-    <https://engineering.fb.com/2017/03/29/data-infrastructure/faiss-a-library-for-efficient-similarity-search/>. 
+    <https://engineering.fb.com/2017/03/29/data-infrastructure/faiss-a-library-for-efficient-similarity-search/>.  # noqa: E501
 
     returns dist, nbors
         nbors: node ids with similar embeddings
         dist: corresponding similarity scores
-    
+
     :graph_emb (type: 2D numpy array): Numpy array of size (n, embedding_size),
     n: number of nodes in graph
-    embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
+    embedding_size: Length of the node embeddings when they are mapped to
+        d-dimensional euclidean space.
     :top_k_nbors (type: int): Returns top-k nearest neighbors of all embeddings.
     :nlist (type: int): Number of clusters to be generated.
     :search_type (type: str): We support two types of search for now:
     1. exact search: For precise similarity search but at the cost of scalability.
-    2. approx search: For scalable similarity search but at the cost of some precision loss.
+    2. approx search: For scalable similarity search but at the cost of some precision
+        loss.
     """
 
     emb_dim = graph_emb.shape[1]
     # exact search
-    if search_type == 'exact':
+    if search_type == "exact":
         # inner product search index
         search_index = faiss.IndexFlatIP(emb_dim)
         # normalize the embeddings and add to search index
         faiss.normalize_L2(graph_emb)
         search_index.add(graph_emb)
         # search top-k nearest neighbors of all embeddings
-        dist, nbors = search_index.search(graph_emb, k=top_k_nbors+1)
+        dist, nbors = search_index.search(graph_emb, k=top_k_nbors + 1)
 
     # approx search
-    elif search_type == 'approx':
+    elif search_type == "approx":
         # assign embeddings to specific cluster using Inner product
         quantiser = faiss.IndexFlatIP(emb_dim)
         # L2 normalization of embeddings
         faiss.normalize_L2(graph_emb)
         # creates partitioned search index
-        search_index = faiss.IndexIVFFlat(quantiser, emb_dim, nlist, faiss.METRIC_INNER_PRODUCT)
+        search_index = faiss.IndexIVFFlat(
+            quantiser, emb_dim, nlist, faiss.METRIC_INNER_PRODUCT
+        )
         # train and add the embeddings to the index
         search_index.train(graph_emb)
         search_index.add(graph_emb)
         # nbors: top k node ids with highest cosine sim
         # dist: cosine distances of top k nodes
-        dist, nbors = search_index.search(graph_emb, k=top_k_nbors+1)
+        dist, nbors = search_index.search(graph_emb, k=top_k_nbors + 1)
 
     else:
-        assert search_type == 'exact' or search_type == 'approx', "pass search type either exact or approx"
+        raise Exception("Pass search type either exact or approx.")
 
     return dist, nbors
-
-
-                
diff --git a/fastgraphml/graph_embeddings/downstream_tasks/visualization.py b/fastgraphml/graph_embeddings/downstream_tasks/visualization.py
index 123af7e..d44368b 100644
--- a/fastgraphml/graph_embeddings/downstream_tasks/visualization.py
+++ b/fastgraphml/graph_embeddings/downstream_tasks/visualization.py
@@ -1,30 +1,47 @@
-import umap
-import matplotlib.pyplot as plt
-import seaborn as sns
 import math
+from typing import Any, Dict, Optional, Union
+
+import matplotlib.pyplot as plt
 import numpy as np
+import numpy.typing as npt
+import seaborn as sns
+import umap
+from torch_geometric.data import Data
 
 
-def visualize_embeddings(graph, graph_emb, class_mapping=None, node_type=None, emb_percent=0.1):
-    """Performs dimensionality reduction (2D) and visualization of graph embeddings using U-Map
-    
-    :graph (type: PyG data object): PyG data object which can be accessed using model.G inside the framework.
+def visualize_embeddings(
+    graph: Data,
+    graph_emb: Union[npt.NDArray[np.float64], Any],
+    class_mapping: Optional[Dict[int, str]] = None,
+    node_type: Optional[str] = None,
+    emb_percent: float = 0.1,
+) -> None:
+    """Performs dimensionality reduction (2D) and visualization of graph embeddings
+    using U-Map.
+
+    :graph (type: PyG data object): PyG data object which can be accessed using model.G
+        inside the framework.
     :graph_emb (type: 2D numpy array): Numpy array of size (n, embedding_size),
     n: number of nodes in graph,
-    embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
-    :class_mapping (type: dict): It is a dictionary where class names are mapped to integer labels.
-    If class_mappings are not provided the method uses integers labels as legend inside the figure. for e.g.
-    {0: 'Desktops',1: 'Data Storage',2: 'Laptops',3: 'Monitors',4: 'Computer Components',
-    5: 'Video Projectors',6: 'Routers',7: 'Tablets',8: 'Networking Products',9: 'Webcams'} # amazon computer dataset.
-    :node_type (type: str):  Node type for which we want to retreive embeddings to perform visualization. Used to store Hetero graph embeddings.
-    :emb_perc (type: float): Percentage of embeddings to visualize, 0.1 means 10% of data will be visualized.
+    embedding_size: Length of the node embeddings when they are mapped to d-dimensional
+        euclidean space.
+    :class_mapping (type: dict): It is a dictionary where class names are mapped to
+        integer labels. If class_mappings are not provided the method uses integers
+        labels as legend inside the figure. for e.g.
+        {0: 'Desktops',1: 'Data Storage',2: 'Laptops',3: 'Monitors',
+        4: 'Computer Components', 5: 'Video Projectors',6: 'Routers',7: 'Tablets',
+        8: 'Networking Products',9: 'Webcams'} # amazon computer dataset.
+    :node_type (type: str):  Node type for which we want to retreive embeddings to
+         perform visualization. Used to store Hetero graph embeddings.
+    :emb_perc (type: float): Percentage of embeddings to visualize, 0.1 means 10% of
+         data will be visualized.
     """
 
     if node_type is not None:
         y_np = graph[node_type].y.cpu().numpy()
     else:
         y_np = graph.y.cpu().numpy()
-        
+
     num_nodes = math.floor(graph_emb.shape[0] * emb_percent)
     num_y = math.floor(y_np.shape[0] * emb_percent)
     graph_emb = graph_emb[:num_nodes]
@@ -33,17 +50,21 @@ def visualize_embeddings(graph, graph_emb, class_mapping=None, node_type=None, e
     plt.figure(figsize=(8, 8))
     if class_mapping is not None:
         palette = {}
-        class_names = [class_mapping[l] for l in labels]
+        class_names = [class_mapping[label] for label in labels]
         for n, y in enumerate(set(np.array(class_names))):
-            palette[y] = f'C{n}'
-        sns.scatterplot(x=umap_embd.T[0], y=umap_embd.T[1], hue=np.array(class_names), palette=palette)
-        plt.legend(bbox_to_anchor=(1,1), loc='upper left')
+            palette[y] = f"C{n}"
+        sns.scatterplot(
+            x=umap_embd.T[0],
+            y=umap_embd.T[1],
+            hue=np.array(class_names),
+            palette=palette,
+        )
+        plt.legend(bbox_to_anchor=(1, 1), loc="upper left")
         plt.savefig("./umap_embd_visualization.png", dpi=120)
     else:
         palette = {}
         for n, y in enumerate(set(labels)):
-            palette[y] = f'C{n}'
+            palette[y] = f"C{n}"
         sns.scatterplot(x=umap_embd.T[0], y=umap_embd.T[1], hue=labels, palette=palette)
-        plt.legend(bbox_to_anchor=(1,1), loc='upper left')
+        plt.legend(bbox_to_anchor=(1, 1), loc="upper left")
         plt.savefig("./umap_embd_visualization.png", dpi=120)
-        
\ No newline at end of file
diff --git a/fastgraphml/graph_embeddings/models/dmgi.py b/fastgraphml/graph_embeddings/models/dmgi.py
index d1be8fd..ac4d763 100644
--- a/fastgraphml/graph_embeddings/models/dmgi.py
+++ b/fastgraphml/graph_embeddings/models/dmgi.py
@@ -1,55 +1,83 @@
+import shutil
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 import torch
 import torch.nn.functional as F
+from arango.database import Database
 from sklearn.linear_model import LogisticRegression
-from torch.optim import Adam
-import torch_geometric.transforms as T
+from torch import Tensor
+from torch_geometric.data import Data
 from torch_geometric.nn import GCNConv
-import shutil
-from arango.database import Database
+from torch_geometric.typing import Adj, EdgeType, OptPairTensor
+
 from ..utils import GraphUtils
 
 # check for gpu
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-class DMGI(torch.nn.Module):
-    """ Deep Multiplex Graph Infomax (DMGI) model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/hetero/dmgi_unsup.py>'
-    
-        :database (type: Database): A python-arango database instance.
-        :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG.
-        :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define metagraph as 
-                    a dictionary defining vertex & edge collections to import to PyG, along 
-                    with collection-level specifications to indicate which ArangoDB attributes will become PyG features/labels. 
-                    It also supports different encoders such as identity and categorical encoder on database attributes. Detailed information regarding different
-                    use cases and metagraph definitons can be found on adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
-        :metapaths (type: list(list[Tuple(str,str,str))]): It is described as list of list of tuples  . Adds additional edge types to a Hetero Graph between the source node type and the 
-        destination node type of a given metapath. The metapath defined as (src_node_type, rel_type, dst_node_type) tuples. 
-        e.g. Adding two metapaths (reference: <https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.AddMetaPaths>):
-        # 1. From "paper" to "paper" through "conference"
-        # 2. From "author" to "conference" through "paper"
-        metapaths = [[("paper", "conference"), ("conference", "paper")],
-             [("author", "paper"), ("paper", "conference")]] 
-        :key_node (type: str): Node type on which we want to test the performance of generated graph embeddings. Performance is tested using node classification task.
-        :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs (via PyG data objects) directy rather than ArangoDB graphs.
-                    When generating graph embeddings via PyG graphs, database=arango_graph=metagraph=None.
-        :embedding_size (type: int): Length of the node embeddings when they are mapped to d-dimensional euclidean space.
-        :dropout_perc (type: float): Handles overfitting inside the model.
-        :transform (type: torch_geometric.transforms): It is used to transform PyG data objects. Various transformation methods can be chained together using Compose.
-                    for e.g. transform = T.Compose([
-                            T.NormalizeFeatures(),
-                            T.RandomNodeSplit(num_val=0.2, num_test=0.1)])   
-        :num_val (type: float): Percentage of nodes selected for validation set.
-        :num_test (type: float): Percentage of nodes selected for test set.
-        
-        Note: After selecting the percentage for validation and test nodes, rest percentage of the nodes are considered as training nodes.
-        """
 
-    def __init__(self, database = None, arango_graph = None, metagraph = None,
-     metapaths = None, key_node=None, pyg_graph = None, embedding_size=64, dropout_perc = 0.5,
-     transform = None, num_val = 0.1, num_test = 0.1):
+class DMGI(torch.nn.Module):
+    """Deep Multiplex Graph Infomax (DMGI) model
+    modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/hetero/dmgi_unsup.py>'.  # noqa: E501
+
+    :database (type: Database): A python-arango database instance.
+    :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG
+    :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define
+        metagraph as a dictionary defining vertex & edge collections to import to PyG,
+        along with collection-level specifications to indicate which ArangoDB attributes
+        will become PyG features/labels. It also supports different encoders such as
+        identity and categorical encoder on database attributes. Detailed information
+        regarding different use cases and metagraph definitons can be found on
+        adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
+    :metapaths (type: list(list[Tuple(str,str,str))]): It is described as list of list
+        of tuples. Adds additional edge types to a Hetero Graph between the source node
+        type and the destination node type of a given metapath.
+        The metapath defined as (src_node_type, rel_type, dst_node_type) tuples.
+    e.g. Adding two metapaths (reference: <https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.AddMetaPaths>):# noqa: E501
+    # 1. From "paper" to "paper" through "conference"
+    # 2. From "author" to "conference" through "paper"
+    metapaths = [[("paper", "conference"), ("conference", "paper")],
+         [("author", "paper"), ("paper", "conference")]]
+    :key_node (type: str): Node type on which we want to test the performance of
+        generated graph embeddings. Performance is tested using node classification task
+    :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs
+        (via PyG data objects) directy rather than ArangoDB graphs. When generating
+        graph embeddings via PyG graphs, database=arango_graph=metagraph=None.
+    :embedding_size (type: int): Length of the node embeddings when they are mapped to
+        d-dimensional euclidean space.
+    :dropout_perc (type: float): Handles overfitting inside the model.
+    :transform (type: torch_geometric.transforms): It is used to transform PyG data
+        objects. Various transformation methods can be chained together using Compose.
+                for e.g. transform = T.Compose([
+                        T.NormalizeFeatures(),
+                        T.RandomNodeSplit(num_val=0.2, num_test=0.1)])
+    :num_val (type: float): Percentage of nodes selected for validation set.
+    :num_test (type: float): Percentage of nodes selected for test set.
+
+    Note: After selecting the percentage for validation and test nodes,
+    rest percentage of the nodes are considered as training nodes.
+    """
+
+    def __init__(
+        self,
+        database: Database = None,
+        arango_graph: Optional[str] = None,
+        metagraph: Union[Dict[str, object], None] = None,
+        metapaths: Optional[List[EdgeType]] = None,
+        key_node: Optional[str] = None,
+        pyg_graph: Data = None,
+        embedding_size: int = 64,
+        dropout_perc: float = 0.5,
+        transform: Optional[List[Callable[..., Any]]] = None,
+        num_val: float = 0.1,
+        num_test: float = 0.1,
+    ):
         super().__init__()
 
-        if (database is not None or arango_graph is not None or metagraph is not None) and pyg_graph is not None:
-            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"
+        if (
+            database is not None or arango_graph is not None or metagraph is not None
+        ) and pyg_graph is not None:
+            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"  # noqa: E501
             raise Exception(msg)
 
         if database is not None:
@@ -58,8 +86,17 @@ def __init__(self, database = None, arango_graph = None, metagraph = None,
                 raise TypeError(msg)
 
         # arango to PyG
-        self.graph_util = GraphUtils(arango_graph, metagraph, database, pyg_graph, 
-        num_val, num_test, transform, key_node, metapaths)
+        self.graph_util = GraphUtils(
+            arango_graph,
+            metagraph,
+            database,
+            pyg_graph,
+            num_val,
+            num_test,
+            transform,
+            key_node,
+            metapaths,
+        )
         # get PyG graph
         G = self.graph_util.graph
         G = G.to(device)
@@ -67,30 +104,33 @@ def __init__(self, database = None, arango_graph = None, metagraph = None,
         self.x = G[key_node].x.float()
         self.edge_indices = G.edge_index_dict.values()
         self.key_node = key_node
-        
+
         self.metadata = G.metadata()
         self.dropout_perc = dropout_perc
 
         num_nodes = G[key_node].num_nodes
         in_channels = G[key_node].x.size(-1)
         out_channels = embedding_size
-        num_relations = len(G.edge_types) 
+        num_relations = len(G.edge_types)
         # relation type specific node encoder
         # generates realtion type specific node embedding
         self.convs = torch.nn.ModuleList(
-            [GCNConv(in_channels, out_channels) for _ in range(num_relations)])
+            [GCNConv(in_channels, out_channels) for _ in range(num_relations)]
+        )
         self.M = torch.nn.Bilinear(out_channels, out_channels, 1)
         self.Z = torch.nn.Parameter(torch.Tensor(num_nodes, out_channels))
         self.reset_parameters()
 
-    def reset_parameters(self):
+    def reset_parameters(self) -> None:
         for conv in self.convs:
             conv.reset_parameters()
         torch.nn.init.xavier_uniform_(self.M.weight)
         self.M.bias.data.zero_()
         torch.nn.init.xavier_uniform_(self.Z)
 
-    def forward(self, x, edge_indices):
+    def forward(
+        self, x: Union[Tensor, OptPairTensor], edge_indices: Adj
+    ) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]:
         pos_hs, neg_hs, summaries = [], [], []
         for conv, edge_index in zip(self.convs, edge_indices):
             pos_h = F.dropout(x, p=self.dropout_perc, training=self.training)
@@ -106,8 +146,10 @@ def forward(self, x, edge_indices):
 
         return pos_hs, neg_hs, summaries
 
-    def loss(self, pos_hs, neg_hs, summaries):
-        loss = 0.
+    def loss(
+        self, pos_hs: List[Tensor], neg_hs: List[Tensor], summaries: List[Tensor]
+    ) -> Tensor:
+        loss = 0.0
         for pos_h, neg_h, s in zip(pos_hs, neg_hs, summaries):
             s = s.expand_as(pos_h)
             loss += -torch.log(self.M(pos_h, s).sigmoid() + 1e-15).mean()
@@ -124,33 +166,42 @@ def loss(self, pos_hs, neg_hs, summaries):
 
     # save checkpoints whenever there is an increase in validation accuracy
     @staticmethod
-    def save_checkpoints(state, is_best, ckp_path, best_model_path):
+    def save_checkpoints(
+        state: Dict[str, str], is_best: bool, ckp_path: str, best_model_path: str
+    ) -> None:
         file_path = ckp_path
         torch.save(state, file_path)
-        #if it is a best model, min train loss
+        # if it is a best model, min train loss
         if is_best:
             best_file_path = best_model_path
-            # copy best checkpoint file to best model path 
+            # copy best checkpoint file to best model path
             shutil.copyfile(file_path, best_file_path)
 
-
-    def _train(self, model, ckp_path ="./latest_model_checkpoint.pt", best_model_path = "./best_model.pt", 
-    epochs = 51, lr=0.0005, **kwargs):
+    def _train(
+        self,
+        ckp_path: str = "./latest_model_checkpoint.pt",
+        best_model_path: str = "./best_model.pt",
+        epochs: int = 51,
+        lr: float = 0.0005,
+        **kwargs: Any,
+    ) -> None:
         """Train GraphML model.
 
-        :model: Graph embedding model.
-        :ckp_path (type: str): Path to save model's latest checkpoints (i.e. at each epoch). Pytorch models are saved with .pt file extension.
-         By default it saves model in cwd.
-        :best_model_path (type: str): Path to save model whenever there is an increase in validation accuracy. By default it saves model in cwd.
+        :ckp_path (type: str): Path to save model's latest checkpoints
+            (i.e. at each epoch). Pytorch models are saved with .pt file extension.
+            By default it saves model in cwd.
+        :best_model_path (type: str): Path to save model whenever there is an increase
+            in validation accuracy. By default it saves model in cwd.
         :epochs (type: int): Number of times training data go through the model.
         :lr (type: float): Learning rate.
-        :**kwargs: Additional arguments for the Adam optimizer for e.g. weight_decay, betas, etc.
+        :**kwargs: Additional arguments for the Adam optimizer for
+            e.g. weight_decay, betas, etc.
         """
 
-        model = model.to(device)
+        model = self.to(device)
         optimizer = torch.optim.Adam(model.parameters(), lr, **kwargs)
         best_acc = 0.0
-        print('Training started .........')
+        print("Training started .........")
         for epoch in range(1, epochs):
             model.train()
             optimizer.zero_grad()
@@ -164,16 +215,20 @@ def _train(self, model, ckp_path ="./latest_model_checkpoint.pt", best_model_pat
             # validate model #
             ##################
 
-            val_acc, test_acc = self.val(model)
+            val_acc, test_acc = self.val()
 
-            print(f'Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, 'f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')
+            print(
+                f"Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, "
+                f"Val: {val_acc:.4f}, Test: {test_acc:.4f}"
+            )
 
             # create checkpoint variable and add crucial model information
             model_checkpoint = {
-            'epoch': epoch,
-            'best_acc': val_acc,
-            'state_dict': model.state_dict(),
-            'optimizer': optimizer.state_dict(),}
+                "epoch": epoch,
+                "best_acc": val_acc,
+                "state_dict": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+            }
 
             # save current/last checkpoint
             # used for resuming training
@@ -181,47 +236,56 @@ def _train(self, model, ckp_path ="./latest_model_checkpoint.pt", best_model_pat
 
             # save model if val acc increases
             if val_acc > best_acc:
-                print('Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...'.format(best_acc, val_acc))
+                print(
+                    "Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...".format(
+                        best_acc, val_acc
+                    )
+                )
                 # save checkpoint as best model
                 self.save_checkpoints(model_checkpoint, True, ckp_path, best_model_path)
                 best_acc = val_acc
 
     @torch.no_grad()
-    def val(self, model):
-        """Tests the performance of a generated graph embeddings using Node Classification as a downstream task.
-        
+    def val(
+        self,
+    ) -> Tuple[float, float]:
+        """Tests the performance of a generated graph embeddings using Node
+        Classification as a downstream task.
+
         returns validation and test accuracy.
 
         model: Graph embedding model.
         """
-        model.eval()
-        train_emb = model.Z[self.G[self.key_node].train_mask].cpu()
-        val_emb = model.Z[self.G[self.key_node].val_mask].cpu()
-        test_emb = model.Z[self.G[self.key_node].test_mask].cpu()
+        self.eval()
+        train_emb = self.Z[self.G[self.key_node].train_mask].cpu()
+        val_emb = self.Z[self.G[self.key_node].val_mask].cpu()
+        test_emb = self.Z[self.G[self.key_node].test_mask].cpu()
 
         train_y = self.G[self.key_node].y[self.G[self.key_node].train_mask].cpu()
         val_y = self.G[self.key_node].y[self.G[self.key_node].val_mask].cpu()
         test_y = self.G[self.key_node].y[self.G[self.key_node].test_mask].cpu()
-        clf = LogisticRegression(max_iter=400, class_weight='balanced')
+        clf = LogisticRegression(max_iter=400, class_weight="balanced")
         clf.fit(train_emb, train_y)
-        val_acc =  clf.score(val_emb, val_y)
+        val_acc = clf.score(val_emb, val_y)
         test_acc = clf.score(test_emb, test_y)
 
         return val_acc, test_acc
 
     @torch.no_grad()
-    def get_embeddings(self, model):
+    def get_embeddings(
+        self,
+    ) -> Dict[Optional[str], Any]:
         """Returns Graph Embeddings for the key node only.
-          
-           Embeddings size: (n, embedding_size), where
-           n: number of nodes present for key node inside graph.
-           embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
 
-        model: Graph embedding model.
+        Embeddings size: (n, embedding_size), where
+        n: number of nodes present for key node inside graph.
+        embedding_size: Length of the node embeddings when they are mapped to
+            d-dimensional euclidean space.
+
         """
         emb = {}
-        model.eval()
-        z = model.Z.detach().cpu().numpy()
+        self.eval()
+        z = self.Z.detach().cpu().numpy()
         emb[self.key_node] = z
 
-        return emb
\ No newline at end of file
+        return emb
diff --git a/fastgraphml/graph_embeddings/models/gat.py b/fastgraphml/graph_embeddings/models/gat.py
index 83ae533..45536ed 100644
--- a/fastgraphml/graph_embeddings/models/gat.py
+++ b/fastgraphml/graph_embeddings/models/gat.py
@@ -1,83 +1,114 @@
+import shutil
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from arango.database import Database
+from sklearn.linear_model import LogisticRegression
+from torch import Tensor
+from torch.nn import Linear as Lin
 from torch_cluster import random_walk
+from torch_geometric.data import Data
 from torch_geometric.loader import NeighborSampler as RawNeighborSampler
 from torch_geometric.nn import GATConv
-from sklearn.metrics import accuracy_score
-from sklearn.linear_model import LogisticRegression
-from torch.nn import Linear as Lin
-from torch import Tensor
-import shutil
-from arango.database import Database
-import numpy as np
-from ..utils import GraphUtils
 
+from ..utils import GraphUtils
 
 # check for gpu
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 # neighborhood sampling
 class NeighborSampler(RawNeighborSampler):
-    """ For each node in batch, it sample a direct neighbor (as positive
-        example) and a random node (as negative example):
+    """For each node in batch, it sample a direct neighbor (as positive example) and a
+    random node (as negative example):
 
-        returns sampled neighborhood
+    returns sampled neighborhood
     """
 
-    def sample(self, batch):
+    def sample(self, batch: Tensor) -> Union[Tensor, Any]:
         batch = torch.tensor(batch)
         row, col, _ = self.adj_t.coo()
 
         # For each node in `batch`, we sample a direct neighbor (as positive
         # example) and a random node (as negative example):
-        pos_batch = random_walk(row, col, batch, walk_length=1,
-                                coalesced=False)[:, 1]
-        neg_batch = torch.randint(0, self.adj_t.size(1), (batch.numel(), ),
-                                dtype=torch.long)
+        pos_batch = random_walk(row, col, batch, walk_length=1, coalesced=False)[:, 1]
+        neg_batch = torch.randint(
+            0, self.adj_t.size(1), (batch.numel(),), dtype=torch.long
+        )
         batch = torch.cat([batch, pos_batch, neg_batch], dim=0)
         return super().sample(batch)
 
 
-
 class GAT(torch.nn.Module):
-    """ GATCONV model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/gat.py>'
-    
+    """GATCONV model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/gat.py>'.  # noqa: E501
+
     :database (type: Database): A python-arango database instance.
-    :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG.
-    :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define metagraph as 
-                a dictionary defining vertex & edge collections to import to PyG, along 
-                with collection-level specifications to indicate which ArangoDB attributes will become PyG features/labels. 
-                It also supports different encoders such as identity and categorical encoder on database attributes. Detailed information regarding different
-                use cases and metagraph definitons can be found on adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
-    :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs (via PyG data objects) directy rather than ArangoDB graphs.
-                When generating graph embeddings via PyG graphs, database=arango_graph=metagraph=None.
-    :embedding_size (type: int): Length of the node embeddings when they are mapped to d-dimensional euclidean space.
-    :heads (type: int): Number of attention heads. Model learns to give attention to only important nodes in node's neighborhood.
+    :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG
+    :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define
+        metagraph as a dictionary defining vertex & edge collections to import to PyG,
+        along with collection-level specifications to indicate which ArangoDB attributes
+        will become PyG features/labels. It also supports different encoders such as
+        identity and categorical encoder on database attributes. Detailed information
+        regarding different use cases and metagraph definitons can be found on
+        adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
+    :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs
+        (via PyG data objects) directy rather than ArangoDB graphs.
+        When generating graph embeddings via PyG graphs,
+        database=arango_graph=metagraph=None.
+    :embedding_size (type: int): Length of the node embeddings when they are mapped to
+        d-dimensional euclidean space.
+    :heads (type: int): Number of attention heads. Model learns to give attention to
+        only important nodes in node's neighborhood.
     :num_layers (type: int): Number of GAT Layers.
-    :sizes (type: [int, int]): Number of neighbors to select at each layer for every node (uniform random sampling) in order to perform neighborhood sampling. 
-    :batch_size (type: int): Number of nodes to be present inside batch along with their neighborhood. Used while performing neighborhood sampling.
+    :sizes (type: [int, int]): Number of neighbors to select at each layer for every
+        node (uniform random sampling) in order to perform neighborhood sampling.
+    :batch_size (type: int): Number of nodes to be present inside batch along with
+        their neighborhood. Used while performing neighborhood sampling.
     :dropout_perc (type: float): Handles overfitting inside the model
-    :shuffle (type: bool): If set to True, it shuffles data before performing neighborhood sampling.
-    :transform (type: torch_geometric.transforms): It is used to transform PyG data objects. Various transformation methods can be chained together using Compose.
+    :shuffle (type: bool): If set to True, it shuffles data before performing
+        neighborhood sampling.
+    :transform (type: torch_geometric.transforms): It is used to transform PyG data
+        objects. Various transformation methods can be chained together using Compose.
                for e.g. transform = T.Compose([
                         T.NormalizeFeatures(),
-                        T.RandomNodeSplit(num_val=0.2, num_test=0.1)])  
+                        T.RandomNodeSplit(num_val=0.2, num_test=0.1)])
     :num_val (type: float): Percentage of nodes selected for validation set.
     :num_test (type: float): Percentage of nodes selected for test set.
-    
-    Note: After selecting the percentage for validation and test nodes, rest percentage of the nodes are considered as training nodes.
 
-    :**kwargs: Additional arguments of the GATConv class. For more arguments please refer the following link
-     <https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GATConv>
+    Note: After selecting the percentage for validation and test nodes,
+    rest percentage of the nodes are considered as training nodes.
+
+    :**kwargs: Additional arguments of the GATConv class. For more arguments please
+    refer the following link
+     <https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GATConv>  # noqa: E501
     """
 
-    def __init__(self, database = None, arango_graph = None, metagraph = None, pyg_graph = None, embedding_size = 64, heads = 2,
-        num_layers = 2, sizes = [10, 10], batch_size = 256, dropout_perc = 0.5,  shuffle = True, transform = None, num_val = 0.1, num_test = 0.1, **kwargs):
+    def __init__(
+        self,
+        database: Database = None,
+        arango_graph: Optional[str] = None,
+        metagraph: Union[Dict[str, object], None] = None,
+        pyg_graph: Data = None,
+        embedding_size: int = 64,
+        heads: int = 2,
+        num_layers: int = 2,
+        sizes: List[int] = [10, 10],
+        batch_size: int = 256,
+        dropout_perc: float = 0.5,
+        shuffle: bool = True,
+        transform: Optional[List[Callable[..., Any]]] = None,
+        num_val: float = 0.1,
+        num_test: float = 0.1,
+        **kwargs: Any,
+    ):
         super().__init__()
 
-        if (database is not None or arango_graph is not None or metagraph is not None) and pyg_graph is not None:
-            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"
+        if (
+            database is not None or arango_graph is not None or metagraph is not None
+        ) and pyg_graph is not None:
+            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"  # noqa: E501
             raise Exception(msg)
 
         if database is not None:
@@ -86,7 +117,9 @@ def __init__(self, database = None, arango_graph = None, metagraph = None, pyg_g
                 raise TypeError(msg)
 
         # arango to Pyg
-        self.graph_util = GraphUtils(arango_graph, metagraph, database, pyg_graph, num_val, num_test, transform)
+        self.graph_util = GraphUtils(
+            arango_graph, metagraph, database, pyg_graph, num_val, num_test, transform
+        )
         # get PyG graph
         G = self.graph_util.graph
         self.in_channels = G.num_node_features
@@ -97,34 +130,49 @@ def __init__(self, database = None, arango_graph = None, metagraph = None, pyg_g
         self.num_nodes = G.num_nodes
         self.num_layers = num_layers
         self.sizes = sizes
-        self.batch_size = batch_size 
+        self.batch_size = batch_size
         self.dropout_perc = dropout_perc
         self.shuffle = shuffle
 
-        # create train loader 
-        self.train_loader = NeighborSampler(self.edge_index, batch_size=self.batch_size, sizes=self.sizes,
-                                            shuffle=self.shuffle, num_nodes=self.num_nodes)
-
+        # create train loader
+        self.train_loader = NeighborSampler(
+            self.edge_index,
+            batch_size=self.batch_size,
+            sizes=self.sizes,
+            shuffle=self.shuffle,
+            num_nodes=self.num_nodes,
+        )
 
         # defining GAT layers
         self.convs = nn.ModuleList()
         for i in range(num_layers):
             if i == 0:
-                self.convs.append(GATConv(self.in_channels, self.hidden_channels, heads))
+                self.convs.append(
+                    GATConv(self.in_channels, self.hidden_channels, heads, **kwargs)
+                )
             else:
-                self.convs.append(GATConv(heads * self.hidden_channels, self.hidden_channels, heads))
-            
+                self.convs.append(
+                    GATConv(
+                        heads * self.hidden_channels,
+                        self.hidden_channels,
+                        heads,
+                        **kwargs,
+                    )
+                )
+
         # adding skip connections
         self.skips = torch.nn.ModuleList()
         for i in range(num_layers):
             if i == 0:
                 self.skips.append(Lin(self.in_channels, self.hidden_channels * heads))
             else:
-                self.skips.append(Lin(self.hidden_channels * heads, self.hidden_channels * heads))
+                self.skips.append(
+                    Lin(self.hidden_channels * heads, self.hidden_channels * heads)
+                )
 
-    def forward(self, x, adjs):
+    def forward(self, x: Tensor, adjs: Tensor) -> Tensor:
         for i, (edge_index, _, size) in enumerate(adjs):
-            x_target = x[:size[1]]  # Target nodes are always placed first.
+            x_target = x[: size[1]]  # Target nodes are always placed first.
             x = self.convs[i]((x, x_target), edge_index)
             x = x + self.skips[i](x_target)
             if i != self.num_layers - 1:
@@ -132,7 +180,7 @@ def forward(self, x, adjs):
                 x = F.dropout(x, p=self.dropout_perc, training=self.training)
         return x
 
-    def full_forward(self, x, edge_index):
+    def full_forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
         for i, conv in enumerate(self.convs):
             x = conv(x, edge_index) + self.skips[i](x)
             if i != self.num_layers - 1:
@@ -142,32 +190,42 @@ def full_forward(self, x, edge_index):
 
     # save checkpoints whenever there is an increase in validation accuracy
     @staticmethod
-    def save_checkpoints(state, is_best, ckp_path, best_model_path):
+    def save_checkpoints(
+        state: Dict[str, str], is_best: bool, ckp_path: str, best_model_path: str
+    ) -> None:
         file_path = ckp_path
         torch.save(state, file_path)
-        #if it is a best model, min train loss
+        # if it is a best model, min train loss
         if is_best:
             best_file_path = best_model_path
-            # copy best checkpoint file to best model path 
+            # copy best checkpoint file to best model path
             shutil.copyfile(file_path, best_file_path)
-    
 
-    def _train(self, model, ckp_path = "./latest_model_checkpoint.pt", best_model_path = "./best_model.pt", epochs = 51, lr = 0.001, **kwargs):
+    def _train(
+        self,
+        ckp_path: str = "./latest_model_checkpoint.pt",
+        best_model_path: str = "./best_model.pt",
+        epochs: int = 51,
+        lr: float = 0.001,
+        **kwargs: Any,
+    ) -> None:
         """Train GraphML model.
 
-        :model: Graph embedding model.
-        :ckp_path (type: str): Path to save model's latest checkpoints (i.e. at each epoch). Pytorch models are saved with .pt file extension.
-         By default it saves model in cwd.
-        :best_model_path (type: str): Path to save model whenever there is an increase in validation accuracy. By default it saves model in cwd.
+        :ckp_path (type: str): Path to save model's latest checkpoints
+            (i.e. at each epoch). Pytorch models are saved with .pt file extension.
+            By default it saves model in cwd.
+        :best_model_path (type: str): Path to save model whenever there is an increase
+            in validation accuracy. By default it saves model in cwd.
         :epochs (type: int): Number of times training data go through the model.
         :lr (type: float): Learning rate.
-        :**kwargs: Additional arguments for the Adam optimizer for e.g. weight_decay, betas, etc.
+        :**kwargs: Additional arguments for the Adam optimizer for
+            e.g. weight_decay, betas, etc.
         """
 
-        model = model.to(device)
+        model = self.to(device)
         optimizer = torch.optim.Adam(model.parameters(), lr, **kwargs)
         best_acc = 0.0
-        print('Training started .........')
+        print("Training started .........")
         for epoch in range(1, epochs):
             model.train()
             total_loss = 0
@@ -187,22 +245,26 @@ def _train(self, model, ckp_path = "./latest_model_checkpoint.pt", best_model_pa
 
                 total_loss += float(loss) * out.size(0)
             # print unsupervised GraphSage loss
-            train_loss = total_loss / self.num_nodes
+            # train_loss = total_loss / self.num_nodes
 
             ##################
             # validate model #
             ##################
 
-            val_acc, test_acc = self.val(model)
+            val_acc, test_acc = self.val()
 
-            print(f'Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, 'f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')
+            print(
+                f"Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, "
+                f"Val: {val_acc:.4f}, Test: {test_acc:.4f}"
+            )
 
             # create checkpoint variable and add crucial model information
             model_checkpoint = {
-            'epoch': epoch,
-            'best_acc': val_acc,
-            'state_dict': model.state_dict(),
-            'optimizer': optimizer.state_dict(),}
+                "epoch": epoch,
+                "best_acc": val_acc,
+                "state_dict": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+            }
 
             # save current/last checkpoint
             # used for resuming training
@@ -210,23 +272,29 @@ def _train(self, model, ckp_path = "./latest_model_checkpoint.pt", best_model_pa
 
             # save model if val acc increases
             if val_acc > best_acc:
-                print('Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...'.format(best_acc, val_acc))
+                print(
+                    "Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...".format(
+                        best_acc, val_acc
+                    )
+                )
                 # save checkpoint as best model
                 self.save_checkpoints(model_checkpoint, True, ckp_path, best_model_path)
                 best_acc = val_acc
 
-    
     @torch.no_grad()
-    def val(self, model):
-        """Tests the performance of a generated graph embeddings using Node Classification as a downstream task.
-        
+    def val(
+        self,
+    ) -> Tuple[float, float]:
+        """Tests the performance of a generated graph embeddings using Node
+        Classification as a downstream task.
+
         returns validation and test accuracy.
 
         model: Graph embedding model.
         """
-        model.eval()
-        out = model.full_forward(self.x, self.edge_index).cpu()
-        clf = LogisticRegression(max_iter=400, class_weight='balanced')
+        self.eval()
+        out = self.full_forward(self.x, self.edge_index).cpu()
+        clf = LogisticRegression(max_iter=400, class_weight="balanced")
         clf.fit(out[self.G.train_mask], self.G.y[self.G.train_mask])
         val_acc = clf.score(out[self.G.val_mask], self.G.y[self.G.val_mask])
         test_acc = clf.score(out[self.G.test_mask], self.G.y[self.G.test_mask])
@@ -234,13 +302,15 @@ def val(self, model):
         return val_acc, test_acc
 
     @torch.no_grad()
-    def get_embeddings(self, model):
+    def get_embeddings(
+        self,
+    ) -> Union[Tensor, Any]:
         """Returns Graph Embeddings of size (n, embedding_size),
-           n: number of nodes in graph.
-           embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
+        n: number of nodes in graph.
+        embedding_size: Length of the node embeddings when they are mapped to
+            d-dimensional euclidean space.
 
-        model: Graph embedding model.
         """
-        model.eval()
-        emb = model.full_forward(self.x, self.edge_index).detach().cpu().numpy()
-        return emb
\ No newline at end of file
+        self.eval()
+        emb = self.full_forward(self.x, self.edge_index).detach().cpu().numpy()
+        return emb
diff --git a/fastgraphml/graph_embeddings/models/graph_sage.py b/fastgraphml/graph_embeddings/models/graph_sage.py
index d1935a1..2896090 100644
--- a/fastgraphml/graph_embeddings/models/graph_sage.py
+++ b/fastgraphml/graph_embeddings/models/graph_sage.py
@@ -1,84 +1,108 @@
-from typing import Dict, Tuple, Callable, List, Any
-from torch_geometric.data import Data
-from torch import Tensor
+import shutil
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from arango.database import Database
+from sklearn.linear_model import LogisticRegression
+from torch import Tensor
 from torch_cluster import random_walk
+from torch_geometric.data import Data
 from torch_geometric.loader import NeighborSampler as RawNeighborSampler
 from torch_geometric.nn import SAGEConv
-from sklearn.metrics import accuracy_score
-from sklearn.linear_model import LogisticRegression
-from torch import Tensor, dropout
-import shutil
-from arango.database import Database
-import numpy as np
-from ..utils import GraphUtils
 
+from ..utils import GraphUtils
 
 # check for gpu
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 # neighborhood sampling
 class NeighborSampler(RawNeighborSampler):
-    """ For each node in batch, it sample a direct neighbor (as positive
-        example) and a random node (as negative example):
+    """For each node in batch, it sample a direct neighbor (as positive example) and a
+    random node (as negative example):
 
-        returns sampled neighborhood
+    returns sampled neighborhood
     """
 
-    def sample(self, batch: int):
+    def sample(self, batch: Tensor) -> Union[Tensor, Any]:
         batch = torch.tensor(batch)
         row, col, _ = self.adj_t.coo()
 
-       
-        pos_batch = random_walk(row, col, batch, walk_length=1,
-                                coalesced=False)[:, 1]
-        neg_batch = torch.randint(0, self.adj_t.size(1), (batch.numel(), ),
-                                dtype=torch.long)
+        pos_batch = random_walk(row, col, batch, walk_length=1, coalesced=False)[:, 1]
+        neg_batch = torch.randint(
+            0, self.adj_t.size(1), (batch.numel(),), dtype=torch.long
+        )
         batch = torch.cat([batch, pos_batch, neg_batch], dim=0)
         return super().sample(batch)
 
 
 class SAGE(nn.Module):
-    """ GraphSage model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/graph_sage_unsup.py>'
-    
+    """GraphSage model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/graph_sage_unsup.py>'. # noqa: E501
+
     :database (type: Database): A python-arango database instance.
-    :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG.
-    :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define metagraph as 
-                a dictionary defining vertex & edge collections to import to PyG, along 
-                with collection-level specifications to indicate which ArangoDB attributes will become PyG features/labels. 
-                It also supports different encoders such as identity and categorical encoder on database attributes. Detailed information regarding different
-                use cases and metagraph definitons can be found on adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
-    :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs (via PyG data objects) directy rather than ArangoDB graphs.
-                When generating graph embeddings via PyG graphs, database=arango_graph=metagraph=None.
-    :embedding_size (type: int): Length of the node embeddings when they are mapped to d-dimensional euclidean space.
+    :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG
+    :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define
+        metagraph as a dictionary defining vertex & edge collections to import to PyG,
+        along with collection-level specifications to indicate which ArangoDB
+        attributes will become PyG features/labels. It also supports different encoders
+        such as identity and categorical encoder on database attributes. Detailed
+        information regarding different use cases and metagraph definitons can be found
+        on adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
+    :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs
+        (via PyG data objects) directy rather than ArangoDB graphs. When generating
+        graph embeddings via PyG graphs, database=arango_graph=metagraph=None.
+    :embedding_size (type: int): Length of the node embeddings when they are mapped to
+        d-dimensional euclidean space.
     :num_layers (type: int): Number of GraphSage Layers.
-    :sizes (type: [int, int]): Number of neighbors to select at each layer for every node (uniform random sampling) in order to perform neighborhood sampling. 
+    :sizes (type: [int, int]): Number of neighbors to select at each layer for every
+        node (uniform random sampling) in order to perform neighborhood sampling.
     If set -1 all neighbors are included in layer l.
-    :batch_size (type: int): Number of nodes to be present inside batch along with their neighborhood. Used while performing neighborhood sampling.
+    :batch_size (type: int): Number of nodes to be present inside batch along with their
+        neighborhood. Used while performing neighborhood sampling.
     :dropout_perc (type: float): Handles overfitting inside the model
-    :shuffle (type: bool): If set to True, it shuffles data before performing neighborhood sampling.
-    :transform (type: torch_geometric.transforms): It is used to transform PyG data objects. Various transformation methods can be chained together using Compose.
+    :shuffle (type: bool): If set to True, it shuffles data before performing
+        neighborhood sampling.
+    :transform (type: torch_geometric.transforms): It is used to transform PyG data
+        objects. Various transformation methods can be chained together using Compose.
                for e.g. transform = T.Compose([
                         T.NormalizeFeatures(),
-                        T.RandomNodeSplit(num_val=0.2, num_test=0.1)])         
+                        T.RandomNodeSplit(num_val=0.2, num_test=0.1)])
     :num_val (type: float): Percentage of nodes selected for validation set.
     :num_test (type: float): Percentage of nodes selected for test set.
-    
-    Note: After selecting the percentage for validation and test nodes, rest percentage of the nodes are considered as training nodes.
 
-    :**kwargs: Additional arguments of the SAGEConv class. For more arguments please refer the following link
-     <https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.SAGEConv>
+    Note: After selecting the percentage for validation and test nodes,
+        rest percentage of the nodes are considered as training nodes.
+
+    :**kwargs: Additional arguments of the SAGEConv class.
+       For more arguments please refer the following link
+     <https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.SAGEConv> # noqa: E501
     """
 
-    def __init__(self, database: Database = None, arango_graph: str = None, metagraph: Dict = None, 
-        pyg_graph: Data = None, embedding_size: int = 64, num_layers: int = 2, sizes: Tuple[int, int] = [10, 10], batch_size: int = 256, 
-        dropout_perc: float = 0.5, shuffle: bool = True, transform: List[Callable] = None, num_val: float = 0.1, num_test: float = 0.1, **kwargs):
+    def __init__(
+        self,
+        database: Database = None,
+        arango_graph: Optional[str] = None,
+        metagraph: Union[Dict[str, object], None] = None,
+        pyg_graph: Data = None,
+        embedding_size: int = 64,
+        num_layers: int = 2,
+        sizes: List[int] = [10, 10],
+        batch_size: int = 256,
+        dropout_perc: float = 0.5,
+        shuffle: bool = True,
+        transform: Optional[List[Callable[..., Any]]] = None,
+        num_val: float = 0.1,
+        num_test: float = 0.1,
+        **kwargs: Any,
+    ):
         super().__init__()
 
-        if (database is not None or arango_graph is not None or metagraph is not None) and pyg_graph is not None:
-            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"
+        if (
+            database is not None or arango_graph is not None or metagraph is not None
+        ) and pyg_graph is not None:
+            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"  # noqa: E501
             raise Exception(msg)
 
         if database is not None:
@@ -87,7 +111,9 @@ def __init__(self, database: Database = None, arango_graph: str = None, metagrap
                 raise TypeError(msg)
 
         # arango to PyG
-        self.graph_util = GraphUtils(arango_graph, metagraph, database, pyg_graph, num_val, num_test, transform)
+        self.graph_util = GraphUtils(
+            arango_graph, metagraph, database, pyg_graph, num_val, num_test, transform
+        )
         # get PyG graph
         G = self.graph_util.graph
         self.in_channels = G.num_node_features
@@ -98,20 +124,27 @@ def __init__(self, database: Database = None, arango_graph: str = None, metagrap
         self.num_nodes = G.num_nodes
         self.num_layers = num_layers
         self.sizes = sizes
-        self.batch_size = batch_size 
+        self.batch_size = batch_size
         self.dropout_perc = dropout_perc
         self.shuffle = shuffle
-        # create train loader 
-        self.train_loader = NeighborSampler(self.edge_index, batch_size=self.batch_size, sizes=self.sizes,
-                                            shuffle=self.shuffle, num_nodes=self.num_nodes)
+        # create train loader
+        self.train_loader = NeighborSampler(
+            self.edge_index,
+            batch_size=self.batch_size,
+            sizes=self.sizes,
+            shuffle=self.shuffle,
+            num_nodes=self.num_nodes,
+        )
         self.convs = nn.ModuleList()
         for i in range(num_layers):
             self.in_channels = self.in_channels if i == 0 else self.hidden_channels
-            self.convs.append(SAGEConv(self.in_channels, self.hidden_channels, **kwargs))
+            self.convs.append(
+                SAGEConv(self.in_channels, self.hidden_channels, **kwargs)
+            )
 
     def forward(self, x: Tensor, adjs: Tensor) -> Tensor:
         for i, (edge_index, _, size) in enumerate(adjs):
-            x_target = x[:size[1]]  # Target nodes are always placed first.
+            x_target = x[: size[1]]  # Target nodes are always placed first.
             x = self.convs[i]((x, x_target), edge_index)
             if i != self.num_layers - 1:
                 x = x.relu()
@@ -127,35 +160,44 @@ def full_forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
                 x = F.dropout(x, p=self.dropout_perc, training=self.training)
         return x
 
-
     # save checkpoints whenever there is an increase in validation accuracy
     @staticmethod
-    def save_checkpoints(state: Dict, is_best: bool, ckp_path: str, best_model_path: str):
+    def save_checkpoints(
+        state: Dict[str, str], is_best: bool, ckp_path: str, best_model_path: str
+    ) -> None:
         file_path = ckp_path
         torch.save(state, file_path)
-        #if it is a best model, min train loss
+        # if it is a best model, min train loss
         if is_best:
             best_file_path = best_model_path
-            # copy best checkpoint file to best model path 
+            # copy best checkpoint file to best model path
             shutil.copyfile(file_path, best_file_path)
-    
 
-    def _train(self, model: Any, ckp_path: str = "./latest_model_checkpoint.pt", best_model_path: str = "./best_model.pt", epochs: int = 51, lr: float = 0.001, **kwargs):
+    def _train(
+        self,
+        ckp_path: str = "./latest_model_checkpoint.pt",
+        best_model_path: str = "./best_model.pt",
+        epochs: int = 51,
+        lr: float = 0.001,
+        **kwargs: Any,
+    ) -> None:
         """Train GraphML model.
 
-        :model: Graph embedding model.
-        :ckp_path (type: str): Path to save model's latest checkpoints (i.e. at each epoch). Pytorch models are saved with .pt file extension.
-         By default it saves model in cwd.
-        :best_model_path (type: str): Path to save model whenever there is an increase in validation accuracy. By default it saves model in cwd.
+        :ckp_path (type: str): Path to save model's latest checkpoints
+            (i.e. at each epoch). Pytorch models are saved with .pt file extension.
+            By default it saves model in cwd.
+        :best_model_path (type: str): Path to save model whenever there is an increase
+            in validation accuracy. By default it saves model in cwd.
         :epochs (type: int): Number of times training data go through the model.
         :lr (type: float): Learning rate.
-        :**kwargs: Additional arguments for the Adam optimizer for e.g. weight_decay, betas, etc.
+        :**kwargs: Additional arguments for the Adam optimizer for
+            e.g. weight_decay, betas, etc.
         """
 
-        model = model.to(device)
+        model = self.to(device)
         optimizer = torch.optim.Adam(model.parameters(), lr, **kwargs)
         best_acc = 0.0
-        print('Training started .........')
+        print("Training started .........")
         for epoch in range(1, epochs):
             model.train()
             total_loss = 0
@@ -175,22 +217,26 @@ def _train(self, model: Any, ckp_path: str = "./latest_model_checkpoint.pt", bes
 
                 total_loss += float(loss) * out.size(0)
             # print unsupervised GraphSage loss
-            train_loss = total_loss / self.num_nodes
+            # train_loss = total_loss / self.num_nodes
 
             ##################
             # validate model #
             ##################
 
-            val_acc, test_acc = self.val(model)
+            val_acc, test_acc = self.val()
 
-            print(f'Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, 'f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')
+            print(
+                f"Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, "
+                f"Val: {val_acc:.4f}, Test: {test_acc:.4f}"
+            )
 
             # create checkpoint variable and add crucial model information
             model_checkpoint = {
-            'epoch': epoch,
-            'best_acc': val_acc,
-            'state_dict': model.state_dict(),
-            'optimizer': optimizer.state_dict(),}
+                "epoch": epoch,
+                "best_acc": val_acc,
+                "state_dict": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+            }
 
             # save current/last checkpoint
             # used for resuming training
@@ -198,23 +244,27 @@ def _train(self, model: Any, ckp_path: str = "./latest_model_checkpoint.pt", bes
 
             # save model if val acc increases
             if val_acc > best_acc:
-                print('Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...'.format(best_acc, val_acc))
+                print(
+                    "Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...".format(
+                        best_acc, val_acc
+                    )
+                )
                 # save checkpoint as best model
                 self.save_checkpoints(model_checkpoint, True, ckp_path, best_model_path)
                 best_acc = val_acc
 
-    
     @torch.no_grad()
-    def val(self, model: Any) -> Tuple[float, float]:
+    def val(
+        self,
+    ) -> Tuple[float, float]:
         """Tests the performance of a generated graph embeddings using Node Classification as a downstream task.
-        
+
         returns validation and test accuracy.
 
-        model: Graph embedding model.
-        """
-        model.eval()
-        out = model.full_forward(self.x, self.edge_index).cpu()
-        clf = LogisticRegression(max_iter=400, class_weight='balanced')
+        """  # noqa: E501
+        self.eval()
+        out = self.full_forward(self.x, self.edge_index).cpu()
+        clf = LogisticRegression(max_iter=400, class_weight="balanced")
         clf.fit(out[self.G.train_mask], self.G.y[self.G.train_mask])
         val_acc = clf.score(out[self.G.val_mask], self.G.y[self.G.val_mask])
         test_acc = clf.score(out[self.G.test_mask], self.G.y[self.G.test_mask])
@@ -222,14 +272,15 @@ def val(self, model: Any) -> Tuple[float, float]:
         return val_acc, test_acc
 
     @torch.no_grad()
-    def get_embeddings(self, model: Any) -> Tensor:
+    def get_embeddings(
+        self,
+    ) -> Union[Tensor, Any]:
         """Returns Graph Embeddings of size (n, embedding_size),
-           n: number of nodes in graph.
-           embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
+        n: number of nodes in graph.
+        embedding_size: Length of the node embeddings when they are mapped to
+            d-dimensional euclidean space.
 
-        model: Graph embedding model.
         """
-        model.eval()
-        emb = model.full_forward(self.x, self.edge_index).detach().cpu().numpy()
+        self.eval()
+        emb = self.full_forward(self.x, self.edge_index).detach().cpu().numpy()
         return emb
-
diff --git a/fastgraphml/graph_embeddings/models/metapath2vec.py b/fastgraphml/graph_embeddings/models/metapath2vec.py
index 50e9497..13cdefd 100644
--- a/fastgraphml/graph_embeddings/models/metapath2vec.py
+++ b/fastgraphml/graph_embeddings/models/metapath2vec.py
@@ -1,71 +1,114 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch_geometric.transforms as T
-from torch_geometric.nn import MetaPath2Vec
-from sklearn.linear_model import LogisticRegression
-from ..utils import GraphUtils
 import shutil
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
 from arango.database import Database
+from sklearn.linear_model import LogisticRegression
+from torch import Tensor
+from torch_geometric.data import Data
+from torch_geometric.nn import MetaPath2Vec
+from torch_geometric.typing import EdgeType
 
+from ..utils import GraphUtils
 
 # check for gpu
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 
 class METAPATH2VEC:
-    """ metapath2vec model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/hetero/metapath2vec.py>'
-    
-        :database (type: Database): A python-arango database instance.
-        :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG.
-        :metagraph (type: Dict): It exports ArangoDB graphs to PyG data objects. We define metagraph as 
-                    a dictionary defining vertex & edge collections to import to PyG, along 
-                    with collection-level specifications to indicate which ArangoDB attributes will become PyG features/labels. 
-                    It also supports different encoders such as identity and categorical encoder on database attributes. Detailed information regarding different
-                    use cases and metagraph definitons can be found on adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
-        :metapaths (type: list[Tuple(str,str,str)]): The metapath defined as (src_node_type, rel_type, dst_node_type) tuples. M2V uses metapaths to 
-        perform random walks on the graph and then uses skip-grapm to compute graph embeddings.
-        :key_node (type: str): Node type on which we want to test the performance of generated graph embeddings. Performance is tested using node classification task.
-        :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs (via PyG data objects) directy rather than ArangoDB graphs.
-                    When generating graph embeddings via PyG graphs, database=arango_graph=metagraph=None.
-        :embedding_size (type: int): Length of the node embeddings when they are mapped to d-dimensional euclidean space.
-        :walk_length (type: int): The walk length. 
-        :context_size (type: int): The actual context size which is considered for positive samples. 
-                      This parameter increases the effective sampling rate by reusing samples across different source nodes.
-        :walks_per_node (type: float): The number of walks to sample for each node.
-        :num_negative_samples (type: bool): The number of negative samples to use for each positive sample.
-        :num_nodes_dict (type: Dict): Dictionary holding the number of nodes for each node type.
-        :sparse (type: bool): If set to True, gradients w.r.t. to the weight matrix will be sparse.
-        :transform (type: torch_geometric.transforms): It is used to transform PyG data objects. Various transformation methods can be chained together using Compose.
-                   for e.g. transform = T.Compose([
-                            T.NormalizeFeatures(),
-                            T.RandomNodeSplit(num_val=0.2, num_test=0.1)]).
-        :num_val (type: float): Percentage of nodes selected for validation set.
-        :num_test (type: float): Percentage of nodes selected for test set.
-        :shuffle (type: bool): If set to True, it shuffles data before training.
-        
-        Note: After selecting the percentage for validation and test nodes, rest percentage of the nodes are considered as training nodes.
-
-        :Detailed information about Metapath2Vec args can be found in
-        <https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.models.MetaPath2Vec>
-        """
+    """metapath2vec model modified from '<https://github.com/pyg-team/pytorch_geometric/blob/6267de93c6b04f46a306aa58e414de330ef9bb10/examples/hetero/metapath2vec.py>'.  # noqa: E501
+
+    :database (type: Database): A python-arango database instance.
+    :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG
+    :metagraph (type: Dict): It exports ArangoDB graphs to PyG data objects. We define
+        metagraph as a dictionary defining vertex & edge collections to import to PyG,
+        along with collection-level specifications to indicate which ArangoDB attributes
+        will become PyG features/labels. It also supports different encoders such as
+        identity and categorical encoder on database attributes. Detailed information
+        regarding different use cases and metagraph definitons can be found on
+        adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
+    :metapaths (type: list[Tuple(str,str,str)]): The metapath defined as
+        (src_node_type, rel_type, dst_node_type) tuples. M2V uses metapaths to
+        perform random walks on the graph and then uses skip-grapm to compute
+        graph embeddings.
+    :key_node (type: str): Node type on which we want to test the performance of
+        generated graph embeddings. Performance is tested using node classification task.
+    :pyg_graph (type: PyG data object): It generates graph embeddings using PyG graphs
+        (via PyG data objects) directy rather than ArangoDB graphs. When generating graph
+        embeddings via PyG graphs, database=arango_graph=metagraph=None.
+    :embedding_size (type: int): Length of the node embeddings when they are mapped to
+        d-dimensional euclidean space.
+    :walk_length (type: int): The walk length.
+    :context_size (type: int): The actual context size which is considered for positive
+        samples. This parameter increases the effective sampling rate by reusing samples
+        across different source nodes.
+    :walks_per_node (type: float): The number of walks to sample for each node.
+    :num_negative_samples (type: bool): The number of negative samples to use for each
+        positive sample.
+    :num_nodes_dict (type: Dict): Dictionary holding the number of nodes for each
+        node type.
+    :sparse (type: bool): If set to True, gradients w.r.t. to the weight matrix will be
+        sparse.
+    :transform (type: torch_geometric.transforms): It is used to transform PyG data
+        objects. Various transformation methods can be chained together using Compose.
+               for e.g. transform = T.Compose([
+                        T.NormalizeFeatures(),
+                        T.RandomNodeSplit(num_val=0.2, num_test=0.1)]).
+    :num_val (type: float): Percentage of nodes selected for validation set.
+    :num_test (type: float): Percentage of nodes selected for test set.
+    :shuffle (type: bool): If set to True, it shuffles data before training.
+
+    Note: After selecting the percentage for validation and test nodes,
+    rest percentage of the nodes are considered as training nodes.
+
+    :Detailed information about Metapath2Vec args can be found in
+    <https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.models.MetaPath2Vec>  # noqua: E501
+    """
 
-    def __init__(self, database = None, arango_graph = None, metagraph = None, metapaths = None, key_node=None,
-                pyg_graph=None, embedding_size=64, walk_length=5,
-                context_size=6, walks_per_node=5, num_negative_samples=5,
-                num_nodes_dict=None, sparse=False, batch_size=64, transform = None,
-                num_val=0.1, num_test=0.1, shuffle=True):
-        
-        if (database is not None or arango_graph is not None or metagraph is not None) and pyg_graph is not None:
-            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"
+    def __init__(
+        self,
+        database: Database = None,
+        arango_graph: Optional[str] = None,
+        metagraph: Union[Dict[str, object], None] = None,
+        metapaths: List[EdgeType] = [0],
+        key_node: Union[str, None] = None,
+        pyg_graph: Data = None,
+        embedding_size: int = 64,
+        walk_length: int = 5,
+        context_size: int = 6,
+        walks_per_node: int = 5,
+        num_negative_samples: int = 5,
+        num_nodes_dict: Optional[Dict[str, int]] = None,
+        sparse: bool = False,
+        batch_size: int = 64,
+        transform: Optional[List[Callable[..., Any]]] = None,
+        num_val: float = 0.1,
+        num_test: float = 0.1,
+        shuffle: bool = True,
+    ):
+
+        if (
+            database is not None or arango_graph is not None or metagraph is not None
+        ) and pyg_graph is not None:
+            msg = "when generating graph embeddings via PyG data objects, database=arango_graph=metagraph=None and vice versa"  # noqa: E501
             raise Exception(msg)
 
         if database is not None:
             if not issubclass(type(database), Database):
                 msg = "**db** parameter must inherit from arango.database.Database"
                 raise TypeError(msg)
-        
+
         # arango tp pyg
-        self.graph_util = GraphUtils(arango_graph, metagraph, database, pyg_graph, num_val, num_test, transform, key_node)
+        self.graph_util = GraphUtils(
+            arango_graph,
+            metagraph,
+            database,
+            pyg_graph,
+            num_val,
+            num_test,
+            transform,
+            key_node,
+        )
         # get PyG graph
         G = self.graph_util.graph
         self.G = G
@@ -87,82 +130,119 @@ def __init__(self, database = None, arango_graph = None, metagraph = None, metap
 
     # save checkpoints whenever there is an increase in validation accuracy
     @staticmethod
-    def save_checkpoints(state, is_best, ckp_path, best_model_path):
+    def save_checkpoints(
+        state: Dict[str, str], is_best: bool, ckp_path: str, best_model_path: str
+    ) -> None:
         file_path = ckp_path
         torch.save(state, file_path)
-        #if it is a best model, min train loss
+        # if it is a best model, min train loss
         if is_best:
             best_file_path = best_model_path
-            # copy best checkpoint file to best model path 
+            # copy best checkpoint file to best model path
             shutil.copyfile(file_path, best_file_path)
-    
-    
-    def _train(self, ckp_path ="./latest_model_checkpoint.pt", best_model_path = "./best_model.pt", 
-    epochs=51, lr=0.03, log_steps=100, eval_steps=2000, **kwargs):
+
+    def _train(
+        self,
+        ckp_path: str = "./latest_model_checkpoint.pt",
+        best_model_path: str = "./best_model.pt",
+        epochs: int = 51,
+        lr: float = 0.03,
+        log_steps: int = 100,
+        eval_steps: int = 2000,
+        **kwargs: Any,
+    ) -> None:
         """Train GraphML model.
 
-        :ckp_path (type: str): Path to save model's latest checkpoints (i.e. at each epoch). Pytorch models are saved with .pt file extension.
-         By default it saves model in cwd.
-        :best_model_path (type: str): Path to save model whenever there is an increase in validation accuracy. By default it saves model in cwd.
+        :ckp_path (type: str): Path to save model's latest checkpoints
+            (i.e. at each epoch). Pytorch models are saved with .pt file extension.
+            By default it saves model in cwd.
+        :best_model_path (type: str): Path to save model whenever there is an increase
+            in validation accuracy. By default it saves model in cwd.
         :epochs (type: int): Number of times training data go through the model.
         :lr (type: float): Learning rate.
         : log_steps (type: int): Logs epochs and loss.
-        : eval_steps (type: int): Evaluate model performance using validation and test data.
-        :**kwargs: Additional arguments for the Adam optimizer for e.g. weight_decay, betas, etc.
+        : eval_steps (type: int): Evaluate model performance using validation and
+            test data.
+        :**kwargs: Additional arguments for the Adam optimizer for
+            e.g. weight_decay, betas, etc.
         """
 
         # model
-        model = MetaPath2Vec(edge_index_dict=self.edge_index_dict, embedding_dim=self.embeddin_size, metapath=self.metapath, walk_length=self.walk_length,
-                            context_size=self.context_size, walks_per_node=self.walks_per_node, num_negative_samples=self.num_negative_samples, 
-                            num_nodes_dict=self.num_node_dict, sparse=self.sparse)
-        
+        model = MetaPath2Vec(
+            edge_index_dict=self.edge_index_dict,
+            embedding_dim=self.embeddin_size,
+            metapath=self.metapath,
+            walk_length=self.walk_length,
+            context_size=self.context_size,
+            walks_per_node=self.walks_per_node,
+            num_negative_samples=self.num_negative_samples,
+            num_nodes_dict=self.num_node_dict,
+            sparse=self.sparse,
+        )
+
         # transfer model to gpu
         model = model.to(device)
-        # returns data loader that creates both positive and negative random walks on the heterogeneous graph.
+        # returns data loader that creates both positive and negative
+        # random walks on the heterogeneous graph.
         loader = model.loader(batch_size=self.batch_size, shuffle=self.shuffle)
-        if self.sparse == True:
-            optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=lr, **kwargs)
+        if self.sparse is True:
+            optimizer = torch.optim.SparseAdam(
+                list(model.parameters()), lr=lr, **kwargs
+            )
         else:
             optimizer = torch.optim.Adam(model.parameters(), lr=lr, **kwargs)
         best_acc = 0.0
-        print('Training started .........')
-        
+        print("Training started .........")
+
         for epoch in range(1, epochs):
             model.train()
             total_loss = 0
             for i, (pos_rw, neg_rw) in enumerate(loader):
                 optimizer.zero_grad()
-                
+
                 loss = model.loss(pos_rw.to(device), neg_rw.to(device))
-                
+
                 loss.backward()
                 optimizer.step()
 
                 total_loss += loss.item()
                 if (i + 1) % log_steps == 0:
-                    print((f'Epoch: {epoch}, Step: {i + 1:05d}/{len(loader)}, '
-                        f'Loss: {total_loss / log_steps:.4f}'))
+                    print(
+                        (
+                            f"Epoch: {epoch}, Step: {i + 1:05d}/{len(loader)}, "
+                            f"Loss: {total_loss / log_steps:.4f}"
+                        )
+                    )
                     total_loss = 0
 
                 if (i + 1) % eval_steps == 0:
                     val_acc, test_acc = self.val(model)
-                    print((f'Epoch: {epoch}, Step: {i + 1:05d}/{len(loader)}, 'f'Val_Acc: {val_acc:.4f}', 
-                    f'Test_Acc:{test_acc:.4f}'))
-                
+                    print(
+                        (
+                            f"Epoch: {epoch}, Step: {i + 1:05d}/{len(loader)}, "
+                            f"Val_Acc: {val_acc:.4f}",
+                            f"Test_Acc:{test_acc:.4f}",
+                        )
+                    )
+
             ##################
             # validate model #
             ##################
 
             val_acc, test_acc = self.val(model)
 
-            print(f'Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, 'f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')
+            print(
+                f"Epoch: {epoch:03d}, Train_Loss: {loss:.4f}, "
+                f"Val: {val_acc:.4f}, Test: {test_acc:.4f}"
+            )
 
             # create checkpoint variable and add crucial model information
             model_checkpoint = {
-            'epoch': epoch,
-            'best_acc': val_acc,
-            'state_dict': model.state_dict(),
-            'optimizer': optimizer.state_dict(),}
+                "epoch": epoch,
+                "best_acc": val_acc,
+                "state_dict": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+            }
 
             # save current/last checkpoint
             # used for resuming training
@@ -170,16 +250,21 @@ def _train(self, ckp_path ="./latest_model_checkpoint.pt", best_model_path = "./
 
             # save model if val acc increases
             if val_acc > best_acc:
-                print('Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...'.format(best_acc, val_acc))
+                print(
+                    "Val Acc increased ({:.5f} --> {:.5f}).  Saving model ...".format(
+                        best_acc, val_acc
+                    )
+                )
                 # save checkpoint as best model
                 self.save_checkpoints(model_checkpoint, True, ckp_path, best_model_path)
                 best_acc = val_acc
 
     @torch.no_grad()
-    def val(self, model):
-        """Tests the performance of a generated graph embeddings using Node Classification as a downstream task.
-        Performance is tested on a node type (associated with labels) mentioned using key_node argument. 
-        
+    def val(self, model: Any) -> Tuple[float, float]:
+        """Tests the performance of a generated graph embeddings using Node
+        Classification as a downstream task. Performance is tested on a node type
+        (associated with labels) mentioned using key_node argument.
+
         returns validation and test accuracy on key node.
 
         model: Graph embedding model.
@@ -191,33 +276,47 @@ def val(self, model):
         train_node_idx = self.G[self.key_node].train_mask
         val_node_idx = self.G[self.key_node].val_mask
         test_node_idx = self.G[self.key_node].test_mask
-        clf = LogisticRegression(max_iter=400, class_weight='balanced')
-        clf.fit(z[train_node_idx].detach().cpu().numpy(), y[train_node_idx].detach().cpu().numpy())
-        val_acc = clf.score(z[val_node_idx].detach().cpu().numpy(), y[val_node_idx].detach().cpu().numpy())
-        test_acc = clf.score(z[test_node_idx].detach().cpu().numpy(), y[test_node_idx].detach().cpu().numpy())
+        clf = LogisticRegression(max_iter=400, class_weight="balanced")
+        clf.fit(
+            z[train_node_idx].detach().cpu().numpy(),
+            y[train_node_idx].detach().cpu().numpy(),
+        )
+        val_acc = clf.score(
+            z[val_node_idx].detach().cpu().numpy(),
+            y[val_node_idx].detach().cpu().numpy(),
+        )
+        test_acc = clf.score(
+            z[test_node_idx].detach().cpu().numpy(),
+            y[test_node_idx].detach().cpu().numpy(),
+        )
 
         self.m2v_model = model
         return val_acc, test_acc
 
     @torch.no_grad()
-    def get_embeddings(self, ):
-        """Returns Graph Embeddings as a dictionary {node_type1:emb1, node_type2:emb2, ....} 
-           where embeddings for each node type is present, if that node type is used in metapath. 
-           Embeddings size: (n, embedding_size), where
-           n: number of nodes present for specific type inside graph.
-           embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
+    def get_embeddings(
+        self,
+    ) -> Dict[str, Tensor]:
+        """Returns Graph Embeddings as a dictionary {node_type1:emb1, node_type2:emb2,
+
+        ....} where embeddings for each node type is present, if that node type is used
+        in metapath.
+
+        Embeddings size: (n, embedding_size), where
+        n: number of nodes present for specific type inside graph.
+        embedding_size: Length of the node embeddings when they are mapped to
+            d-dimensional euclidean space.
         """
 
         emb = {}
         metapath_nodes = []
         for idx in range(len(self.metapath)):
-            src_node, _ , dest_node = self.metapath[idx]
+            src_node, _, dest_node = self.metapath[idx]
             metapath_nodes.append(src_node)
             metapath_nodes.append(dest_node)
-        
+
         metapath_nodes = list(set(metapath_nodes))
         for node in metapath_nodes:
             emb[node] = self.m2v_model(node).detach().cpu().numpy()
-        
-        return emb
 
+        return emb
diff --git a/fastgraphml/graph_embeddings/utils.py b/fastgraphml/graph_embeddings/utils.py
index b95dcc2..9143436 100644
--- a/fastgraphml/graph_embeddings/utils.py
+++ b/fastgraphml/graph_embeddings/utils.py
@@ -1,21 +1,36 @@
 # utils file
-import torch.nn as nn
-from arango import ArangoClient
-import torch
-from torch_geometric.data import Data
-from tqdm import tqdm
+from typing import Any, Dict, List, Optional, Union
+
 import numpy as np
+import numpy.typing as npt
+import torch.nn as nn
 import torch_geometric.transforms as T
 from adbpyg_adapter import ADBPyG_Adapter
-from rich.progress import track    
+from arango.database import Database
+from arango.exceptions import CollectionCreateError
+from rich.progress import track
+from torch_geometric.data import Data
+from torch_geometric.typing import EdgeType
+
 from .downstream_tasks.similarity_search import similarity_search
-        
 
-# various Graph ML utilites 
+
+# various Graph ML utilites
 class GraphUtils(nn.Module):
-    """Various graph utility hepler methods"""
-    def __init__(self, arango_graph, metagraph, database, pyg_graph,
-        num_val, num_test, transform, key_node=None, metapaths=None):
+    """Various graph utility hepler methods."""
+
+    def __init__(
+        self,
+        arango_graph: Optional[str],
+        metagraph: Union[Dict[str, object], None],
+        database: Database,
+        pyg_graph: Data,
+        num_val: float,
+        num_test: float,
+        transform: Any,
+        key_node: Union[str, None] = None,
+        metapaths: Optional[List[EdgeType]] = None,
+    ):
         super().__init__()
         self.num_val = num_val
         self.num_test = num_test
@@ -32,58 +47,74 @@ def __init__(self, arango_graph, metagraph, database, pyg_graph,
             self.graph = self.pyg_preprocess(pyg_graph)
             print(self.graph_stats())
 
-    def pyg_preprocess(self, pyg_data):
-        """Takes PyG data object and preprocess it. By default it performs following preporocessing:
+    def pyg_preprocess(self, pyg_data: Data) -> Data:
+        """Takes PyG data object and preprocess it.
+        By default it performs following preporocessing:
         1. AddMetaPaths
         2. Perform Random Node Split i.e. splitting data into train, val and test set.
-        
+
         :pyg_data (tpye: PyG data object): PyG data object
         """
         # hetero graph
-        if self.key_node is not None and self.metapaths is not None:     
-            if hasattr(pyg_data[self.key_node], 'train_mask'):
+        if self.key_node is not None and self.metapaths is not None:
+            if hasattr(pyg_data[self.key_node], "train_mask"):
                 pyg_data = pyg_data
-                pyg_data = T.AddMetaPaths(self.metapaths, drop_orig_edges=True, drop_unconnected_nodes=True)(pyg_data)
+                pyg_data = T.AddMetaPaths(
+                    self.metapaths, drop_orig_edges=True, drop_unconnected_nodes=True
+                )(pyg_data)
                 if self.transform is not None:
                     pyg_data = self.transform(pyg_data)
             else:
-                pyg_data = T.RandomNodeSplit(num_val=self.num_val, num_test=self.num_test)(pyg_data)
-                pyg_data = T.AddMetaPaths(self.metapaths, drop_orig_edges=True, drop_unconnected_nodes=True)(pyg_data)
+                pyg_data = T.RandomNodeSplit(
+                    num_val=self.num_val, num_test=self.num_test
+                )(pyg_data)
+                pyg_data = T.AddMetaPaths(
+                    self.metapaths, drop_orig_edges=True, drop_unconnected_nodes=True
+                )(pyg_data)
                 if self.transform is not None:
                     pyg_data = self.transform(pyg_data)
-        
+
         elif self.key_node is not None and self.metapaths is None:
-            if hasattr(pyg_data[self.key_node], 'train_mask'):
+            if hasattr(pyg_data[self.key_node], "train_mask"):
                 pyg_data = pyg_data
                 if self.transform is not None:
                     pyg_data = self.transform(pyg_data)
             else:
-                pyg_data = T.RandomNodeSplit(num_val=self.num_val, num_test=self.num_test)(pyg_data)
+                pyg_data = T.RandomNodeSplit(
+                    num_val=self.num_val, num_test=self.num_test
+                )(pyg_data)
                 if self.transform is not None:
                     pyg_data = self.transform(pyg_data)
         # homo graph
         else:
-            if hasattr(pyg_data, 'train_mask'):
+            if hasattr(pyg_data, "train_mask"):
                 pyg_data = pyg_data
                 if self.transform is not None:
                     pyg_data = self.transform(pyg_data)
             else:
-                pyg_data = T.RandomNodeSplit(num_val=self.num_val, num_test=self.num_test)(pyg_data)
+                pyg_data = T.RandomNodeSplit(
+                    num_val=self.num_val, num_test=self.num_test
+                )(pyg_data)
                 if self.transform is not None:
                     pyg_data = self.transform(pyg_data)
-       
-        return pyg_data
 
+        return pyg_data
 
-    def arango_to_pyg(self, arango_graph, metagraph):
+    def arango_to_pyg(
+        self, arango_graph: Optional[str], metagraph: Union[Dict[str, object], None]
+    ) -> Data:
         """Exports ArangoDB graph to PyG data object using ArangoDB PyG Adapter.
-        
-        :arango_graph (type: str): The name of ArangoDB graph which we want to export to PyG.
-        :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects. We define metagraph as 
-            a dictionary defining vertex & edge collections to import to PyG, along 
-            with collection-level specifications to indicate which ArangoDB attributes will become PyG features/labels. 
-            It also supports different encoders such as identity and categorical encoder on database attributes. Detailed information regarding different
-            use cases and metagraph definitons can be found on adbpyg_adapter github page i.e <https://github.com/arangoml/pyg-adapter>.
+
+        :arango_graph (type: str): The name of ArangoDB graph which we want to
+        export to PyG.
+        :metagraph (type: dict): It exports ArangoDB graphs to PyG data objects.
+            We define metagraph as a dictionary defining vertex & edge collections
+            to import to PyG, along with collection-level specifications to
+            indicate which ArangoDB attributes will become PyG features/labels.
+            It also supports different encoders such as identity and categorical
+            encoder on database attributes. Detailed information regarding different
+            use cases and metagraph definitons can be found on adbpyg_adapter github
+            page i.e <https://github.com/arangoml/pyg-adapter>.
         """
 
         adbpyg = ADBPyG_Adapter(self.database)
@@ -91,150 +122,186 @@ def arango_to_pyg(self, arango_graph, metagraph):
         pyg_obj = self.pyg_preprocess(pyg_obj)
         return pyg_obj
 
-    # print graph statistics information about input graph for e.g. num of nodes, edges, etc.
-    def graph_stats(self,):
+    # print graph statistics information about input graph for
+    # e.g. num of nodes, edges, etc.
+    def graph_stats(
+        self,
+    ) -> Dict[str, Union[str, int, float]]:
         if self.key_node is None:
             print("Homogeneous Graph Detected ........ \n")
             g_info = {}
-            g_info['Nodes'] = self.graph.num_nodes
-            g_info['Node_feature_size'] = self.graph.num_node_features
-            g_info['Number_of_classes'] = self.graph.y.unique().size(0)
-            g_info['Edges'] = self.graph.num_edges
-            if 'edge_attr' in self.graph:
+            g_info["Nodes"] = self.graph.num_nodes
+            g_info["Node_feature_size"] = self.graph.num_node_features
+            g_info["Number_of_classes"] = self.graph.y.unique().size(0)
+            g_info["Edges"] = self.graph.num_edges
+            if "edge_attr" in self.graph:
                 e_size = self.graph.edge_attr.size(1)
-                g_info['Edge_feature_fize'] = e_size
+                g_info["Edge_feature_fize"] = e_size
             else:
-                g_info['Edge_feature_fize'] = None
+                g_info["Edge_feature_fize"] = None
             if self.graph.is_directed():
-                g_info['Graph Directionality'] = "Directed"
+                g_info["Graph Directionality"] = "Directed"
             else:
-                g_info['Graph Directionality'] = "Undirected"
-            g_info['Average node degree'] = f'{self.graph.num_edges/ self.graph.num_nodes:.2f}'
-            g_info['Number of training nodes'] = self.graph.train_mask.sum().item()
-            g_info['Number of val nodes'] = self.graph.val_mask.sum().item()
-            g_info['Number of test nodes'] = self.graph.test_mask.sum().item()
-            g_info['Has isolated nodes'] = self.graph.has_isolated_nodes()
+                g_info["Graph Directionality"] = "Undirected"
+            g_info[
+                "Average node degree"
+            ] = f"{self.graph.num_edges/ self.graph.num_nodes:.2f}"
+            g_info["Number of training nodes"] = self.graph.train_mask.sum().item()
+            g_info["Number of val nodes"] = self.graph.val_mask.sum().item()
+            g_info["Number of test nodes"] = self.graph.test_mask.sum().item()
+            g_info["Has isolated nodes"] = self.graph.has_isolated_nodes()
             return g_info
         else:
             print("Heterogeneous Graph Detected .......... \n")
             g_info = {}
             node_types = self.graph.metadata()[0]
             edge_types = self.graph.metadata()[1]
-            
+
             # adding node information
             node_data = {}
             for node in node_types:
                 total_nodes = 0
                 total_nodes = self.graph[node].num_nodes
                 if total_nodes is None:
-                    # counting documents in ArangoDB collections and adding it to num_nodes attribute
+                    # counting documents in ArangoDB collections and adding it
+                    # to num_nodes attribute
                     total_nodes = self.database.collection(node).count()
                     self.graph[node].num_nodes = total_nodes
-                if hasattr(self.graph[node], 'train_mask'):
-                    node_data['Number of '+ node +' nodes'] = total_nodes
-                    node_data['Number of train '+ node +' nodes'] = self.graph[node].train_mask.sum().item()
-                    node_data['Number of val '+ node +' nodes'] = self.graph[node].val_mask.sum().item()
-                    node_data['Number of test '+ node +' nodes'] = self.graph[node].test_mask.sum().item()
-                    node_data['Number of classes in '+ node +' nodes'] = self.graph[node].y.unique().size(0)
+                if hasattr(self.graph[node], "train_mask"):
+                    node_data["Number of " + node + " nodes"] = total_nodes
+                    node_data["Number of train " + node + " nodes"] = (
+                        self.graph[node].train_mask.sum().item()
+                    )
+                    node_data["Number of val " + node + " nodes"] = (
+                        self.graph[node].val_mask.sum().item()
+                    )
+                    node_data["Number of test " + node + " nodes"] = (
+                        self.graph[node].test_mask.sum().item()
+                    )
+                    node_data["Number of classes in " + node + " nodes"] = (
+                        self.graph[node].y.unique().size(0)
+                    )
                 else:
-                    node_data['number of '+ node +' nodes'] = total_nodes
-            
-            # add total nodes        
-            g_info['Nodes'] = self.graph.num_nodes
+                    node_data["number of " + node + " nodes"] = total_nodes
+
+            # add total nodes
+            g_info["Nodes"] = self.graph.num_nodes
             # add total edges
-            g_info['Edges'] = self.graph.num_edges
+            g_info["Edges"] = self.graph.num_edges
             # add all node types
-            g_info['node_types'] = node_types
+            g_info["node_types"] = node_types
             # add all edge_types
-            g_info['edge_types'] = edge_types
+            g_info["edge_types"] = edge_types
             # check graph directionality
             if self.graph.is_directed():
-                g_info['Graph Directionality'] = "Directed"
+                g_info["Graph Directionality"] = "Directed"
             else:
-                g_info['Graph Directionality'] = "Undirected"
+                g_info["Graph Directionality"] = "Undirected"
             # check for isolated nodes
-            g_info['Has isolated nodes'] = self.graph.has_isolated_nodes()
-            g_info['node stats'] = node_data
+            g_info["Has isolated nodes"] = self.graph.has_isolated_nodes()
+            g_info["node stats"] = node_data
             return g_info
-            
-    
+
     # store graph embeddings inside arangodb
-    def store_embeddings(self, graph_emb, collection_name=None, batch_size=100, class_mapping=None, 
-        node_type=None, nearest_nbors_search=False, top_k_nbors=10, nlist=10, search_type='exact'):
+    def store_embeddings(
+        self,
+        graph_emb: Union[npt.NDArray[np.float64], Any],
+        collection_name: str,
+        batch_size: int = 100,
+        class_mapping: Optional[Dict[int, str]] = None,
+        node_type: Optional[str] = None,
+        nearest_nbors_search: bool = False,
+        top_k_nbors: int = 10,
+        nlist: int = 10,
+        search_type: str = "exact",
+    ) -> None:
         """Store generated graph embeddings inside ArangoDB.
 
         :graph_emb (type: 2D numpy array): Numpy array of size (n, embedding_size),
         n: number of nodes in graph
-        embedding_size: Length of the node embeddings when they are mapped to d-dimensional euclidean space.
-        :collection_name (type: str): Name of the document collection where we want to store graph embeddings.
+        embedding_size: Length of the node embeddings when they are mapped to
+            d-dimensional euclidean space.
+        :collection_name (type: str): Name of the document collection where we want to
+            store graph embeddings.
         :batch_size (type: int): Batch size.
-        :class_mapping (type: dict): It is a dictionary where class names are mapped to integer labels.
-        If class_mappings are not provided the method uses integers labels as legend inside the figure. for e.g.
-        {0: 'Desktops',1: 'Data Storage',2: 'Laptops',3: 'Monitors',4: 'Computer Components',
-        5: 'Video Projectors',6: 'Routers',7: 'Tablets',8: 'Networking Products',9: 'Webcams'} # amazon computer dataset.
-        :node_type (type: str):  Node type for which we want to store embeddings. Used to store Hetero graph embeddings.
-        :nearest_nbors_search (type: bool): If nearest_nbors_search=True, store_embeddings method saves generated Graph Embeddings in ArangoDB 
-        along with top_k nearest neighbors (node ids with similar embeddings) and their corresponding similarity scores (i.e. cosine distance).
+        :class_mapping (type: dict): It is a dictionary where class names are mapped
+            to integer labels.
+        If class_mappings are not provided the method uses integers labels as legend
+        inside the figure. for e.g.
+        {0: 'Desktops',1: 'Data Storage',2: 'Laptops',3: 'Monitors',
+        4: 'Computer Components', 5: 'Video Projectors',6: 'Routers',7: 'Tablets',
+        8: 'Networking Products',9: 'Webcams'} # amazon computer dataset.
+        :node_type (type: str):  Node type for which we want to store embeddings.
+            Used to store Hetero graph embeddings.
+        :nearest_nbors_search (type: bool): If nearest_nbors_search=True,
+            store_embeddings method saves generated Graph Embeddings in ArangoDB
+            along with top_k nearest neighbors (node ids with similar embeddings)
+            and their corresponding similarity scores (i.e. cosine distance).
         :top_k_nbors (type: int): Returns top-k nearest neighbors of all embeddings.
         :nlist (type: int): Number of clusters to be generated.
         :search_type (type: str): We support two types of search for now:
         1. exact search: For precise similarity search but at the cost of scalability.
-        2. approx search: For scalable similarity search but at the cost of some precision loss.
+        2. approx search: For scalable similarity search but at the cost of some
+            precision loss.
         """
 
-        assert collection_name is not None, "pass arangodb collection name to store embeddings "
+        if collection_name is None:
+            raise Exception("Pass arangodb collection name to store embeddings.")
 
         # create document collection with name "collection_name" in arangodb
         if not self.database.has_collection(collection_name):
-            self.database.create_collection(collection_name, replication_factor=3)
-        
+            try:
+                self.database.create_collection(collection_name, replication_factor=3)
+            except CollectionCreateError as exec:
+                raise CollectionCreateError(exec)
+
         batch = []
         batch_idx = 1
         index = 0
         emb_collection = self.database[collection_name]
 
-        if nearest_nbors_search == True:
+        if nearest_nbors_search is True:
             dist, nbors = similarity_search(graph_emb)
 
         for idx in track(range(graph_emb.shape[0])):
             insert_doc = {}
-            insert_doc["_id"] = collection_name + '/' + str(idx)
+            insert_doc["_id"] = collection_name + "/" + str(idx)
             insert_doc["embedding"] = graph_emb[idx].tolist()
-            ## add class names
+            # add class names
             if class_mapping is not None and self.key_node is None:
-                insert_doc['label'] = self.graph.y[idx].item()
-                insert_doc['class_name'] = class_mapping[self.graph.y[idx].item()]
+                insert_doc["label"] = self.graph.y[idx].item()
+                insert_doc["class_name"] = class_mapping[self.graph.y[idx].item()]
 
             elif class_mapping is not None and self.key_node is not None:
-                if hasattr(self.graph[node_type], 'y'):
-                    insert_doc['label'] = self.graph[node_type].y[idx].item()
-                    insert_doc['class_name'] = class_mapping[self.graph[node_type].y[idx].item()]
+                if hasattr(self.graph[node_type], "y"):
+                    insert_doc["label"] = self.graph[node_type].y[idx].item()
+                    insert_doc["class_name"] = class_mapping[
+                        self.graph[node_type].y[idx].item()
+                    ]
                 else:
                     pass
 
             elif class_mapping is None and self.key_node is not None:
-                if hasattr(self.graph[node_type], 'y'):
-                    insert_doc['label'] = self.graph[self.key_node].y[idx].item()
+                if hasattr(self.graph[node_type], "y"):
+                    insert_doc["label"] = self.graph[self.key_node].y[idx].item()
                 else:
                     pass
 
             elif class_mapping is None and self.key_node is None:
-                insert_doc['label'] = self.graph.y[idx].item()
+                insert_doc["label"] = self.graph.y[idx].item()
             else:
                 pass
-            ## add top-k nearest neighbours
-            if nearest_nbors_search == True:
+            # add top-k nearest neighbours
+            if nearest_nbors_search is True:
                 insert_doc["cosine_sim"] = dist[idx].tolist()
                 insert_doc["similar_nodes"] = nbors[idx].tolist()
             batch.append(insert_doc)
-            index +=1
-            last_record = (idx == (graph_emb.shape[0] - 1))
+            index += 1
+            last_record = idx == (graph_emb.shape[0] - 1)
             if index % batch_size == 0:
                 batch_idx += 1
                 emb_collection.import_bulk(batch)
-                batch = []   
+                batch = []
             if last_record and len(batch) > 0:
                 print("Inserting batch the last batch!")
                 emb_collection.import_bulk(batch)
-                
-                
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..9016033
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,16 @@
+# pyproject.toml
+[tool.pytest.ini_options]
+minversion = "6.0"
+addopts = "-ra -q"
+testpaths = [
+    "tests"
+]
+
+[tool.black]
+line-length=88
+
+[tool.coverage.report]
+omit = ["*tests*"]
+
+[tool.coverage.run]
+omit = ["*tests*"]
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..8c8c281
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,16 @@
+[isort]
+profile = black
+
+[flake8]
+max-line-length = 88
+extend-ignore = E203, W503, E251
+exclude =.git .idea .*_cache dist venv
+
+[mypy]
+strict = True
+ignore_missing_imports = True
+implicit_reexport = True
+scripts_are_modules = True 
+follow_imports = skip
+disallow_subclassing_any = False
+disallow_untyped_decorators = False
diff --git a/setup.py b/setup.py
index 53634d4..92629a5 100644
--- a/setup.py
+++ b/setup.py
@@ -1,37 +1,62 @@
 from setuptools import setup
 
-with open("./README.md") as fp:
+with open("./README.md", encoding="utf-8") as fp:
     long_description = fp.read()
 
+
+dev_requirements = [
+    "black==22.8.0",
+    "pytest==7.1.3",
+    "sphinx==5.1.1",
+    "pre-commit==2.20.0",
+    "types-requests==2.28.11.2",
+    "flake8==5.0.4",
+    "pytest==7.1.3",
+    "pytest-cov==4.0.0",
+    "isort==5.10.1",
+    "mypy==0.982",
+    "types-setuptools==65.5.0.2",
+    "bandit==1.7.4",
+]
+
+
 setup(
     name="fastgraphml",
     author="Sachin Sharma",
     author_email="sachin@arangodb.com",
-    description="Given an input graph the libraray generates graph embeddings using Low-Code built on top of PyG",
+    description="Given an input graph the library generates graph embeddings using Low-Code built on top of PyG",  # noqa: E501
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/arangoml/fastgraphml",
-    keywords=["arangodb", "pyg", "graph deep learning" ,"graph embeddings", "pytorch geometric",
-    "graph machine learning", "graph neural networks"],
-    packages=["graph_embeddings"],
+    keywords=[
+        "arangodb",
+        "pyg",
+        "graph deep learning",
+        "graph embeddings",
+        "pytorch geometric",
+        "graph machine learning",
+        "graph neural networks",
+    ],
+    packages=["fastgraphml"],
     include_package_data=True,
     python_requires=">=3.8",
     license="MIT License",
     install_requires=[
-        "numpy",
-        "sklearn",
-        "tqdm",
-        "umap-learn",
-        "matplotlib",
-        "matplotlib-inline",
-        "seaborn",
-        "adbpyg-adapter",
-        "python-arango>=7.4.1",
-        "setuptools>=45",
+        "requests==2.28.1",
+        "sklearn==0.0",
+        "matplotlib==3.5.2",
+        "seaborn==0.11.2",
+        "matplotlib-inline==0.1.2",
+        "numpy==1.22.4",
+        "tqdm==4.64.0",
+        "umap-learn==0.5.3",
+        "adbpyg-adapter==1.1.0",
+        "python-arango==7.4.1",
+        "setuptools==45",
+        "arango_datasets==1.0.0",
     ],
     extras_require={
-        "dev": [
-        ],
+        "dev": dev_requirements,
     },
     classifiers=[
         "Intended Audience :: Data Scientists/ML Engineers",
@@ -43,4 +68,4 @@
         "Topic :: Utilities",
         "Typing :: Typed",
     ],
-)
\ No newline at end of file
+)


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 271bfba..1efb7e5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,20 +1,23 @@
-# configuration test
+from typing import Any
+
+import pytest
+from adbpyg_adapter import ADBPyG_Adapter
 from arango import ArangoClient
-from arango.database import StandardDatabase
 from arango.http import DefaultHTTPClient
+from torch_geometric.datasets import IMDB, Planetoid
 
-from adbpyg_adapter import ADBPyG_Adapter
-from torch_geometric.datasets import Planetoid, IMDB
 
-def pytest_addoption(parser) -> None:
+def pytest_addoption(parser: pytest.Parser) -> None:
     parser.addoption("--url", action="store", default="http://localhost:8529")
-    parser.addoption("--dbName", action="store", default="_system")
+    parser.addoption("--dbName", action="store", default="fastgraphml")
     parser.addoption("--username", action="store", default="root")
     parser.addoption("--password", action="store", default="")
 
 
-def pytest_configure(config) -> None:
-    global con
+global con
+
+
+def pytest_configure(config: Any) -> None:
     con = {
         "url": config.getoption("url"),
         "username": config.getoption("username"),
@@ -29,26 +32,30 @@ def pytest_configure(config) -> None:
     print("Database: " + con["dbName"])
     print("----------------------------------------")
 
-    class NoTimeoutHTTPClient(DefaultHTTPClient):  # type: ignore
+    class NoTimeoutHTTPClient(DefaultHTTPClient):
         REQUEST_TIMEOUT = None
 
     global db
-    db = ArangoClient(hosts=con["url"], http_client=NoTimeoutHTTPClient()).db(
-        con["dbName"], con["username"], con["password"], verify=True
-    )
-    adbpyg = ADBPyG_Adapter(db)
-    # loading cora dataset (Homogeneous)
+    db = ArangoClient(  # type: ignore
+        hosts=con["url"], http_client=NoTimeoutHTTPClient()
+    ).db(con["dbName"], con["username"], con["password"], verify=True)
+
+    adbpyg = ADBPyG_Adapter(db)  # type: ignore
     dataset_cora = Planetoid("./", "Cora")[0]
-    if db.has_graph('cora'):
-        db.delete_graph("cora", drop_collections=True, ignore_missing=True)
+    if db.has_graph("cora"):  # type: ignore
+        db.delete_graph(  # type: ignore
+            "cora", drop_collections=True, ignore_missing=True
+        )
         adbpyg.pyg_to_arangodb("cora", dataset_cora)
     else:
-        adbpyg.pyg_to_arangodb("cora", dataset_cora)        
-        
+        adbpyg.pyg_to_arangodb("cora", dataset_cora)
+
     # loading IMDB dataset (Heterogeneous)
-    dataset_imdb = IMDB('./imdb')[0]
-    if db.has_graph('imdb'):
-        db.delete_graph('imdb', drop_collections=True, ignore_missing=True) 
+    dataset_imdb = IMDB("./imdb")[0]
+    if db.has_graph("imdb"):  # type: ignore
+        db.delete_graph(  # type: ignore
+            "imdb", drop_collections=True, ignore_missing=True
+        )
         adbpyg.pyg_to_arangodb("imdb", dataset_imdb)
     else:
-        adbpyg.pyg_to_arangodb("imdb", dataset_imdb)
\ No newline at end of file
+        adbpyg.pyg_to_arangodb("imdb", dataset_imdb)
diff --git a/tests/test_ge_models.py b/tests/test_ge_models.py
index 0c39dc4..8c3af21 100644
--- a/tests/test_ge_models.py
+++ b/tests/test_ge_models.py
@@ -1,93 +1,135 @@
-from .conftest import db
-from fastgraphml.graph_embeddings import SAGE, GAT, METAPATH2VEC, DMGI
-from fastgraphml.graph_embeddings import downstream_tasks
+import pytest
+from torch_geometric.datasets import Planetoid
 
-# test sage model
-def test_sage():
+from fastgraphml.graph_embeddings import DMGI, GAT, METAPATH2VEC, SAGE, downstream_tasks
+
+from .conftest import db  # type: ignore
+
+
+def test_sage() -> None:
     metagraph = {
-    "vertexCollections": {
+        "vertexCollections": {
             "cora_N": {"x": "x", "y": "y"},
         },
         "edgeCollections": {
             "cora_E": {},
         },
     }
-    model = SAGE(db, 'graph', metagraph, embedding_size=64) # define model
-    model._train(model, epochs=2) # train
-    embeddings = model.get_embeddings(model=model) # get embeddings
+    model = SAGE(db, "graph", metagraph, embedding_size=64)  # define model
+    model._train(epochs=2)  # train
+    embeddings = model.get_embeddings()  # get embeddings
     # check embeddings size
     assert embeddings.size == int(173312)
     # check similarity search
-    dist, nbors = downstream_tasks.similarity_search(embeddings, top_k_nbors=10, nlist=10, search_type='exact')
+    dist, nbors = downstream_tasks.similarity_search(
+        embeddings, top_k_nbors=10, nlist=10, search_type="exact"
+    )
     assert nbors.size == int(29788)
     # check similarity search
-    dist, nbors = downstream_tasks.similarity_search(embeddings, top_k_nbors=10, nlist=10, search_type='approx')
+    dist, nbors = downstream_tasks.similarity_search(
+        embeddings, top_k_nbors=10, nlist=10, search_type="approx"
+    )
     assert nbors.size == int(29788)
+    # check similarity search
+    with pytest.raises(Exception):
+        dist, nbors = downstream_tasks.similarity_search(
+            embeddings, top_k_nbors=10, nlist=10, search_type="wrong"
+        )
+
     # check vis
-    class_names = {0: 'a', 1: 'b', 2: 'c', 3: 'd', 4: 'e', 5: 'f', 6: 'g'}
-    downstream_tasks.visualize_embeddings(model.G, embeddings, class_mapping=class_names)
+    class_names = {0: "a", 1: "b", 2: "c", 3: "d", 4: "e", 5: "f", 6: "g"}
+    downstream_tasks.visualize_embeddings(
+        model.G, embeddings, class_mapping=class_names
+    )
     # check storing emb
-    if db.has_collection('cora_emb'):
-        db.delete_collection('cora_emb')
-    model.graph_util.store_embeddings(embeddings, collection_name='cora_emb', class_mapping=class_names, nearest_nbors_search=True)
+    if db.has_collection("cora_emb"):
+        db.delete_collection("cora_emb")
+    model.graph_util.store_embeddings(
+        embeddings,
+        collection_name="cora_emb",
+        class_mapping=class_names,
+        nearest_nbors_search=True,
+    )
+
+
+# test sage model on PyG dataset directly
+def test_sage_pyg() -> None:
+    dataset_cora = Planetoid("./", "Cora")[0]
+    model = SAGE(pyg_graph=dataset_cora, embedding_size=64)
+    model._train(epochs=2)
+    embeddings = model.get_embeddings()
+    assert embeddings.size == int(173312)
+
 
 # test gat model
-def test_gat():
+def test_gat() -> None:
     metagraph = {
-    "vertexCollections": {
+        "vertexCollections": {
             "cora_N": {"x": "x", "y": "y"},
         },
         "edgeCollections": {
             "cora_E": {},
         },
     }
-    model = GAT(db, 'graph', metagraph, embedding_size=64, heads=1) # define model
-    model._train(model, epochs=2) # train
-    embeddings = model.get_embeddings(model=model) # get embeddings
+    model = GAT(db, "graph", metagraph, embedding_size=64, heads=1)  # define model
+    model._train(epochs=2)  # train
+    embeddings = model.get_embeddings()  # get embeddings
     # check embeddings size
     assert embeddings.size == int(173312)
 
-def test_m2v():
+
+def test_m2v() -> None:
     metagraph = {
-    "vertexCollections": {
-    
-        "movie": { "x": "x", "y": "y"},  
-        "director": {"x": "x"},
-        "actor": {"x": "x"},
-    },
-    "edgeCollections": {
-        "to": {},
-    },
+        "vertexCollections": {
+            "movie": {"x": "x", "y": "y"},
+            "director": {"x": "x"},
+            "actor": {"x": "x"},
+        },
+        "edgeCollections": {
+            "to": {},
+        },
     }
-    metapaths = [('movie', 'to','actor'),
-             ('actor', 'to', 'movie'), ] # MAM # co-actor relationship
-    
-    model = METAPATH2VEC(db, 'imdb', metagraph, metapaths, key_node='movie', embedding_size=4,
-                     walk_length=5, context_size=6, walks_per_node=5, num_negative_samples=5,
-                     sparse=True)
-    model._train(epochs=2, lr=0.03) # train
+    metapaths = [
+        ("movie", "to", "actor"),
+        ("actor", "to", "movie"),
+    ]  # MAM # co-actor relationship
+
+    model = METAPATH2VEC(
+        db,
+        "imdb",
+        metagraph,
+        metapaths,
+        key_node="movie",
+        embedding_size=4,
+        walk_length=5,
+        context_size=6,
+        walks_per_node=5,
+        num_negative_samples=5,
+        sparse=True,
+    )
+    model._train(epochs=2, lr=0.03)
     embeddings = model.get_embeddings()
     # check embeddings size
-    assert embeddings['movie'].size == int(17112)
+    assert embeddings["movie"].size == int(17112)
+
 
-def test_dmgi():
+def test_dmgi() -> None:
     metagraph = {
-    "vertexCollections": {
-    
-        "movie": { "x": "x", "y": "y"},  
-        "director": {"x": "x"},
-        "actor": {"x": "x"},
-    },
-    "edgeCollections": {
-        "to": {},
-    },
+        "vertexCollections": {
+            "movie": {"x": "x", "y": "y"},
+            "director": {"x": "x"},
+            "actor": {"x": "x"},
+        },
+        "edgeCollections": {
+            "to": {},
+        },
     }
-    metapaths = [
-        [('movie', 'actor'), ('actor', 'movie')]
-        ]  # MAM  
-    model = DMGI(db, 'imdb', metagraph, metapaths, key_node='movie', embedding_size=4)
-    model._train(model, epochs=2, lr=0.0005)
-    embeddings = model.get_embeddings(model=model)
+    metapaths = [[("movie", "actor"), ("actor", "movie")]]  # MAM
+    model = DMGI(db, "imdb", metagraph, metapaths, key_node="movie", embedding_size=4)
+    model._train(epochs=2, lr=0.0005)
+    embeddings = model.get_embeddings()
     # check embeddings size
-    assert embeddings['movie'].size == int(17112)
-    downstream_tasks.visualize_embeddings(model.G, embeddings['movie'], node_type='movie')
\ No newline at end of file
+    assert embeddings["movie"].size == int(17112)
+    downstream_tasks.visualize_embeddings(
+        model.G, embeddings["movie"], node_type="movie"
+    )