Update datasets

azimuth/__init__.py

@@ -1,12 +1,12 @@
 from pathlib import Path
 
-from datasets import set_progress_bar_enabled
+from datasets import disable_progress_bar
 from datasets.utils.logging import set_verbosity_error
 
 from azimuth.utils.logs import set_logger_config
 
 PROJECT_ROOT = str(Path(__file__).parent)
 
 set_verbosity_error()
-set_progress_bar_enabled(False)
+disable_progress_bar()
 set_logger_config()

azimuth/routers/v1/utterances.py

@@ -353,9 +353,8 @@ def get_similar(
     item_scores = dict(
         source_ds.select([index])[f"neighbors_{neighbors_dataset_split_name}"][0][:limit]
     )
-    # NOTE: idx may be float in the HF Dataset.
     items: Dict[int, Dict] = {
-        idx: neighbors_ds_with_class_names[int(idx)] for idx in item_scores.keys()  # type: ignore
+        idx: neighbors_ds_with_class_names[int(idx)] for idx in item_scores.keys()
     }
     # Build utterances from `items`
     similar_utterances = [

pyproject.toml

@@ -36,7 +36,7 @@ spacy = "3.1.5"
 en_core_web_sm = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.1.0/en_core_web_sm-3.1.0.tar.gz#egg=en_core_web_sm"}
 fr_core_news_md = {url = "https://github.com/explosion/spacy-models/releases/download/fr_core_news_md-3.1.0/fr_core_news_md-3.1.0.tar.gz#egg=fr_core_news_md"}
 nlpaug = "1.1.10"
-datasets = "1.16.1"
+datasets = "2.1.0"
 
 # Misc
 filelock = "^3.0.12"

tests/test_loading_resources.py

@@ -3,7 +3,7 @@
 # in the root directory of this source tree.
 import string
 from dataclasses import dataclass
-from typing import Callable, List
+from typing import Callable, Dict, List
 
 import numpy as np
 import torch
@@ -73,9 +73,13 @@ def load_intent_data(train_path, test_path, python_loader) -> Dataset:
     return load_dataset(python_loader, data_files={"train": train_path, "test": test_path})
 
 
-def load_file_dataset(*args, azimuth_config, **kwargs):
+def load_file_dataset(data_files: Dict[str, str], azimuth_config):
     # Load a file dataset and cast the label column as a ClassLabel.
-    ds_dict = load_dataset(*args, **kwargs)
+    # Train and test needs to be loaded separately because they don't always share the same columns.
+    # Train sometimes don't have predictions. HF will complain if we load both together.
+    ds_dict = load_dataset("csv", data_files={"train": data_files["train"]})
+    ds_dict_test = load_dataset("csv", data_files={"test": data_files["test"]})
+    ds_dict.update(ds_dict_test)
     features: Features = [v.features for v in ds_dict.values()][0]
     if not isinstance(features[azimuth_config.columns.label], ClassLabel):
         # Get all classes from both set and apply the same mapping to every dataset.

tests/test_modules/test_dataset_analysis/test_dataset_warnings.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import pytest
+from datasets import ClassLabel
 
 from azimuth.dataset_split_manager import DatasetSplitManager
 from azimuth.modules.dataset_analysis.dataset_warnings import DatasetWarningsModule
@@ -88,8 +89,13 @@ def add_rejection_class(mod, monkeypatch):
     # Adding a rejection class
     eval_dm: DatasetSplitManager = mod.get_dataset_split_manager(DatasetSplitName.eval)
     train_dm: DatasetSplitManager = mod.get_dataset_split_manager(DatasetSplitName.train)
-    eval_dm._base_dataset_split.features["label"].names.append("NO_INTENT")  # Should be 2
-    train_dm._base_dataset_split.features["label"].names.append("NO_INTENT")  # Should be 2
+    existing_classes = eval_dm.get_class_names(labels_only=True)
+    eval_dm._base_dataset_split.features["label"] = ClassLabel(
+        num_classes=3, names=existing_classes + ["NO_INTENT"]
+    )
+    train_dm._base_dataset_split.features["label"] = ClassLabel(
+        num_classes=3, names=existing_classes + ["NO_INTENT"]
+    )
     eval_dm._base_dataset_split = eval_dm._base_dataset_split.map(
         lambda u, i: {"label": 2 if i % 10 == 0 else u["label"]}, with_indices=True
     )

tests/utils.py

@@ -79,7 +79,6 @@ def file_based_ds_from_paths(
 ):
     return {
         "class_name": "tests.test_loading_resources.load_file_dataset",
-        "args": ["csv"],
         "kwargs": {
             "data_files": {
                 "train": train,