konfai 1.0.7__py3-none-any.whl → 1.0.9__py3-none-any.whl

konfai/predictor.py

@@ -6,12 +6,12 @@ import torch
 import tqdm
 import os
 
-from konfai import MODELS_DIRECTORY, PREDICTIONS_DIRECTORY, CONFIG_FILE, MODEL, DEEP_LEARNING_API_ROOT
+from konfai import MODELS_DIRECTORY, PREDICTIONS_DIRECTORY, CONFIG_FILE, MODEL, KONFAI_ROOT
 from konfai.utils.config import config
 from konfai.utils.utils import State, get_patch_slices_from_nb_patch_per_dim, NeedDevice, _getModule, DistributedObject, DataLog, description
 from konfai.utils.dataset import Dataset, Attribute
-from konfai.data.dataset import DataPrediction, DatasetIter
-from konfai.data.HDF5 import Accumulator, PathCombine
+from konfai.data.data_manager import DataPrediction, DatasetIter
+from konfai.data.patching import Accumulator, PathCombine
 from konfai.network.network import ModelLoader, Network, NetState, CPU_Model
 from konfai.data.transform import Transform, TransformLoader
 
@@ -22,6 +22,8 @@ import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 import importlib
+import copy
+from collections import defaultdict
 
 class OutDataset(Dataset, NeedDevice, ABC):
 
@@ -50,13 +52,13 @@ class OutDataset(Dataset, NeedDevice, ABC):
             
             if _transform_type is not None:
                 for classpath, transform in _transform_type.items():
-                    transform = transform.getTransform(classpath, DL_args =  "{}.outsDataset.{}.OutDataset.{}".format(DEEP_LEARNING_API_ROOT(), name_layer, name))
+                    transform = transform.getTransform(classpath, DL_args =  "{}.outsDataset.{}.OutDataset.{}".format(KONFAI_ROOT(), name_layer, name))
                     transform.setDatasets(datasets)
                     transform_type.append(transform)
 
         if self._patchCombine is not None:
-            module, name = _getModule(self._patchCombine, "data.HDF5")
-            self.patchCombine = getattr(importlib.import_module(module), name)(config = None, DL_args =  "{}.outsDataset.{}.OutDataset".format(DEEP_LEARNING_API_ROOT(), name_layer))
+            module, name = _getModule(self._patchCombine, "data.patching")
+            self.patchCombine = getattr(importlib.import_module(module), name)(config = None, DL_args =  "{}.outsDataset.{}.OutDataset".format(KONFAI_ROOT(), name_layer))
     
     def setPatchConfig(self, patchSize: Union[list[int], None], overlap: Union[int, None], nb_data_augmentation: int) -> None:
         if patchSize is not None and overlap is not None:
@@ -92,7 +94,7 @@ class OutDataset(Dataset, NeedDevice, ABC):
 class OutSameAsGroupDataset(OutDataset):
 
     @config("OutDataset")
-    def __init__(self, dataset_filename: str = "Dataset:h5", group: str = "default", sameAsGroup: str = "default", pre_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, post_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, final_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, patchCombine: Union[str, None] = None, redution: str = "mean", inverse_transform: bool = True) -> None:
+    def __init__(self, dataset_filename: str = "./Dataset:mha", group: str = "default", sameAsGroup: str = "default", pre_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, post_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, final_transforms : dict[str, TransformLoader] = {"default:Normalize": TransformLoader()}, patchCombine: Union[str, None] = None, redution: str = "mean", inverse_transform: bool = True) -> None:
         super().__init__(dataset_filename, group, pre_transforms, post_transforms, final_transforms, patchCombine)
         self.group_src, self.group_dest = sameAsGroup.split(":")
         self.redution = redution
@@ -152,7 +154,7 @@ class OutSameAsGroupDataset(OutDataset):
         if self.redution == "mean":
             result = torch.mean(result.float(), dim=0).to(dtype)
         elif self.redution == "median":
