DPO: merge SFT LoRA then train

train_dpo.py

@@ -1,74 +1,82 @@
 # /// script
-# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "torch"]
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "datasets", "transformers", "accelerate", "torch"]
 # ///
 """DPO on offline compile-ok vs compile-fail preferences (compiler-as-judge)."""
 
 import os
 
 from datasets import load_dataset
-from peft import LoraConfig
+from peft import LoraConfig, PeftModel
 from trl import DPOConfig, DPOTrainer
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 DATASET_ID = os.environ.get("DATASET_ID", "gonzalolinares/cpp-compiler-prefs")
-BASE_MODEL = os.environ.get("BASE_MODEL", "gonzalolinares/qwen25-1.5b-cpp-sft")
-FALLBACK_MODEL = os.environ.get("FALLBACK_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
+SFT_ADAPTER = os.environ.get("BASE_MODEL", "gonzalolinares/qwen25-1.5b-cpp-sft")
+BASE_MODEL = os.environ.get("FALLBACK_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
 HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "gonzalolinares/qwen25-1.5b-cpp-dpo")
 OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "qwen25-1.5b-cpp-dpo")
 
 
-def resolve_model(model_id: str) -> str:
+def load_policy():
+    """Load base NL model, merge SFT LoRA if present."""
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype="auto")
     try:
-        AutoTokenizer.from_pretrained(model_id)
-        return model_id
-    except Exception:
-        print(f"Could not load {model_id}, falling back to {FALLBACK_MODEL}")
-        return FALLBACK_MODEL
+        model = PeftModel.from_pretrained(model, SFT_ADAPTER)
+        model = model.merge_and_unload()
+        print(f"Merged SFT adapter from {SFT_ADAPTER}")
+    except Exception as e:
+        print(f"No SFT adapter merge ({e}); training from base {BASE_MODEL}")
+    return model, tokenizer
 
 
 def main() -> None:
-    model_id = resolve_model(BASE_MODEL)
     ds = load_dataset(DATASET_ID, split="train")
-    # Normalize to prompt/chosen/rejected text if needed
-    def to_dpo(example):
-        def as_text(x):
-            if isinstance(x, list):
-                # list of chat messages -> last assistant or join
-                parts = []
-                for m in x:
-                    if isinstance(m, dict):
-                        parts.append(f"{m.get('role', '')}: {m.get('content', '')}")
-                    else:
-                        parts.append(str(m))
-                return "\n".join(parts)
-            return str(x)
 
+    def to_dpo(example):
         prompt = example.get("prompt")
         chosen = example.get("chosen")
         rejected = example.get("rejected")
-        # Prefer conversational: if prompt is message list without assistant
+
+        def last_assistant(messages):
+            if isinstance(messages, list) and messages:
+                last = messages[-1]
+                if isinstance(last, dict):
+                    return last.get("content", str(last))
+            return str(messages)
+
+        def user_text(messages):
+            if isinstance(messages, list):
+                parts = []
+                for m in messages:
+                    if isinstance(m, dict) and m.get("role") in {"system", "user"}:
+                        parts.append(m.get("content", ""))
+                return "\n\n".join(parts)
+            return str(messages)
+
+        # TRL DPO conversational: prompt=list[messages], chosen/rejected=list with assistant
         if isinstance(prompt, list) and prompt and isinstance(prompt[0], dict):
-            # TRL DPO can take conversational format
-            return {
-                "prompt": prompt,
-                "chosen": chosen if isinstance(chosen, list) else [{"role": "assistant", "content": str(chosen)}],
-                "rejected": rejected
-                if isinstance(rejected, list)
-                else [{"role": "assistant", "content": str(rejected)}],
-            }
+            ch = chosen if isinstance(chosen, list) else [{"role": "assistant", "content": str(chosen)}]
+            rj = rejected if isinstance(rejected, list) else [{"role": "assistant", "content": str(rejected)}]
+            return {"prompt": prompt, "chosen": ch, "rejected": rj}
+
         return {
-            "prompt": as_text(prompt),
-            "chosen": as_text(chosen),
-            "rejected": as_text(rejected),
+            "prompt": user_text(prompt),
+            "chosen": last_assistant(chosen),
+            "rejected": last_assistant(rejected),
         }
 
-    ds = ds.map(to_dpo)
+    ds = ds.map(to_dpo, remove_columns=[c for c in ds.column_names if c not in {"prompt", "chosen", "rejected"}])
+    # keep only needed columns after map - re-add by selecting
+    keep = {"prompt", "chosen", "rejected"}
+    drop = [c for c in ds.column_names if c not in keep]
+    if drop:
+        ds = ds.remove_columns(drop)
     split = ds.train_test_split(test_size=0.1, seed=42)
 
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+    model, tokenizer = load_policy()
 
     trainer = DPOTrainer(
         model=model,