diff --git a/.github/workflows/doc-build.yml b/.github/workflows/doc-build.yml
index 2177e25d2..f9f0c99d3 100644
--- a/.github/workflows/doc-build.yml
+++ b/.github/workflows/doc-build.yml
@@ -15,6 +15,7 @@ jobs:
     uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main
     with:
       commit_sha: ${{ github.sha }}
-      package: optimum-neuron
+      package: optimum.neuron
+      package_name: optimum-neuron
     secrets:
       token: ${{ secrets.HUGGINGFACE_PUSH }}
\ No newline at end of file
diff --git a/.github/workflows/doc-pr-build.yml b/.github/workflows/doc-pr-build.yml
index f76e78677..58213da54 100644
--- a/.github/workflows/doc-pr-build.yml
+++ b/.github/workflows/doc-pr-build.yml
@@ -18,6 +18,7 @@ jobs:
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}
-      package: optimum-neuron
+      package: optimum.neuron
+      package_name: optimum-neuron
     secrets:
       token: ${{ secrets.HUGGINGFACE_PUSH }}
\ No newline at end of file
diff --git a/Makefile b/Makefile
index 41854380e..39008ff2a 100644
--- a/Makefile
+++ b/Makefile
@@ -47,20 +47,6 @@ build_dist:
 pypi_upload: build_dist
 	python -m twine upload dist/*
 
-build_doc_docker_image:
-	docker build -t doc_maker --build-arg commit_sha=$(COMMIT_SHA_SUBPACKAGE) --build-arg clone_url=$(REAL_CLONE_URL) ./docs
-
-doc: build_doc_docker_image
-	@test -n "$(BUILD_DIR)" || (echo "BUILD_DIR is empty." ; exit 1)
-	@test -n "$(VERSION)" || (echo "VERSION is empty." ; exit 1)
-	docker run -v $(CURRENT_DIR):/doc_folder --workdir=/doc_folder doc_maker \
-	doc-builder build optimum.neuron /optimum-neuron/docs/source/ \
-		--build_dir $(BUILD_DIR) \
-		--version $(VERSION) \
-		--version_tag_suffix "" \
-		--html \
-		--clean
-
 test_installs:
 	python -m pip install .[tests]
 	python -m pip install git+https://github.com/huggingface/transformers.git
diff --git a/docs/README.md b/docs/README.md
index 97e9d8f42..3c51ef2c5 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -7,11 +7,11 @@ pip install hf-doc-builder==0.4.0 watchdog --upgrade
 
 2. Local Development
 ```bash
-doc-builder preview optimum docs/source/
+doc-builder preview optimum.neuron docs/source/
 ```
 3. Build Docs
 ```bash
-doc-builder build optimum docs/source/ --build_dir build/ 
+doc-builder build optimum.neuron docs/source/ --build_dir build/ 
 ```
 
 ## Add assets/Images
diff --git a/docs/source/package_reference/trainer.mdx b/docs/source/package_reference/trainer.mdx
index 2bcb59b4b..76a21983d 100644
--- a/docs/source/package_reference/trainer.mdx
+++ b/docs/source/package_reference/trainer.mdx
@@ -47,6 +47,6 @@ Another way to customize the training loop behavior for the PyTorch [`TrainiumTr
 
 ## TrainiumTrainer
 
-[[autodoc]] optimum.neuron.trainers.TrainiumTrainer
+[[autodoc]] trainers.TrainiumTrainer
 
-[[autodoc]] optimum.neuron.trainers.Seq2SeqTrainiumTrainer
+[[autodoc]] trainers.Seq2SeqTrainiumTrainer
diff --git a/optimum/neuron/__init__.py b/optimum/neuron/__init__.py
index 9c4dd8823..976aef62f 100644
--- a/optimum/neuron/__init__.py
+++ b/optimum/neuron/__init__.py
@@ -16,6 +16,7 @@
 
 from .trainers import Seq2SeqTrainiumTrainer, TrainiumTrainer
 from .utils import patch_transformers_for_neuron_sdk
+from .version import __version__
 
 
 if not os.environ.get("DISABLE_TRANSFORMERS_PATCHING", False):
diff --git a/setup.py b/setup.py
index 431d693db..b579ba0e5 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,8 @@
 QUALITY_REQUIRES = [
     "black",
     "ruff",
-    "hf_doc_builder @ git+https://github.com/huggingface/doc-builder.git",
+    "isort",
+    "hf-doc-builder==0.4.0",
 ]
 
 EXTRAS_REQUIRE = {