-            result, _ = torch.median(result.float(), dim=0).to(dtype)
+            result, _ = torch.median(result.float(), dim=0)
         else:
             raise NameError("Reduction method does not exist (mean, median)")
         for transform in self.final_transforms:
@@ -199,19 +201,19 @@ class OutDatasetLoader():
 
 class _Predictor():
 
-    def __init__(self, world_size: int, global_rank: int, local_rank: int, predict_path: str, data_log: Union[list[str], None], outsDataset: dict[str, OutDataset], model: DDP, dataloader_prediction: DataLoader) -> None:
+    def __init__(self, world_size: int, global_rank: int, local_rank: int, predict_path: str, data_log: Union[list[str], None], outsDataset: dict[str, OutDataset], modelComposite: DDP, dataloader_prediction: DataLoader) -> None:
         self.world_size = world_size        
         self.global_rank = global_rank
         self.local_rank = local_rank
 
-        self.model = model
+        self.modelComposite = modelComposite
         self.dataloader_prediction = dataloader_prediction
         self.outsDataset = outsDataset
 
         
         self.it = 0
 
-        self.device = self.model.device
+        self.device = self.modelComposite.device
         self.dataset: DatasetIter = self.dataloader_prediction.dataset
         patch_size, overlap = self.dataset.getPatchConfig()
         for outDataset in self.outsDataset.values():
@@ -220,7 +222,7 @@ class _Predictor():
         if data_log is not None:
             for data in data_log:
                 self.data_log[data.split("/")[0].replace(":", ".")] = (DataLog.__getitem__(data.split("/")[1]).value[0], int(data.split("/")[2]))
-        self.tb = SummaryWriter(log_dir = predict_path+"Metric/") if len([network for network in self.model.module.getNetworks().values() if network.measure is not None]) or len(self.data_log) else None
+        self.tb = SummaryWriter(log_dir = predict_path+"Metric/") if len([network for network in self.modelComposite.module.getNetworks().values() if network.measure is not None]) or len(self.data_log) else None
         
     def __enter__(self):
         return self
@@ -234,15 +236,15 @@ class _Predictor():
     
     @torch.no_grad()
     def run(self):
-        self.model.eval()  
-        self.model.module.setState(NetState.PREDICTION)
-        desc = lambda : "Prediction : {}".format(description(self.model))
+        self.modelComposite.eval()  
+        self.modelComposite.module.setState(NetState.PREDICTION)
+        desc = lambda : "Prediction : {}".format(description(self.modelComposite))
         self.dataloader_prediction.dataset.load()
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=False, desc = desc(), total=len(self.dataloader_prediction), disable=self.global_rank != 0 and "DL_API_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_prediction), leave=False, desc = desc(), total=len(self.dataloader_prediction), disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
             dist.barrier()
             for it, data_dict in batch_iter:
                 input = self.getInput(data_dict)
-                for name, output in self.model(input, list(self.outsDataset.keys())):
+                for name, output in self.modelComposite(input, list(self.outsDataset.keys())):
                     self._predict_log(data_dict)
                     outDataset = self.outsDataset[name]
                     for i, (index, patch_augmentation, patch_index) in enumerate([(int(index), int(patch_augmentation), int(patch_index)) for index, patch_augmentation, patch_index in zip(list(data_dict.values())[0][1], list(data_dict.values())[0][2], list(data_dict.values())[0][3])]):
@@ -254,7 +256,7 @@ class _Predictor():
                 self.it += 1
 
     def _predict_log(self, data_dict : dict[str, tuple[torch.Tensor, int, int, int]]):
-        measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank, {"" : self.model.module}, 1)
+        measures = DistributedObject.getMeasure(self.world_size, self.global_rank, self.local_rank, {"" : self.modelComposite.module}, 1)
         
         if self.global_rank == 0:
             images_log = []
@@ -265,30 +267,67 @@ class _Predictor():
                     else:
                         images_log.append(name.replace(":", "."))
 
