load from hub (#33)

* load from hub

* type and lint

* move `datasets` from dev dependencies to common

---------

Co-authored-by: voorhs <[email protected]>

Author: Samoed

CONTRIBUTING.md

@@ -70,7 +70,7 @@ autointent hydra.job_logging.root.level=ERROR
 Еще можно изменить параметры логгера через yaml файлы:
 1. Создадим папку с конфиг. файлами: test_config
 2. test_config/config.yaml:
-```
+```yaml
 defaults:
   - optimization_config
   - _self_

autointent/context/data_handler/schemas.py

@@ -1,7 +1,10 @@
 from enum import Enum
 from functools import cached_property
+from typing import Any
 
+import datasets
 from pydantic import BaseModel
+from typing_extensions import Self
 
 from autointent.custom_types import LabelType
 
@@ -83,3 +86,35 @@ def n_classes(self) -> int:
 
     def to_multilabel(self) -> "Dataset":
         return Dataset(utterances=[utterance.to_multilabel() for utterance in self.utterances], intents=self.intents)
+
+    @classmethod
+    def from_datasets(
+        cls,
+        dataset_name: str,
+        split: str = "train",
+        utterances_kwargs: dict[str, Any] | None = None,
+        intents_kwargs: dict[str, Any] | None = None,
+        # tags_kwargs: dict[str, Any] | None = None,
+    ) -> Self:
+        configs = datasets.get_dataset_config_names(dataset_name)
+
+        utterances = []
+        intents = []
+        if "utterances" in configs:
+            utterance_ds = datasets.load_dataset(
+                dataset_name, name="utterances", split=split, **(utterances_kwargs or {})
+            )
+            utterances = [Utterance(**item) for item in utterance_ds]
+        # tags = []
+        # if "tags" in configs:
+        #     tags_ds = datasets.load_dataset(dataset_name, name="tags", split=split, **(tags_kwargs or {}))
+        if "intents" in configs:
+            intents_ds = datasets.load_dataset(dataset_name, name="intents", split=split, **(intents_kwargs or {}))
+            intents = [Intent(**item) for item in intents_ds]
+        return cls(utterances=utterances, intents=intents)
+
+    def push_to_hub(self, dataset_name: str, split: str = "train") -> None:
+        utterances_ds = datasets.Dataset.from_list([utterance.model_dump() for utterance in self.utterances])
+        intents_ds = datasets.Dataset.from_list([intent.model_dump() for intent in self.intents])
+        utterances_ds.push_to_hub(dataset_name, config_name="utterances", split=split)
+        intents_ds.push_to_hub(dataset_name, config_name="intents", split=split)

pyproject.toml

@@ -19,13 +19,13 @@ pydantic = "^2.9.2"
 hydra-core = "^1.3.2"
 faiss-cpu = "^1.9.0"
 openai = "^1.52.1"
+datasets = "2.20.0"
 
 
 [tool.poetry.group.dev]
 optional = true
 
 [tool.poetry.group.dev.dependencies]
-datasets = "2.20.0"
 tach = "^0.11.3"
 ipykernel = "^6.29.5"
 ipywidgets = "^8.1.5"
@@ -139,6 +139,7 @@ module = [
     "hydra.*",
     "transformers",
     "faiss",
+    "datasets",
     "joblib",
 ]
 ignore_missing_imports = true

scripts/transform_json_to_dataset.py

@@ -0,0 +1,31 @@
+from datasets import Dataset, load_dataset, DatasetDict
+
+
+def transform_dataset(
+    path: str,
+) -> tuple[Dataset | None, Dataset | None, Dataset | None]:
+    ds: DatasetDict = load_dataset("json", data_files=path)["train"]
+    utterance_ds = None
+    tags_ds = None
+    intents_ds = None
+    if "utterances" in ds.column_names:
+        utterance_ds = Dataset.from_list(ds["utterances"][0])
+    if "tags" in ds.column_names:
+        tags_ds = Dataset.from_list(ds["tags"][0])
+    if "intents" in ds.column_names:
+        intents_ds = Dataset.from_list(ds["intents"][0])
+    return utterance_ds, tags_ds, intents_ds
+
+
+def push_json_to_hub(path: str, ds_name: str) -> None:
+    utterance_ds, tags_ds, intents_ds = transform_dataset(path)
+    if utterance_ds is not None:
+        utterance_ds.push_to_hub(ds_name, config_name="utterances")
+    if tags_ds is not None:
+        tags_ds.push_to_hub(ds_name, config_name="tags")
+    if intents_ds is not None:
+        intents_ds.push_to_hub(ds_name, config_name="intents")
+
+
+if __name__ == "__main__":
+    push_json_to_hub("../tests/assets/data/clinc_subset_multilabel.json", "clinc_subset_multilabel")