feat(train): add bc5cdr/ncbi_disease aliases; robust dataset loader and retry

app.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi, create_repo
 
 
 DEFAULT_BASE_MODEL = "dmis-lab/biobert-base-cased-v1.2"
-DEFAULT_DATASET = "conll2003"  # fallback; medical sets may require custom preprocessing
+DEFAULT_DATASET = "wikiann:en"  # robust default; medical sets may require custom preprocessing
 TARGET_REPO = os.getenv("MEDVLLM_TARGET_REPO", "Junaidi-AI/med-vllm")
 
 
@@ -29,6 +29,8 @@ def _train_ner_lora(
     Minimal LoRA token-classification trainer.
     Uses conll2003 by default to be robust in Spaces. Extend to medical datasets later.
     """
+    # Avoid importing any local dataset scripts even if present in working dir
+    os.environ.setdefault("HF_DATASETS_DISABLE_LOCAL_IMPORTS", "1")
     from datasets import load_dataset
     from transformers import (
         AutoTokenizer,
@@ -49,8 +51,39 @@ def _train_ner_lora(
 
     set_seed(42)
 
-    log(f"Loading dataset: {dataset_name}")
-    ds = load_dataset(dataset_name)
+    ds_spec = (dataset_name or "").strip()
+    log(f"Loading dataset: {ds_spec}")
+    # Support optional config via 'name:config' (e.g., 'wikiann:en')
+    try:
+        # Medical aliases -> BigBio NER configs
+        alias_map = {
+            "bc5cdr": ("bigbio/bc5cdr", "bigbio_ner"),
+            "ncbi_disease": ("bigbio/ncbi_disease", "bigbio_ner"),
+        }
+        lower_spec = ds_spec.lower()
+        if lower_spec in alias_map:
+            ds_name, ds_config = alias_map[lower_spec]
+            log(f"Using alias mapping: {ds_spec} -> {ds_name}:{ds_config}")
+            ds = load_dataset(ds_name, ds_config)
+        elif ":" in ds_spec:
+            ds_name, ds_config = [s.strip() for s in ds_spec.split(":", 1)]
+            ds = load_dataset(ds_name, ds_config)
+        else:
+            ds = load_dataset(ds_spec)
+    except Exception as e:
+        # Fallback: if it looks like 'name:config' but was treated as a local path, try explicit two-arg call
+        err_msg = str(e)
+        log(f"Dataset load failed: {err_msg}")
+        if ":" in ds_spec:
+            try:
+                ds_name, ds_config = [s.strip() for s in ds_spec.split(":", 1)]
+                log(f"Retrying with split name/config: {ds_name}, {ds_config}")
+                ds = load_dataset(ds_name, ds_config)
+            except Exception as e2:
+                log(f"Retry failed: {e2}")
+                raise
+        else:
+            raise
 
     if "train" not in ds:
         raise RuntimeError("Dataset must have a train split")