-            for name, network in self.model.module.getNetworks().items():
+            for name, network in self.modelComposite.module.getNetworks().items():
                 if network.measure is not None:
                     self.tb.add_scalars("Prediction/{}/Loss".format(name), {k : v[1] for k, v in measures["{}{}".format(name, "")][0].items()}, self.it)
                     self.tb.add_scalars("Prediction/{}/Metric".format(name), {k : v[1] for k, v in measures["{}{}".format(name, "")][1].items()}, self.it)
                 if len(images_log):
                     for name, layer, _ in self.model.module.get_layers([v.to(0) for k, v in self.getInput(data_dict).items() if k[1]], images_log):
                         self.data_log[name][0](self.tb, "Prediction/{}".format(name), layer[:self.data_log[name][1]].detach().cpu().numpy(), self.it)
+
+class ModelComposite(Network):
+
+    def __init__(self, model: Network, nb_models: int, method: str):
+        super().__init__(model.in_channels, model.optimizer, model.schedulers, model.outputsCriterionsLoader, model.patch, model.nb_batch_per_step, model.init_type, model.init_gain, model.dim)
+        self.method = method
+        for i in range(nb_models):
+            self.add_module("Model_{}".format(i), copy.deepcopy(model), in_branch=[0], out_branch=["output_{}".format(i)])
+
+    def load(self, state_dicts : list[dict[str, dict[str, torch.Tensor]]]):
+        for i, state_dict in enumerate(state_dicts):
+            self["Model_{}".format(i)].load(state_dict, init=False)
+            self["Model_{}".format(i)].setName("{}_{}".format(self["Model_{}".format(i)].getName(), i))
+            
+    def forward(self, data_dict: dict[tuple[str, bool], torch.Tensor], output_layers: list[str] = []) -> list[tuple[str, torch.Tensor]]:
+        result = {}
+        for name, module in self.items():
+            result[name] = module(data_dict, output_layers)
         
+        aggregated = defaultdict(list)
+        for module_outputs in result.values():
+            for key, tensor in module_outputs:
+                aggregated[key].append(tensor)
+
+        final_outputs = []
+        for key, tensors in aggregated.items():
+            stacked = torch.stack(tensors, dim=0)
+            if self.method == 'mean':
+                agg = torch.mean(stacked, dim=0)
+            elif self.method == 'median':
+                agg = torch.median(stacked, dim=0).values
+            final_outputs.append((key, agg))
+
+        return final_outputs
+    
+
 class Predictor(DistributedObject):
 
     @config("Predictor")
     def __init__(self, 
                     model: ModelLoader = ModelLoader(),
                     dataset: DataPrediction = DataPrediction(),
+                    combine: str = "mean",
                     train_name: str = "name",
                     manual_seed : Union[int, None] = None,
                     gpu_checkpoints: Union[list[str], None] = None,
                     outsDataset: Union[dict[str, OutDatasetLoader], None] = {"default:Default" : OutDatasetLoader()},
                     images_log: list[str] = []) -> None:
-        if os.environ["DEEP_LEANING_API_CONFIG_MODE"] != "Done":
+        if os.environ["KONFAI_CONFIG_MODE"] != "Done":
             exit(0)
         super().__init__(train_name)
         self.manual_seed = manual_seed
         self.dataset = dataset
+        self.combine = combine
 
         self.model = model.getModel(train=False)
         self.it = 0
@@ -305,34 +344,37 @@ class Predictor(DistributedObject):
         
         self.gpu_checkpoints = gpu_checkpoints
 
-    def _load(self) -> dict[str, dict[str, torch.Tensor]]:
-        if MODEL().startswith("https://"):
-            try:
-                state_dict = {MODEL().split(":")[1]: torch.hub.load_state_dict_from_url(url=MODEL().split(":")[0], map_location="cpu", check_hash=True)}
-            except:
-                raise Exception("Model : {} does not exist !".format(MODEL())) 
-        else:
-            if MODEL() != "":
-                path = ""
-                name = MODEL() 
+    def _load(self) -> list[dict[str, dict[str, torch.Tensor]]]:
+        model_paths = MODEL().split(":")
+        state_dicts = []
+        for model_path in model_paths:
+            if model_path.startswith("https://"):
+                try:
+                    state_dicts.append(torch.hub.load_state_dict_from_url(url=model_path, map_location="cpu", check_hash=True))
+                except:
+                    raise Exception("Model : {} does not exist !".format(model_path)) 
             else:
