# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_datasets.ipynb.

# %% ../nbs/00_datasets.ipynb 2
from __future__ import annotations
import math
import numpy as np
import matplotlib.pyplot as plt
from operator import itemgetter
from itertools import zip_longest
import fastcore.all as fc
from torch.utils.data import default_collate
from torch.utils.data import DataLoader, SequentialSampler, RandomSampler, BatchSampler

# %% auto 0
__all__ = ['get_dls', 'inplace', 'collate_dict', 'DataLoaders']

# %% ../nbs/00_datasets.ipynb 3
def get_dls(train_ds, valid_ds, bs, **kwargs):
    return (DataLoader(train_ds, batch_size=bs, shuffle=True, **kwargs),
            DataLoader(valid_ds, batch_size=bs*2, **kwargs))

# %% ../nbs/00_datasets.ipynb 8
def inplace(f):
    def _f(b):
        f(b)
        return b
    return _f

# %% ../nbs/00_datasets.ipynb 11
def collate_dict(ds):
    get = itemgetter(*ds.features)
    def _f(b): return get(default_collate(b))
    return _f

# %% ../nbs/00_datasets.ipynb 13
class DataLoaders:
    def __init__(self, *dls): self.train,self.valid = dls[:2]

    @classmethod
    def from_dd(cls, dd, batch_size, as_tuple=True, **kwargs):
        f = collate_dict(dd['train'])
        return cls(*get_dls(*dd.values(), bs=batch_size, collate_fn=f))
