feat: add simple handling for datasets.Translation

Renumics · Sep 21, 2023 · 700b5b0 · 700b5b0
1 parent 21a8a3d
commit 700b5b0
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/renumics/spotlight_plugins/core/huggingface_datasource.py b/renumics/spotlight_plugins/core/huggingface_datasource.py
@@ -102,6 +102,7 @@ def get_column_values(
             raw_values = self._dataset.data[column_name].take(indices)
 
         feature = self._dataset.features[column_name]
+
         if isinstance(feature, datasets.Audio) or isinstance(feature, datasets.Image):
             return np.array(
                 [
@@ -112,18 +113,23 @@ def get_column_values(
                 ],
                 dtype=object,
             )
+
         if isinstance(feature, dict):
             return np.array(
                 [value["bytes"].as_py() for value in raw_values], dtype=object
             )
+
         if isinstance(feature, datasets.Sequence):
             if is_array_dtype(intermediate_dtype):
                 return raw_values.to_numpy()
             if is_embedding_dtype(intermediate_dtype):
                 return raw_values.to_numpy()
             return np.array([str(value) for value in raw_values])
-        else:
-            return raw_values.to_numpy()
+
+        if isinstance(feature, datasets.Translation):
+            return np.array([str(value) for value in raw_values])
+
+        return raw_values.to_numpy()
 
     def get_column_metadata(self, _: str) -> ColumnMetadata:
         return ColumnMetadata(nullable=True, editable=False)
@@ -201,5 +207,7 @@ def _get_intermediate_dtype(feature: _FeatureType) -> DType:
             return dtypes.file_dtype
         else:
             raise UnsupportedFeature(feature)
+    elif isinstance(feature, datasets.Translation):
+        return dtypes.str_dtype
     else:
         raise UnsupportedFeature(feature)