-                if self.name.endswith(".pt"):
-                    path = MODELS_DIRECTORY()+"/".join(self.name.split("/")[:-1])+"/StateDict/"
-                    name = self.name.split("/")[-1]
+                if model_path != "":
+                    path = ""
+                    name = model_path
                 else:
-                    path = MODELS_DIRECTORY()+self.name+"/StateDict/"
-                    name = sorted(os.listdir(path))[-1]
-            if os.path.exists(path+name):
-                state_dict = torch.load(path+name, weights_only=False)
-            else:
-                raise Exception("Model : {} does not exist !".format(path+name))
-        return state_dict
+                    if self.name.endswith(".pt"):
+                        path = MODELS_DIRECTORY()+"/".join(self.name.split("/")[:-1])+"/StateDict/"
+                        name = self.name.split("/")[-1]
+                    else:
+                        path = MODELS_DIRECTORY()+self.name+"/StateDict/"
+                        name = sorted(os.listdir(path))[-1]
+                if os.path.exists(path+name):
+                    state_dicts.append(torch.load(path+name, weights_only=False))
+                else:
+                    raise Exception("Model : {} does not exist !".format(path+name))
+        return state_dicts
     
     def setup(self, world_size: int):
         for dataset_filename in self.datasets_filename:
             path = self.predict_path +dataset_filename
             if os.path.exists(path):
-                if os.environ["DL_API_OVERWRITE"] != "True":
+                if os.environ["KONFAI_OVERWRITE"] != "True":
                     accept = builtins.input("The prediction {} already exists ! Do you want to overwrite it (yes,no) : ".format(path))
                     if accept != "yes":
                         return
@@ -345,11 +387,11 @@ class Predictor(DistributedObject):
         self.model.init(autocast=False, state = State.PREDICTION)
         self.model.init_outputsGroup()
         self.model._compute_channels_trace(self.model, self.model.in_channels, None, self.gpu_checkpoints)
-        self.model.load(self._load(), init=False)
-        
-        if len(list(self.outsDataset.keys())) == 0 and len([network for network in self.model.getNetworks().values() if network.measure is not None]) == 0:
-            exit(0)
+        self.modelComposite = ModelComposite(self.model, len(MODEL().split(":")), self.combine)
+        self.modelComposite.load(self._load())
 
