|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Verify candidate Hugging Face datasets for the binary/multi-class/multi-label CLI |
| 3 | +classification benchmark, without downloading full data -- uses |
| 4 | +load_dataset_builder().info to inspect schema and split sizes. |
| 5 | +""" |
| 6 | +from __future__ import annotations |
| 7 | + |
| 8 | +import json |
| 9 | + |
| 10 | +from datasets import ClassLabel, Sequence, Value, load_dataset_builder |
| 11 | + |
| 12 | +LABEL_COL_CANDIDATES = ("label", "labels", "topic", "answer", "category", "class") |
| 13 | + |
| 14 | +# (hf_id, config_or_None, expected_type) |
| 15 | +CANDIDATES = [ |
| 16 | + # --- binary (SetFit re-hosts + originals known to load) --- |
| 17 | + ("stanfordnlp/imdb", None, "binary"), |
| 18 | + ("fancyzhx/yelp_polarity", None, "binary"), |
| 19 | + ("fancyzhx/amazon_polarity", None, "binary"), |
| 20 | + ("stanfordnlp/sst2", None, "binary"), |
| 21 | + ("SetFit/sst2", None, "binary"), |
| 22 | + ("SetFit/imdb", None, "binary"), |
| 23 | + ("SetFit/amazon_polarity", None, "binary"), |
| 24 | + ("SetFit/enron_spam", None, "binary"), |
| 25 | + ("SetFit/subj", None, "binary"), |
| 26 | + ("SetFit/CR", None, "binary"), |
| 27 | + ("SetFit/SentEval-CR", None, "binary"), |
| 28 | + ("SetFit/mrpc", None, "binary"), |
| 29 | + ("SetFit/qqp", None, "binary"), |
| 30 | + ("SetFit/qnli", None, "binary"), |
| 31 | + ("SetFit/rte", None, "binary"), |
| 32 | + ("SetFit/wnli", None, "binary"), |
| 33 | + ("SetFit/hate_speech18", None, "binary"), |
| 34 | + ("SetFit/hate_speech_offensive", None, "binary"), |
| 35 | + ("SetFit/ethos_binary", None, "binary"), |
| 36 | + ("SetFit/toxic_conversations", None, "binary"), |
| 37 | + ("SetFit/toxic_conversations_50k", None, "binary"), |
| 38 | + ("SetFit/insincere-questions", None, "binary"), |
| 39 | + ("SetFit/ade_corpus_v2_classification", None, "binary"), |
| 40 | + ("SetFit/onestop_english", None, "multiclass"), |
| 41 | + ("SetFit/wsc_fixed", None, "binary"), |
| 42 | + ("nyu-mll/glue", "sst2", "binary"), |
| 43 | + ("nyu-mll/glue", "qqp", "binary"), |
| 44 | + ("cardiffnlp/tweet_eval", "offensive", "binary"), |
| 45 | + ("cardiffnlp/tweet_eval", "emotion", "multiclass"), |
| 46 | + ("cornell-movie-review-data/rotten_tomatoes", None, "binary"), |
| 47 | + # --- multi-class --- |
| 48 | + ("fancyzhx/ag_news", None, "multiclass"), |
| 49 | + ("fancyzhx/dbpedia_14", None, "multiclass"), |
| 50 | + ("community-datasets/yahoo_answers_topics", None, "multiclass"), |
| 51 | + ("dair-ai/emotion", None, "multiclass"), |
| 52 | + ("SetFit/20_newsgroups", None, "multiclass"), |
| 53 | + ("SetFit/emotion", None, "multiclass"), |
| 54 | + ("SetFit/ag_news", None, "multiclass"), |
| 55 | + ("SetFit/bbc-news", None, "multiclass"), |
| 56 | + ("SetFit/sst5", None, "multiclass"), |
| 57 | + ("SetFit/yelp_review_full", None, "multiclass"), |
| 58 | + ("SetFit/TREC-QC", None, "multiclass"), |
| 59 | + ("SetFit/student-question-categories", None, "multiclass"), |
| 60 | + ("SetFit/tweet_eval_stance", None, "multiclass"), |
| 61 | + ("SetFit/amazon_massive_scenario_en-US", None, "multiclass"), |
| 62 | + ("SetFit/amazon_massive_intent_en-US", None, "multiclass"), |
| 63 | + ("SetFit/amazon_reviews_multi_en", None, "multiclass"), |
| 64 | + ("SetFit/ethos", None, "multiclass"), |
| 65 | + # --- multi-label --- |
| 66 | + ("google-research-datasets/go_emotions", "simplified", "multilabel"), |
| 67 | + ("google-research-datasets/go_emotions", "raw", "multilabel"), |
| 68 | + ("SetFit/go_emotions", None, "multilabel"), |
| 69 | + ("argilla/go_emotions_multi-label", None, "multilabel"), |
| 70 | + ("owaiskha9654/PubMed_MultiLabel_Text_Classification_Dataset_MeSH", None, "multilabel"), |
| 71 | + ("google/civil_comments", None, "multilabel"), |
| 72 | + ("google/jigsaw_toxicity_pred", None, "multilabel"), |
| 73 | + ("mteb/toxic_conversations_50k", None, "binary"), |
| 74 | + ("Arsive/toxicity_classification_jigsaw", None, "multilabel"), |
| 75 | + ("SetFit/ethos", "multilabel", "multilabel"), |
| 76 | +] |
| 77 | + |
| 78 | +results = [] |
| 79 | +for hf_id, config, expected in CANDIDATES: |
| 80 | + entry = {"hf_id": hf_id, "config": config, "expected": expected} |
| 81 | + try: |
| 82 | + builder = load_dataset_builder(hf_id, config) if config else load_dataset_builder(hf_id) |
| 83 | + info = builder.info |
| 84 | + features = info.features or {} |
| 85 | + splits = info.splits |
| 86 | + train_split = None |
| 87 | + if splits: |
| 88 | + for name in ("train", "training"): |
| 89 | + if name in splits: |
| 90 | + train_split = name |
| 91 | + break |
| 92 | + train_rows = splits[train_split].num_examples if train_split else None |
| 93 | + entry["train_rows"] = train_rows |
| 94 | + entry["columns"] = list(features.keys()) |
| 95 | + |
| 96 | + label_col = None |
| 97 | + label_type = None # "single" or "multi" |
| 98 | + num_classes = None |
| 99 | + for name in LABEL_COL_CANDIDATES: |
| 100 | + if name in features: |
| 101 | + feat = features[name] |
| 102 | + label_col = name |
| 103 | + if isinstance(feat, ClassLabel): |
| 104 | + label_type, num_classes = "single", feat.num_classes |
| 105 | + elif isinstance(feat, Sequence) and isinstance(feat.feature, ClassLabel): |
| 106 | + label_type, num_classes = "multi", feat.feature.num_classes |
| 107 | + elif isinstance(feat, Sequence): |
| 108 | + label_type = "multi" |
| 109 | + elif isinstance(feat, (Value,)) and feat.dtype in ("int64", "int32", "bool"): |
| 110 | + label_type = "single" # plain int/bool label column (e.g. SetFit style) |
| 111 | + break |
| 112 | + entry["label_col"] = label_col |
| 113 | + entry["label_type"] = label_type |
| 114 | + entry["num_classes"] = num_classes |
| 115 | + entry["ok"] = bool(train_rows and train_rows >= 10000 and label_col) |
| 116 | + entry["error"] = None |
| 117 | + except Exception as exc: |
| 118 | + entry["ok"] = False |
| 119 | + entry["error"] = f"{type(exc).__name__}: {exc}"[:150] |
| 120 | + results.append(entry) |
| 121 | + status = "OK" if entry.get("ok") else "FAIL" |
| 122 | + print(f"[{status}] {hf_id} ({config}) expected={expected} -> " |
| 123 | + f"rows={entry.get('train_rows')} label_col={entry.get('label_col')} " |
| 124 | + f"label_type={entry.get('label_type')} classes={entry.get('num_classes')} " |
| 125 | + f"err={entry.get('error')}", flush=True) |
| 126 | + |
| 127 | +with open("benchmarks/_dataset_verification.json", "w") as f: |
| 128 | + json.dump(results, f, indent=2) |
| 129 | + |
| 130 | +ok = [r for r in results if r["ok"]] |
| 131 | +print(f"\n{len(ok)}/{len(results)} candidates verified OK") |
| 132 | +for kind in ("binary", "multiclass", "multilabel"): |
| 133 | + matching = [r for r in ok if r["expected"] == kind] |
| 134 | + print(f" {kind}: {len(matching)} verified") |
0 commit comments