Upload train_sft.py with huggingface_hub

train_sft.py

@@ -28,7 +28,7 @@ SCENARIO_FOLDER  = "iac"               # subdir under app/tasks/scenarios/
 MODEL_NAME       = "unsloth/Qwen2.5-7B-Instruct-bnb-4bit"
 MAX_SEQ_LEN      = 1536
 MAX_NEW_TOKENS   = 600
-NUM_EPOCHS       = 5
+NUM_EPOCHS       = 3
 LEARNING_RATE    = 5e-5
 LORA_RANK        = 16
 GRAD_ACCUM_STEPS = 2
@@ -194,11 +194,10 @@ def main():
 
     scenario_ids = [
         "iac_002", "iac_004", "iac_005", "iac_007",
-        "iac_009", "iac_010", "iac_011", "iac_014",
-        "iac_015", "iac_016", "iac_017", "iac_018",
-        "iac_019", "iac_020", "iac_021",
-        "iac_022", "iac_023",
-        "iac_024",
+        "iac_010", "iac_014",
+        "iac_016", "iac_018",
+        "iac_019", "iac_020",
+        "iac_022", "iac_023", "iac_024",
     ]
     examples = []
     for sid in scenario_ids:
@@ -216,7 +215,7 @@ def main():
             full_text = tokenizer.apply_chat_template(
                 messages, tokenize=False, add_generation_prompt=False
             )
-            examples.append({"text": full_text})
+            examples.append({"text": full_text, "scenario_id": sid})
             print(f"  Loaded {sid} ({len(gt['ground_truth'])} findings)")
         except Exception as e:
             print(f"  Skipping {sid}: {e}")
@@ -256,6 +255,16 @@ def main():
     baseline_scores = evaluate(scenario_ids, "before")
     print(f"  Baseline mean: {sum(baseline_scores.values())/len(baseline_scores):.3f}")
 
+    # Curriculum filter — only train on scenarios with baseline <= 0.5
+    # (high-baseline scenarios already work; SFT on them causes regression)
+    CURRICULUM_THRESHOLD = 0.5
+    train_sids = [s for s in scenario_ids if baseline_scores.get(s, 1.0) <= CURRICULUM_THRESHOLD]
+    print(f"\n  Curriculum filter (baseline <= {CURRICULUM_THRESHOLD}): "
+          f"{len(train_sids)}/{len(scenario_ids)} scenarios → {train_sids}")
+    filtered = [{"text": e["text"]} for e in examples if e["scenario_id"] in train_sids]
+    dataset  = Dataset.from_list(filtered)
+    print(f"  Filtered training dataset: {len(filtered)} examples")
+
     # SFT training
     print("\n[5/6] SFT training...")
     FastLanguageModel.for_training(model)