DatasetHub

class DatasetFormat(Enum):
    CSV = auto()
    JSON = auto()
    NUMPY = auto()
    PICKLE = auto()
    TEXT = auto()
    IMAGE = auto()
    AUDIO = auto()
    VIDEO = auto()
    SQL = auto()
    CUSTOM = auto()

@classmethod
def from_extension(cls, extension: str) -> DatasetFormat

class Dataset:
    def __init__(
        self,
        data: Any,
        format: DatasetFormat,
        name: str = None,
        metadata: Dict = None,
        transform: Callable = None
    )

def __len__(self) -> int

def __getitem__(self, idx: Union[int, slice]) -> Any

def to_tensor(self, framework: str = 'auto') -> Any

class DatasetHub:
    def __init__(self, cache_dir: Optional[str] = None)

def load_dataset(
    self,
    source: Union[str, Path],
    format: Optional[DatasetFormat] = None,
    download: bool = True,
    **kwargs
) -> Dataset

def register_dataset(
    self,
    name: str,
    source: Union[str, Path],
    format: DatasetFormat,
    metadata: Optional[Dict] = None
) -> None

def load_dataset(*args, **kwargs) -> Dataset

def register_dataset(*args, **kwargs)

import neurenix as nx
from neurenix.data import DatasetHub, Dataset, DatasetFormat, load_dataset

# Load a CSV dataset
dataset = load_dataset(
    "https://example.com/data.csv",
    format=DatasetFormat.CSV
)

print(f"Dataset size: {len(dataset)}")

# Access individual samples
sample = dataset[0]
print(f"First sample: {sample}")

# Access a batch
batch = dataset[0:10]
print(f"Batch shape: {len(batch)}")

# Load a NumPy dataset
numpy_dataset = load_dataset(
    "data.npy",
    format=DatasetFormat.NUMPY
)

# Convert to tensor
tensor_data = numpy_dataset.to_tensor(framework='neurenix')

# Create DatasetHub with custom cache directory
hub = DatasetHub(cache_dir="./my_cache")

# Register a dataset
hub.register_dataset(
    name="my_dataset",
    source="https://example.com/dataset.json",
    format=DatasetFormat.JSON,
    metadata={"version": "1.0", "author": "Neurenix Team"}
)

# Load registered dataset by name
registered_dataset = hub.load_dataset("my_dataset")

# Load with custom transformations
def preprocess(sample):
    # Custom preprocessing logic
    return sample * 2.0

dataset = Dataset(
    data=my_data,
    format=DatasetFormat.NUMPY,
    transform=preprocess
)

# Transformed data is returned when accessing samples
transformed_sample = dataset[0]

# Load image dataset
image_dataset = load_dataset(
    "images/",
    format=DatasetFormat.IMAGE
)

# Load from Parquet (for large datasets)
parquet_dataset = load_dataset(
    "large_dataset.parquet",
    format=DatasetFormat.CUSTOM
)

# Use with DataLoader for training
from neurenix.data import DataLoader

train_loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

for batch in train_loader:
    # Training logic
    pass