diff --git a/bayesclass/__init__.py b/bayesclass/__init__.py
index bd33230..78a0c35 100644
--- a/bayesclass/__init__.py
+++ b/bayesclass/__init__.py
@@ -1,5 +1,5 @@
-from .bayesclass import TAN, KDB, AODE
-from ._version import __version__
+# from .bayesclass import TAN, KDB, AODE
+from bayesclass._version import __version__
 
 __author__ = "Ricardo Montañana Gómez"
 __copyright__ = "Copyright 2020-2023, Ricardo Montañana Gómez"
diff --git a/bayesclass/bayesclass.py b/bayesclass/bayesclass.py
index 49acd6e..7b8853e 100644
--- a/bayesclass/bayesclass.py
+++ b/bayesclass/bayesclass.py
@@ -86,7 +86,7 @@ class BayesBase(BaseEstimator, ClassifierMixin):
         --------
         >>> import numpy as np
         >>> import pandas as pd
-        >>> from bayesclass import TAN
+        >>> from bayesclass.bayesclass import TAN
         >>> features = ['A', 'B', 'C', 'D', 'E']
         >>> np.random.seed(17)
         >>> values = pd.DataFrame(np.random.randint(low=0, high=2,
@@ -142,7 +142,7 @@ class BayesBase(BaseEstimator, ClassifierMixin):
         --------
         >>> import numpy as np
         >>> import pandas as pd
-        >>> from bayesclass import TAN
+        >>> from bayesclass.bayesclass import TAN
         >>> features = ['A', 'B', 'C', 'D', 'E']
         >>> np.random.seed(17)
         >>> values = pd.DataFrame(np.random.randint(low=0, high=2,
@@ -373,4 +373,4 @@ class AODE(BayesBase, BaseEnsemble):
         dataset = pd.DataFrame(X, columns=self.features_, dtype="int16")
         for index, model in enumerate(self.models_):
             result[:, index] = model.predict(dataset).values.ravel()
-        return mode(result, axis=1).mode.ravel()
+        return mode(result, axis=1, keepdims=False).mode.ravel()
diff --git a/bayesclass/tests/test_AODE.py b/bayesclass/tests/test_AODE.py
index 4bd1d00..fa0606c 100644
--- a/bayesclass/tests/test_AODE.py
+++ b/bayesclass/tests/test_AODE.py
@@ -4,7 +4,7 @@ from sklearn.datasets import load_iris
 from sklearn.preprocessing import KBinsDiscretizer
 
 
-from bayesclass import AODE
+from bayesclass.bayesclass import AODE
 from .._version import __version__
 
 
diff --git a/bayesclass/tests/test_KDB.py b/bayesclass/tests/test_KDB.py
index 3d8eb4f..bb97b2c 100644
--- a/bayesclass/tests/test_KDB.py
+++ b/bayesclass/tests/test_KDB.py
@@ -6,7 +6,7 @@ from matplotlib.testing.decorators import image_comparison
 from matplotlib.testing.conftest import mpl_test_settings
 
 
-from bayesclass import KDB
+from bayesclass.bayesclass import KDB
 from .._version import __version__
 
 
diff --git a/bayesclass/tests/test_TAN.py b/bayesclass/tests/test_TAN.py
index b3b2e5f..7c2bc9e 100644
--- a/bayesclass/tests/test_TAN.py
+++ b/bayesclass/tests/test_TAN.py
@@ -6,7 +6,7 @@ from matplotlib.testing.decorators import image_comparison
 from matplotlib.testing.conftest import mpl_test_settings
 
 
-from bayesclass import TAN
+from bayesclass.bayesclass import TAN
 from .._version import __version__
 
 
diff --git a/bayesclass/tests/test_common.py b/bayesclass/tests/test_common.py
index ee9ccae..4c9089e 100644
--- a/bayesclass/tests/test_common.py
+++ b/bayesclass/tests/test_common.py
@@ -2,7 +2,7 @@ import pytest
 
 from sklearn.utils.estimator_checks import check_estimator
 
-from bayesclass import TAN, KDB, AODE
+from bayesclass.bayesclass import TAN, KDB, AODE
 
 
 @pytest.mark.parametrize("estimator", [TAN(), KDB(k=2), AODE()])
diff --git a/bayesclass/tests/test_doctest.py b/bayesclass/tests/test_doctest.py
new file mode 100644
index 0000000..9eef3a0
--- /dev/null
+++ b/bayesclass/tests/test_doctest.py
@@ -0,0 +1,7 @@
+import doctest
+import bayesclass
+
+
+def load_tests(loader, tests, ignore):
+    tests.addTests(doctest.DocTestSuite(bayesclass))
+    return tests
diff --git a/example.py b/example.py
index 5f4e164..5bc9715 100644
--- a/example.py
+++ b/example.py
@@ -1,36 +1,33 @@
-import sys
 import time
 from sklearn.model_selection import cross_val_score, StratifiedKFold
-from benchmark import Datasets
-from bayesclass import TAN
+from sklearn.preprocessing import KBinsDiscretizer
+from sklearn.datasets import load_wine
+from bayesclass.bayesclass import TAN
 import warnings
 
 
-if len(sys.argv) < 2:
-    print("Usage: python3 example.py <dataset> [n_folds]")
-    exit(1)
 # Warnings are not errors
 warnings.simplefilter("ignore")
 start = time.time()
 random_state = 17
-name = sys.argv[1]
-n_folds = int(sys.argv[2]) if len(sys.argv) == 3 else 5
-dt = Datasets()
-name_list = list(dt) if name == "all" else [name]
+n_folds = 5
 print(f"Accuracy in {n_folds} folds stratified crossvalidation")
-for name in name_list:
-    dataset_start = time.time()
-    X, y = dt.load(name)
-    clf = TAN(random_state=random_state)
-    fit_params = dict(
-        features=dt.get_features(), class_name=dt.get_class_name(), head=0
-    )
-    kfold = StratifiedKFold(
-        n_splits=n_folds, shuffle=True, random_state=random_state
-    )
-    score = cross_val_score(clf, X, y, cv=kfold, fit_params=fit_params)
-    print(
-        f"{name:20s}{'.' * 10}{score.mean():9.7f} "
-        f"({time.time()-dataset_start:7.2f} seconds)"
-    )
+dataset_start = time.time()
+dataset = load_wine()
+Xc = dataset.data
+enc = KBinsDiscretizer(encode="ordinal")
+X = enc.fit_transform(Xc)
+y = dataset.target
+clf = TAN(random_state=random_state)
+fit_params = dict(features=dataset.feature_names, class_name="class", head=0)
+kfold = StratifiedKFold(
+    n_splits=n_folds, shuffle=True, random_state=random_state
+)
+score = cross_val_score(clf, X, y, cv=kfold, fit_params=fit_params)
+print(
+    f"wine {'.' * 10}{score.mean():9.7f} "
+    f"({time.time()-dataset_start:7.2f} seconds)"
+)
+clf.fit(X, y, **fit_params)
+clf.plot("TAN wine")
 print(f"Took {time.time()-start:.2f} seconds")
diff --git a/examples/plot_classifier.py b/examples/plot_classifier.py
index 6e5bb77..57076ef 100644
--- a/examples/plot_classifier.py
+++ b/examples/plot_classifier.py
@@ -7,7 +7,7 @@ An example plot of :class:`bayesclass.TAN`
 """
 import numpy as np
 from matplotlib import pyplot as plt
-from bayesclass import TAN
+from bayesclass.bayesclass import TAN
 
 X = [[0, 0], [1, 1]]
 y = [0, 1]
diff --git a/pyproject.toml b/pyproject.toml
index bffac4b..edea8f6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,14 @@
 [build-system]
-requires = ["flit_core >=3.2,<4"]
-build-backend = "flit_core.buildapi"
+requires = ["setuptools", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+packages = ["bayesclass"]
+license-files = ["LICENSE"]
+
+[tool.setuptools.dynamic]
+version = { attr = "bayesclass.__version__" }
+readme = { file = ["README.md"] }
 
 [project]
 name = "bayesclass"
@@ -8,17 +16,16 @@ description = "A collection of Bayesian Estimators."
 authors = [
   { name = "Ricardo Montañana", email = "ricardo.montanana@alu.uclm.es" },
 ]
-dynamic = ['version']
+dynamic = ['version', 'readme']
 dependencies = [
-  "numpy",
   "scipy",
+  "numpy",
   "pandas",
   "scikit-learn",
   "pgmpy",
   "networkx",
   "matplotlib",
 ]
-readme = "README.md"
 requires-python = ">=3.8"
 classifiers = [
   "Development Status :: 3 - Alpha",
@@ -44,7 +51,7 @@ doc = ["sphinx", "sphinx-gallery", "sphinx_rtd_theme", "numpydoc"]
 Home = "https://github.com/doctorado-ml/bayesclass"
 
 [tool.pytest.ini_options]
-addopts = "--cov --cov-report html --cov-report term-missing --cov-fail-under 95"
+addopts = "--cov --cov-report html --cov-report term-missing --cov-fail-under 95 " #--doctest-modules
 
 [tool.coverage.run]
 source = ["bayesclass"]