Fix bug in flattening list of lists. (#32)

smartnews · Jul 17, 2022 · b525b02 · b525b02
1 parent be5e723
commit b525b02
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/rsdiv/evaluation/diversity_metrics.py b/src/rsdiv/evaluation/diversity_metrics.py
@@ -15,7 +15,8 @@ class DiversityMetrics:
     def _get_histogram(
         items: Union[Iterable[Hashable], Iterable[Sequence[Hashable]]],
     ) -> np.ndarray:
-        if isinstance(next(iter(items)), Sequence):
+        first_element = next(iter(items))
+        if isinstance(first_element, Sequence) and not isinstance(first_element, str):
             items = chain(*items)
         flatten_items = list(items)
         return np.asarray(pd.Series(flatten_items).value_counts())
@@ -80,7 +81,8 @@ def get_lorenz_curve(
     def get_distribution(
         cls, items: Union[Iterable[Hashable], Iterable[Sequence[Hashable]]]
     ) -> pd.DataFrame:
-        if isinstance(next(iter(items)), Sequence):
+        first_element = next(iter(items))
+        if isinstance(first_element, Sequence) and not isinstance(first_element, str):
             items = chain(*items)
         counter: pd.DataFrame = pd.DataFrame(Counter(items).most_common())
         counter.columns = pd.Index(["category", "percentage"])