Upload 137 files

hugging/td_fuse/heal.py

@@ -66,20 +66,25 @@ def load_healing_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
     print("[heal] Loading healing fine-tune data...")
 
     # Merge-specific: use diverse data that exercises all merged capabilities
+    # Each entry: (dataset_id, config_name_or_None, split, count, text_field)
     datasets_to_load = [
-        # General language (from Pile)
-        ("EleutherAI/pile", "validation", 500, "text"),
+        # General language — same calibration data source that works reliably
+        ("neuralmagic/LLM_compression_calibration", None, "train", 500, "text"),
         # Math reasoning (exercises DeepSeek/MiMo contributions)
-        ("openai/gsm8k", "train", 300, "question"),
-        # Code (exercises Llama contribution)
-        ("codeparrot/github-code", "train", 200, "code"),
+        ("openai/gsm8k", "main", "train", 300, "question"),
+        # Code — bigcode/starcoderdata is a modern alternative
+        ("bigcode/starcoderdata", "python", "train", 200, "content"),
     ]
 
     all_texts = []
 
-    for dataset_id, split, count, text_field in datasets_to_load:
+    for entry in datasets_to_load:
+        dataset_id, config_name, split, count, text_field = entry
         try:
-            ds = load_dataset(dataset_id, split=split, streaming=True, trust_remote_code=True)
+            if config_name:
+                ds = load_dataset(dataset_id, config_name, split=split, streaming=True)
+            else:
+                ds = load_dataset(dataset_id, split=split, streaming=True)
             loaded = 0
             for example in ds:
                 if loaded >= count:
@@ -200,7 +205,7 @@ def apply_qlora_unsloth(
 
     trainer = SFTTrainer(
         model=model,
-        tokenizer=tokenizer,
+        processing_class=tokenizer,
         train_dataset=dataset,
         args=training_args,
         max_seq_length=cfg.heal_seq_len,
@@ -328,7 +333,7 @@ def apply_qlora_standard(
 
     trainer = Trainer(
         model=model,
-        tokenizer=tokenizer,
+        processing_class=tokenizer,
         train_dataset=dataset,
         args=training_args,
     )