+        if len(list(self.outsDataset.keys())) == 0 and len([network for network in self.modelComposite.getNetworks().values() if network.measure is not None]) == 0:
+            exit(0)
         
         self.size = (len(self.gpu_checkpoints)+1 if self.gpu_checkpoints else 1)
         self.dataloader = self.dataset.getData(world_size//self.size)
@@ -358,9 +400,9 @@ class Predictor(DistributedObject):
            
            
     def run_process(self, world_size: int, global_rank: int, local_rank: int, dataloaders: list[DataLoader]):
-        model = Network.to(self.model, local_rank*self.size)
-        model = DDP(model, static_graph=True) if torch.cuda.is_available() else CPU_Model(model)
-        with _Predictor(world_size, global_rank, local_rank, self.predict_path, self.images_log, self.outsDataset, model, *dataloaders) as p:
+        modelComposite = Network.to(self.modelComposite, local_rank*self.size)
+        modelComposite = DDP(modelComposite, static_graph=True) if torch.cuda.is_available() else CPU_Model(modelComposite)
+        with _Predictor(world_size, global_rank, local_rank, self.predict_path, self.images_log, self.outsDataset, modelComposite, *dataloaders) as p:
             p.run()
 
         

konfai/trainer.py

@@ -12,8 +12,8 @@ from torch.utils.tensorboard.writer import SummaryWriter
 from torch.optim.swa_utils import AveragedModel
 import torch.distributed as dist
 
-from konfai import MODELS_DIRECTORY, CHECKPOINTS_DIRECTORY, STATISTICS_DIRECTORY, SETUPS_DIRECTORY, CONFIG_FILE, MODEL, DATE, DL_API_STATE
-from konfai.data.dataset import DataTrain
+from konfai import MODELS_DIRECTORY, CHECKPOINTS_DIRECTORY, STATISTICS_DIRECTORY, SETUPS_DIRECTORY, CONFIG_FILE, MODEL, DATE, KONFAI_STATE
+from konfai.data.data_manager import DataTrain
 from konfai.utils.config import config
 from konfai.utils.utils import State, DataLog, DistributedObject, description
 from konfai.network.network import Network, ModelLoader, NetState, CPU_Model
@@ -72,7 +72,7 @@ class _Trainer():
             self.modelEMA.module.setState(NetState.TRAIN)
 
         desc = lambda : "Training : {}".format(description(self.model, self.modelEMA))
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_training), desc = desc(), total=len(self.dataloader_training), leave=False, disable=self.global_rank != 0 and "DL_API_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_training), desc = desc(), total=len(self.dataloader_training), leave=False, disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
             for _, data_dict in batch_iter:
                 with torch.amp.autocast('cuda', enabled=self.autocast):
                     input = self.getInput(data_dict)
@@ -103,7 +103,7 @@ class _Trainer():
         desc = lambda : "Validation : {}".format(description(self.model, self.modelEMA))
         data_dict = None
         self.dataloader_validation.dataset.load()
-        with tqdm.tqdm(iterable = enumerate(self.dataloader_validation), desc = desc(), total=len(self.dataloader_validation), leave=False, disable=self.global_rank != 0 and "DL_API_CLUSTER" not in os.environ) as batch_iter:
+        with tqdm.tqdm(iterable = enumerate(self.dataloader_validation), desc = desc(), total=len(self.dataloader_validation), leave=False, disable=self.global_rank != 0 and "KONFAI_CLUSTER" not in os.environ) as batch_iter:
             for _, data_dict in batch_iter:
                 input = self.getInput(data_dict)
                 self.model(input)
@@ -205,7 +205,7 @@ class Trainer(DistributedObject):
     def __init__(   self,
                     model : ModelLoader = ModelLoader(),
                     dataset : DataTrain = DataTrain(),
-                    train_name : str = "default:name",
+                    train_name : str = "default:TRAIN_01",
                     manual_seed : Union[int, None] = None,
                     epochs: int = 100,
                     it_validation : Union[int, None] = None,
@@ -215,7 +215,7 @@ class Trainer(DistributedObject):
                     ema_decay : float = 0,
                     data_log: Union[list[str], None] = None,
                     save_checkpoint_mode: str= "BEST") -> None:
-        if os.environ["DEEP_LEANING_API_CONFIG_MODE"] != "Done":
+        if os.environ["KONFAI_CONFIG_MODE"] != "Done":
             exit(0)
         super().__init__(train_name)
         self.manual_seed = manual_seed        
@@ -292,9 +292,9 @@ class Trainer(DistributedObject):
         return (1-self.ema_decay) * averaged_model_parameter + self.ema_decay * model_parameter
     
     def setup(self, world_size: int):
-        state = State._member_map_[DL_API_STATE()]
-        if state != State.RESUME and os.path.exists(STATISTICS_DIRECTORY()+self.name+"/"):
-            if os.environ["DL_API_OVERWRITE"] != "True":
+        state = State._member_map_[KONFAI_STATE()]
+        if state != State.RESUME and os.path.exists(CHECKPOINTS_DIRECTORY()+self.name+"/"):
+            if os.environ["KONFAI_OVERWRITE"] != "True":
                 accept = input("The model {} already exists ! Do you want to overwrite it (yes,no) : ".format(self.name))
                 if accept != "yes":
                     return

