Merge pull request #660 from VijayKalmath/Add-SquadSchema-toHuggingFa…

…ceDatasetClass Add SQUAD schema to HuggingFaceDataset class
QData · Jun 8, 2022 · 9688456 · 9688456
2 parents efa21e6 + 12013cf
commit 9688456
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/textattack/datasets/huggingface_dataset.py b/textattack/datasets/huggingface_dataset.py
@@ -40,6 +40,10 @@ def get_datasets_dataset_columns(dataset):
     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"