add: Level 2 training for 1.5B model after Level 1

train.py

@@ -175,8 +175,101 @@ trainer.train()
 wandb.finish()
 print("Training done!")
 
+# Save Level 1 checkpoint
+model.save_pretrained("deceit-1.5b-l1")
+tokenizer.save_pretrained("deceit-1.5b-l1")
+print("Level 1 checkpoint saved locally")
+
+# Load Level 2 dataset
+import pathlib as _pl2
+data_path_l2 = _pl2.Path("/home/trainer/.local/lib/python3.10/site-packages/deceit_env/data/level2.jsonl")
+questions_l2 = []
+with open(data_path_l2) as f:
+    for line in f:
+        line = line.strip()
+        if line:
+            questions_l2.append(json.loads(line))
+
+print(f"Loaded {len(questions_l2)} Level 2 questions")
+
+def make_prompt_l2(q, distractors):
+    context = "\n".join(distractors)
+    msgs = [
+        {"role":"system","content":SYSTEM_PROMPT},
+        {"role":"user","content":f"Question: {q}\n\nContext:\n{context}\n\nTurn 1 of 3. Respond in JSON."},
+    ]
+    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+
+train_dataset_l2 = Dataset.from_list([
+    {"prompt": make_prompt_l2(q["question"], q.get("distractors", [])), "question": q["question"]}
+    for q in questions_l2
+])
+
+# Update env to level 2
+_env_l2 = DeceitEnvironment(grader=_grader)
+
+def reward_fn_l2(completions, prompts=None, **kwargs):
+    rewards = []
+    for text in completions:
+        try:
+            parsed = parse_action(text)
+        except:
+            parsed = FAIL.copy()
+        try:
+            with _env_lock:
+                obs = _env_l2.reset(level=2)
+                current = parsed
+                total = 0.0
+                for turn in range(obs.max_turns):
+                    if turn == obs.max_turns - 1:
+                        current["is_final"] = True
+                    action = DeceitAction(
+                        reasoning=current.get("reasoning",""),
+                        answer=current.get("answer",""),
+                        confidence=float(current.get("confidence",0.5)),
+                        abstain=bool(current.get("abstain",False)),
+                        is_final=bool(current.get("is_final",True)),
+                    )
+                    result = _env_l2.step(action)
+                    total += result.reward
+                    if result.done:
+                        break
+        except Exception as e:
+            print(f"L2 Episode error: {e}")
+            total = -1.3
+        rewards.append(total)
+    return rewards
+
+# Train Level 2
+print("Starting Level 2 training on 1.5B...")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level2")
+
+trainer_l2 = GRPOTrainer(
+    model=model,
+    processing_class=tokenizer,
+    reward_funcs=[reward_fn_l2],
+    args=GRPOConfig(
+        output_dir="./deceit-1.5b-l2",
+        max_steps=80,
+        per_device_train_batch_size=4,
+        num_generations=4,
+        learning_rate=2e-6,
+        warmup_steps=5,
+        logging_steps=1,
+        save_steps=40,
+        report_to="wandb",
+        max_completion_length=256,
+        remove_unused_columns=False,
+    ),
+    train_dataset=train_dataset_l2,
+)
+trainer_l2.train()
+wandb.finish()
+print("Level 2 training done!")
+
+# Save final model
 model.save_pretrained("deceit-1.5b-final")
 tokenizer.save_pretrained("deceit-1.5b-final")
 model.push_to_hub(HF_REPO_ID)
 tokenizer.push_to_hub(HF_REPO_ID)
-print(f"Saved to {HF_REPO_ID}")
+print(f"Final model saved to {HF_REPO_ID}")