konfai/utils/config.py

@@ -3,9 +3,8 @@ import ruamel.yaml
 import inspect 
 import collections
 from copy import deepcopy
-from typing import Union
+from typing import Union, Literal, get_origin, get_args
 import torch
-
 from konfai import CONFIG_FILE
 
 yaml = ruamel.yaml.YAML()
@@ -26,7 +25,7 @@ class Config():
         if not os.path.exists(self.filename):
             result = input("Create a new config file ? [no,yes,interactive] : ")
             if result in ["yes", "interactive"]:
-                os.environ["DEEP_LEANING_API_CONFIG_MODE"] = "interactive" if result == "interactive" else "default"
+                os.environ["KONFAI_CONFIG_MODE"] = "interactive" if result == "interactive" else "default"
             else:
                 exit(0)
             with open(self.filename, "w") as f:
@@ -69,7 +68,7 @@ class Config():
 
     def __exit__(self, type, value, traceback) -> None:
         self.yml.close()
-        if os.environ["DEEP_LEANING_API_CONFIG_MODE"] == "remove":
+        if os.environ["KONFAI_CONFIG_MODE"] == "remove":
             if os.path.exists(CONFIG_FILE()):
                 os.remove(CONFIG_FILE())
             return
@@ -87,22 +86,22 @@ class Config():
         except:
             result = input("\nKeep a default configuration file ? (yes,no) : ")
             if result == "yes":
-                os.environ["DEEP_LEANING_API_CONFIG_MODE"] = "default"
+                os.environ["KONFAI_CONFIG_MODE"] = "default"
             else:
-                os.environ["DEEP_LEANING_API_CONFIG_MODE"] = "remove"
+                os.environ["KONFAI_CONFIG_MODE"] = "remove"
                 exit(0)
         return default.split(":")[1] if len(default.split(":")) > 1 else default
 
     @staticmethod
     def _getInputDefault(name : str, default : Union[str, None], isList : bool = False) -> Union[list[Union[str, None]], str, None]:
         if isinstance(default, str) and (default == "default" or (len(default.split(":")) > 1 and default.split(":")[0] == "default")):
-            if os.environ["DEEP_LEANING_API_CONFIG_MODE"] == "interactive":
+            if os.environ["KONFAI_CONFIG_MODE"] == "interactive":
                 if isList:
                     list_tmp = []
                     key_tmp = "OK"
-                    while key_tmp != "!" and os.environ["DEEP_LEANING_API_CONFIG_MODE"] == "interactive":
+                    while (key_tmp != "!" and key_tmp != " ") and os.environ["KONFAI_CONFIG_MODE"] == "interactive":
                         key_tmp = Config._getInput(name, default)
-                        if key_tmp != "!":
+                        if (key_tmp != "!" and key_tmp != " "):
                             if key_tmp == "":
                                 key_tmp = default.split(":")[1] if len(default.split(":")) > 1 else default
                             list_tmp.append(key_tmp)
@@ -134,6 +133,7 @@ class Config():
                 list_tmp = [] 
                 for key in value_config:
                     list_tmp.extend(Config._getInputDefault(name, key, isList=True))
+                
 
                 value = list_tmp
                 value_config = list_tmp
@@ -155,7 +155,7 @@ class Config():
                     dict_value[key] = value_tmp
                 value = dict_value
         if isinstance(self.config, str):
-            os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] = "True"
+            os.environ['KONFAI_CONFIG_VARIABLE'] = "True"
             return None
 
         self.config[name] = value_config if value_config is not None else "None"
@@ -169,17 +169,30 @@ def config(key : Union[str, None] = None):
             if "config" in kwargs:
                 filename =  kwargs["config"]
                 if filename == None:
-                    filename = os.environ['DEEP_LEARNING_API_CONFIG_FILE']
+                    filename = os.environ['KONFAI_CONFIG_FILE']
                 else:
