Use fast_gather instead of take

Renumics · Dec 7, 2023 · a1770f8 · a1770f8
1 parent 331c831
commit a1770f8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/renumics/spotlight_plugins/core/huggingface_datasource.py b/renumics/spotlight_plugins/core/huggingface_datasource.py
@@ -90,9 +90,9 @@ def get_column_values(
                 raw_values = self._dataset.data[column_name]
             else:
                 actual_indices = list(range(len(self._dataset)))[indices]
-                raw_values = self._dataset.data[column_name].take(actual_indices)
+                raw_values = self._dataset.data.fast_gather(actual_indices)[column_name]
         else:
-            raw_values = self._dataset.data[column_name].take(indices)
+            raw_values = self._dataset.data.fast_gather(indices)[column_name]
 
         feature = self._dataset.features[column_name]