Upload train.py with huggingface_hub

train.py

@@ -173,7 +173,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--hub_model_id",
         type=str,
-        default="gr8monk3ys/paper-classifier-model",
+        default="gr8monk3ys/paper-classifier",
         help="Repository id on the HuggingFace Hub (default: %(default)s).",
     )
 
@@ -189,18 +189,19 @@ def build_label_mappings(label_names: list[str]) -> tuple[dict, dict]:
 
 def load_and_prepare_dataset(
     dataset_name: str,
-    label2id: dict[str, int],
+    fallback_label_names: list[str],
     max_train_samples: int | None = None,
     max_eval_samples: int | None = None,
-) -> DatasetDict:
+) -> tuple[DatasetDict, list[str]]:
     """Load the dataset and normalise the label column.
 
     The function handles two common dataset layouts:
-      1. The dataset already has train / validation / test splits and a
-         numeric ``label`` column whose values match our ``label2id``.
-      2. The dataset has a string ``label`` column that needs mapping.
+      1. The dataset already has a ``ClassLabel`` label column — its own
+         class names are authoritative and are returned as-is.
+      2. The dataset has a string ``label`` column, which is mapped via
+         ``fallback_label_names``.
 
-    Returns a ``DatasetDict`` with ``train`` and ``validation`` splits.
+    Returns ``(DatasetDict with train/validation splits, label_names)``.
     """
     logger.info("Loading dataset: %s", dataset_name)
     raw = load_dataset(dataset_name, trust_remote_code=True)
@@ -226,6 +227,17 @@ def load_and_prepare_dataset(
         label_col = sample_columns[-1]
     logger.info("Using label column: '%s'", label_col)
 
+    # Prefer the dataset's own class names over any hardcoded list — a
+    # mismatched hardcoded list makes the ClassLabel cast below explode.
+    orig_label_feature = next(iter(raw.values())).features[label_col]
+    if isinstance(orig_label_feature, ClassLabel):
+        label_names = list(orig_label_feature.names)
+        logger.info("Using dataset's ClassLabel names (%d classes).", len(label_names))
+    else:
+        label_names = list(fallback_label_names)
+        logger.info("Dataset has no ClassLabel; using fallback names (%d).", len(label_names))
+    label2id = {label: idx for idx, label in enumerate(label_names)}
+
     # Rename columns so downstream code can rely on 'text' and 'label' ---
     def _rename(example):
         return {"text": str(example[text_col]), "label": example[label_col]}
@@ -249,9 +261,7 @@ def load_and_prepare_dataset(
         raw = raw.filter(lambda ex: ex["label"] != -1)
 
     # Ensure we have a ClassLabel feature --------------------------------
-    label_feature = ClassLabel(
-        num_classes=len(label2id), names=list(label2id.keys())
-    )
+    label_feature = ClassLabel(num_classes=len(label_names), names=label_names)
     raw = raw.cast_column("label", label_feature)
 
     # Build train / validation splits ------------------------------------
@@ -274,7 +284,7 @@ def load_and_prepare_dataset(
         len(raw["train"]),
         len(raw["validation"]),
     )
-    return raw
+    return raw, label_names
 
 
 def tokenize_dataset(
@@ -348,19 +358,19 @@ def main() -> None:
     device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
     logger.info("Using device: %s", device)
 
-    # Label mappings
-    label2id, id2label = build_label_mappings(LABEL_NAMES)
-    num_labels = len(LABEL_NAMES)
-    logger.info("Number of labels: %d", num_labels)
-
-    # Dataset
-    dataset = load_and_prepare_dataset(
+    # Dataset (label names come from the dataset itself when available)
+    dataset, label_names = load_and_prepare_dataset(
         dataset_name=args.dataset_name,
-        label2id=label2id,
+        fallback_label_names=LABEL_NAMES,
         max_train_samples=args.max_train_samples,
         max_eval_samples=args.max_eval_samples,
     )
 
+    # Label mappings
+    label2id, id2label = build_label_mappings(label_names)
+    num_labels = len(label_names)
+    logger.info("Number of labels: %d", num_labels)
+
     # Tokenizer
     logger.info("Loading tokenizer: %s", MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -405,7 +415,7 @@ def main() -> None:
         args=training_args,
         train_dataset=tokenized_dataset["train"],
         eval_dataset=tokenized_dataset["validation"],
-        tokenizer=tokenizer,
+        processing_class=tokenizer,
         compute_metrics=build_compute_metrics_fn(),
         callbacks=[
             EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience),