Spaces:
Runtime error
Runtime error
| """ | |
| HuggingFaceDataset Class | |
| ========================= | |
| TextAttack allows users to provide their own dataset or load from HuggingFace. | |
| """ | |
| import collections | |
| import datasets | |
| import textattack | |
| from .dataset import Dataset | |
| def _cb(s): | |
| """Colors some text blue for printing to the terminal.""" | |
| return textattack.shared.utils.color_text(str(s), color="blue", method="ansi") | |
| def get_datasets_dataset_columns(dataset): | |
| """Common schemas for datasets found in dataset hub.""" | |
| schema = set(dataset.column_names) | |
| if {"premise", "hypothesis", "label"} <= schema: | |
| input_columns = ("premise", "hypothesis") | |
| output_column = "label" | |
| elif {"question", "sentence", "label"} <= schema: | |
| input_columns = ("question", "sentence") | |
| output_column = "label" | |
| elif {"sentence1", "sentence2", "label"} <= schema: | |
| input_columns = ("sentence1", "sentence2") | |
| output_column = "label" | |
| elif {"question1", "question2", "label"} <= schema: | |
| input_columns = ("question1", "question2") | |
| output_column = "label" | |
| elif {"question", "sentence", "label"} <= schema: | |
| input_columns = ("question", "sentence") | |
| output_column = "label" | |
| elif {"context", "question", "title", "answers"} <= schema: | |
| # Common schema for SQUAD dataset | |
| input_columns = ("title", "context", "question") | |
| output_column = "answers" | |
| elif {"text", "label"} <= schema: | |
| input_columns = ("text",) | |
| output_column = "label" | |
| elif {"sentence", "label"} <= schema: | |
| input_columns = ("sentence",) | |
| output_column = "label" | |
| elif {"document", "summary"} <= schema: | |
| input_columns = ("document",) | |
| output_column = "summary" | |
| elif {"content", "summary"} <= schema: | |
| input_columns = ("content",) | |
| output_column = "summary" | |
| elif {"label", "review"} <= schema: | |
| input_columns = ("review",) | |
| output_column = "label" | |
| else: | |
| raise ValueError( | |
| f"Unsupported dataset schema {schema}. Try passing your own `dataset_columns` argument." | |
| ) | |
| return input_columns, output_column | |
| class HuggingFaceDataset(Dataset): | |
| """Loads a dataset from 🤗 Datasets and prepares it as a TextAttack dataset. | |
| Args: | |
| name_or_dataset (:obj:`Union[str, datasets.Dataset]`): | |
| The dataset name as :obj:`str` or actual :obj:`datasets.Dataset` object. | |
| If it's your custom :obj:`datasets.Dataset` object, please pass the input and output columns via :obj:`dataset_columns` argument. | |
| subset (:obj:`str`, `optional`, defaults to :obj:`None`): | |
| The subset of the main dataset. Dataset will be loaded as :obj:`datasets.load_dataset(name, subset)`. | |
| split (:obj:`str`, `optional`, defaults to :obj:`"train"`): | |
| The split of the dataset. | |
| dataset_columns (:obj:`tuple(list[str], str))`, `optional`, defaults to :obj:`None`): | |
| Pair of :obj:`list[str]` representing list of input column names (e.g. :obj:`["premise", "hypothesis"]`) | |
| and :obj:`str` representing the output column name (e.g. :obj:`label`). If not set, we will try to automatically determine column names from known designs. | |
| label_map (:obj:`dict[int, int]`, `optional`, defaults to :obj:`None`): | |
| Mapping if output labels of the dataset should be re-mapped. Useful if model was trained with a different label arrangement. | |
| For example, if dataset's arrangement is 0 for `Negative` and 1 for `Positive`, but model's label | |
| arrangement is 1 for `Negative` and 0 for `Positive`, passing :obj:`{0: 1, 1: 0}` will remap the dataset's label to match with model's arrangements. | |
| Could also be used to remap literal labels to numerical labels (e.g. :obj:`{"positive": 1, "negative": 0}`). | |
| label_names (:obj:`list[str]`, `optional`, defaults to :obj:`None`): | |
| List of label names in corresponding order (e.g. :obj:`["World", "Sports", "Business", "Sci/Tech"]` for AG-News dataset). | |
| If not set, labels will printed as is (e.g. "0", "1", ...). This should be set to :obj:`None` for non-classification datasets. | |
| output_scale_factor (:obj:`float`, `optional`, defaults to :obj:`None`): | |
| Factor to divide ground-truth outputs by. Generally, TextAttack goal functions require model outputs between 0 and 1. | |
| Some datasets are regression tasks, in which case this is necessary. | |
| shuffle (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether to shuffle the underlying dataset. | |
| .. note:: | |
| Generally not recommended to shuffle the underlying dataset. Shuffling can be performed using DataLoader or by shuffling the order of indices we attack. | |
| """ | |
| def __init__( | |
| self, | |
| name_or_dataset, | |
| subset=None, | |
| split="train", | |
| dataset_columns=None, | |
| label_map=None, | |
| label_names=None, | |
| output_scale_factor=None, | |
| shuffle=False, | |
| ): | |
| if isinstance(name_or_dataset, datasets.Dataset): | |
| self._dataset = name_or_dataset | |
| else: | |
| self._name = name_or_dataset | |
| self._subset = subset | |
| self._dataset = datasets.load_dataset(self._name, subset)[split] | |
| subset_print_str = f", subset {_cb(subset)}" if subset else "" | |
| textattack.shared.logger.info( | |
| f"Loading {_cb('datasets')} dataset {_cb(self._name)}{subset_print_str}, split {_cb(split)}." | |
| ) | |
| # Input/output column order, like (('premise', 'hypothesis'), 'label') | |
| ( | |
| self.input_columns, | |
| self.output_column, | |
| ) = dataset_columns or get_datasets_dataset_columns(self._dataset) | |
| if not isinstance(self.input_columns, (list, tuple)): | |
| raise ValueError( | |
| "First element of `dataset_columns` must be a list or a tuple." | |
| ) | |
| self.label_map = label_map | |
| self.output_scale_factor = output_scale_factor | |
| if label_names: | |
| self.label_names = label_names | |
| else: | |
| try: | |
| self.label_names = self._dataset.features[self.output_column].names | |
| except (KeyError, AttributeError): | |
| # This happens when the dataset doesn't have 'features' or a 'label' column. | |
| self.label_names = None | |
| # If labels are remapped, the label names have to be remapped as well. | |
| if self.label_names and label_map: | |
| self.label_names = [ | |
| self.label_names[self.label_map[i]] for i in self.label_map | |
| ] | |
| self.shuffled = shuffle | |
| if shuffle: | |
| self._dataset.shuffle() | |
| def _format_as_dict(self, example): | |
| input_dict = collections.OrderedDict( | |
| [(c, example[c]) for c in self.input_columns] | |
| ) | |
| output = example[self.output_column] | |
| if self.label_map: | |
| output = self.label_map[output] | |
| if self.output_scale_factor: | |
| output = output / self.output_scale_factor | |
| return (input_dict, output) | |
| def filter_by_labels_(self, labels_to_keep): | |
| """Filter items by their labels for classification datasets. Performs | |
| in-place filtering. | |
| Args: | |
| labels_to_keep (:obj:`Union[Set, Tuple, List, Iterable]`): | |
| Set, tuple, list, or iterable of integers representing labels. | |
| """ | |
| if not isinstance(labels_to_keep, set): | |
| labels_to_keep = set(labels_to_keep) | |
| self._dataset = self._dataset.filter( | |
| lambda x: x[self.output_column] in labels_to_keep | |
| ) | |
| def __getitem__(self, i): | |
| """Return i-th sample.""" | |
| if isinstance(i, int): | |
| return self._format_as_dict(self._dataset[i]) | |
| else: | |
| # `idx` could be a slice or an integer. if it's a slice, | |
| # return the formatted version of the proper slice of the list | |
| return [ | |
| self._format_as_dict(self._dataset[j]) for j in range(i.start, i.stop) | |
| ] | |
| def shuffle(self): | |
| self._dataset.shuffle() | |
| self.shuffled = True | |