diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py
index 90ec3fbe..b5d20a10 100644
--- a/bertopic/_bertopic.py
+++ b/bertopic/_bertopic.py
@@ -34,7 +34,7 @@
     from typing import Literal
 else:
     from typing_extensions import Literal
-from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable
+from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable, Optional, Dict
 
 # Models
 try:
@@ -4082,8 +4082,8 @@ def _extract_representative_docs(
         topics: Mapping[str, List[Tuple[str, float]]],
         nr_samples: int = 500,
         nr_repr_docs: int = 5,
-        diversity: float = None,
-    ) -> Union[List[str], List[List[int]]]:
+        diversity: Optional[float] = None,
+    ) -> Tuple[Dict[str, List[str]], List[str], List[List[int]], List[List[int]]]:
         """Approximate most representative documents per topic by sampling
         a subset of the documents in each topic and calculating which are
         most representative to their topic based on the cosine similarity between