github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/aistore/sdk/dataset/label_attribute.py (about) 1 # 2 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. 3 # 4 5 from typing import Callable, Tuple 6 from aistore.sdk.dataset.config_attribute import ConfigAttribute 7 8 9 # pylint: disable=too-few-public-methods 10 class LabelAttribute(ConfigAttribute): 11 """ 12 Defines the labeling attributes for a dataset, mapping filenames to corresponding labels 13 14 Args: 15 name (str): The name of the attribute, used as a reference and identifier in the dataset 16 label_identifier (Callable[[str], str]): A function that maps a file name to a label 17 """ 18 19 def __init__(self, name: str, label_identifier: Callable[[str], str]): 20 self.name = name 21 self.label_identifier = label_identifier 22 23 def get_data_for_entry(self, filename: str) -> Tuple[str, any]: 24 """ 25 Get the data for a given filename 26 27 Args: 28 filename (str): The name of the file to retrieve data for 29 30 Returns: 31 Tuple[str, any]: A tuple containing the sample key and the data for the given filename 32 """ 33 data = self.label_identifier(filename) 34 key = self.name 35 return key, data