-                    os.environ['DEEP_LEARNING_API_CONFIG_FILE'] = filename
+                    os.environ['KONFAI_CONFIG_FILE'] = filename
                 key_tmp =  kwargs["DL_args"]+("."+key if key is not None else "") if "DL_args" in kwargs else key
                 without =  kwargs["DL_without"] if "DL_without" in kwargs else []
-                os.environ['DEEP_LEARNING_API_CONFIG_PATH'] = key_tmp
+                os.environ['KONFAI_CONFIG_PATH'] = key_tmp
                 with Config(filename, key_tmp) as config:
-                    os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] = "False"
+                    os.environ['KONFAI_CONFIG_VARIABLE'] = "False"
                     kwargs = {} 
                     for param in list(inspect.signature(function).parameters.values())[len(args):]:
+
                         annotation = param.annotation
+                        # --- support Literal ---
+                        if get_origin(annotation) is Literal:
+                            allowed_values = get_args(annotation)
+                            default_value = param.default if param.default != inspect._empty else allowed_values[0]
+                            value = config.getValue(param.name, f"default:{default_value}")
+                            if value not in allowed_values:
+                                raise ValueError(
+                                    f"[Config] Invalid value '{value}' for parameter '{param.name}'. "
+                                    f"Expected one of: {allowed_values}."
+                                )
+                            kwargs[param.name] = value
+                            continue
                         if str(annotation).startswith("typing.Union") or str(annotation).startswith("typing.Optional"):
                             for i in annotation.__args__:
                                 annotation = i
@@ -188,8 +201,24 @@ def config(key : Union[str, None] = None):
                             continue
                         if not annotation == inspect._empty:
                             if annotation not in [int, str, bool, float, torch.Tensor]:
-                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple"):
-                                    if annotation.__args__[0] in [int, str, bool, float]:
+                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple") or str(annotation).startswith("typing.List"):
+                                    elem_type = annotation.__args__[0]
+                                    values = config.getValue(param.name, param.default)
+                                    if getattr(elem_type, '__origin__', None) is Union:
+                                        valid_types = elem_type.__args__
+                                        result = []
+                                        for v in values:
+                                            for t in valid_types:
+                                                try:
+                                                    if t == torch.Tensor and not isinstance(v, torch.Tensor):
+                                                        v = torch.tensor(v)
+                                                    result.append(t(v) if t != torch.Tensor else v)
+                                                    break
+                                                except Exception:
+                                                    continue
+                                        kwargs[param.name] = result
+
+                                    elif annotation.__args__[0] in [int, str, bool, float]:
                                         values = config.getValue(param.name, param.default)
                                         kwargs[param.name] = values
                                     else:
@@ -204,9 +233,13 @@ def config(key : Union[str, None] = None):
                                     else: 
                                         raise ConfigError()
                                 else:
-                                    kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
-                                    if os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] == "True":
-                                        os.environ['DEEP_LEARNING_API_CONFIG_VARIABLE'] = "False"
+                                    try:
+                                        kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
+                                    except Exception as e:
+                                        raise ValueError("[Config] Failed to instantiate {} with type {}".format(param.name, annotation.__name__))
+                                    
+                                    if os.environ['KONFAI_CONFIG_VARIABLE'] == "True":
+                                        os.environ['KONFAI_CONFIG_VARIABLE'] = "False"
                                         kwargs[param.name] = None
                             else:
                                 kwargs[param.name] = config.getValue(param.name, param.default)

konfai/utils/dataset.py

@@ -750,7 +750,7 @@ class Dataset():
         else:
             with Dataset.File(self.filename, True, self.format) as file:
                 names = file.getNames(groups)
-        return [name for i, name in enumerate(names) if index is None or i in index]
+        return [name for i, name in enumerate(sorted(names)) if index is None or i in index]
     
     def getInfos(self, groups: str, name: str) -> tuple[list[int], Attribute]:
         if self.is_directory: