GitHub issue: Using a registry instead of calling globals for fetching feature types
Author: psmyth94 (CONTRIBUTOR)
State: closed | Created: 2024-03-10 17:47:51

Hello,

When working with bio-data, each feature often has metadata associated with it (e.g. species, lineage, snp position, etc). To store this, I like to use the feature classes with the added `metadata` attribute. However, when saving or loading with custom features, you get an error since that class doesn't exist in the global namespace in `datasets.features.features`. Take for example,



```python
from dataclasses import dataclass, field
from datasets import Dataset
from datasets.features.features import Value, Features

@dataclass
class FeatureA(Value):
    metadata: dict = field(default=dict)
    _type: str = field(default="FeatureA", init=False, repr=False)

@dataclass
class FeatureB(Value):
    metadata: dict = field(default_factory=dict)
    _type: str = field(default="FeatureB", init=False, repr=False)

test_data = {
    "a": [1, 2, 3],
    "b": [4, 5, 6],
}
test_data = Dataset.from_dict(
    test_data, 
    features=Features({
        "a": FeatureA("int32", metadata={"species": "lactobacillus acetotolerans"}),
        "b": FeatureB("int32", metadata={"species": "lactobacillus iners"}),
    })
)

# returns an error since FeatureA and FeatureB are not in the global namespace
test_data.save_to_disk('./test_data')
```


    Saving the dataset (0/1 shards):   0%|          | 0/3 [00:00<?, ? examples/s]



    ---------------------------------------------------------------------------

    KeyError                                  Traceback (most recent call last)

    Cell In[2], line 28
         19 test_data = Dataset.from_dict(
         20     test_data, 
         21     features=Features({
       (...)
         24     })
         25 )
         27 # returns an error since FeatureA and FeatureB are not in the global namespace
    ---> 28 test_data.save_to_disk('./test_data')
    ...
    File ~\Documents\datasets\src\datasets\features\features.py:1361, in generate_from_dict(obj)
       1359     return {key: generate_from_dict(value) for key, value in obj.items()}
       1360 obj = dict(obj)
    -> 1361 class_type = globals()[obj.pop("_type")]
       1363 if class_type == Sequence:
       1364     return Sequence(feature=generate_from_dict(obj["feature"]), length=obj.get("length", -1))
    

    KeyError: 'FeatureA'


We can avoid this by having a registry (like formatters) and doing


```python
from datasets.features.features import register_feature
register_feature(FeatureA, "FeatureA")
register_feature(FeatureB, "FeatureB")
test_data.save_to_disk('./test_data')
```


    Saving the dataset (1/1 shards): 100%|------| 3/3 [00:00<00:00, 211.13 examples/s]


and loading from disk returns with all metadata information


```python
from datasets import load_from_disk
test_data = load_from_disk('./test_data')
test_data.features
```




    {'a': FeatureA(dtype='int32', id=None, metadata={'species': 'lactobacillus acetotolerans'}),
     'b': FeatureB(dtype='int32', id=None, metadata={'species': 'lactobacillus iners'})}