fixes #534

CHANGELOG.md

@@ -17,6 +17,7 @@ Most recent releases are shown at the top. Each release shows:
 
 ### fixed:
 - Update `test_lda.py` due to changes in `numpy` (#533)
+- Ensure `TopicModel.filter` returns `pd.DataFrame` if supplied (#534)
 
 
 ## 0.41.4 (2024-06-18)

ktrain/text/eda.py

@@ -369,11 +369,14 @@ def filter(self, obj):
                 "Length of obj is not consistent with the number of documents "
                 + "supplied to get_topic_model"
             )
-        # obj = np.array(obj) if isinstance(obj, list) else obj
-        # return obj[self.bool_array]
-        from itertools import compress
+        if isinstance(obj, pd.DataFrame):
+            # ensure DataFrame is returned (#534)
+            return obj[self.bool_array]
+        else:
+            # reduces memory footprint (#531)
+            from itertools import compress
 
-        return list(compress(obj, self.bool_array))
+            return list(compress(obj, self.bool_array))
 
     def get_docs(self, topic_ids=[], doc_ids=[], rank=False):
         """