Skip to content

Commit 4341831

Browse files
committed
Fix type annotations for _extract_representative_docs
1 parent c5b4f9f commit 4341831

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

bertopic/_bertopic.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
from typing import Literal
3535
else:
3636
from typing_extensions import Literal
37-
from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable
37+
from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable, Optional, Dict
3838

3939
# Models
4040
try:
@@ -4082,8 +4082,8 @@ def _extract_representative_docs(
40824082
topics: Mapping[str, List[Tuple[str, float]]],
40834083
nr_samples: int = 500,
40844084
nr_repr_docs: int = 5,
4085-
diversity: float = None,
4086-
) -> Union[List[str], List[List[int]]]:
4085+
diversity: Optional[float] = None,
4086+
) -> Tuple[Dict[str, List[str]], List[str], List[List[int]], List[List[int]]]:
40874087
"""Approximate most representative documents per topic by sampling
40884088
a subset of the documents in each topic and calculating which are
40894089
most representative to their topic based on the cosine similarity between

0 commit comments

Comments
 (0)