Upload training/train_dpo.py with huggingface_hub

training/train_dpo.py

@@ -23,14 +23,29 @@ try:
 except ImportError:
     print("Install required packages: pip install transformers peft trl datasets torch")
 
+# Resolve repo root so paths work on any machine (macOS, Linux, etc.)
+REPO_ROOT = Path(__file__).resolve().parent.parent
+
+def detect_device() -> str:
+    """Pick the best available accelerator."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
 # Configuration
 @dataclass
 class DPOTrainingConfig:
     """DPO training configuration"""
     base_model: str = "HuggingFaceTB/SmolLM2-360M-Instruct"
-    dataset_path: str = "/home/workspace/Projects/NeuralAI/data/train_dpo_v7.jsonl"
-    output_dir: str = "/home/workspace/Projects/NeuralAI/checkpoints/dpo_model_v7"
+    dataset_path: str = str(REPO_ROOT / "data" / "train_dpo_v14.jsonl")
+    output_dir: str = str(REPO_ROOT / "checkpoints" / "dpo_model_v14")
     adapter_path: str = ""  # Path to existing adapter if continuing training
+
+    # Hugging Face upload (set push_to_hub=True to auto-sync after training)
+    push_to_hub: bool = False
+    hub_repo: str = "Subject-Emu-5259/NeuralAI"
     
     # DPO parameters
     beta: float = 0.1  # KL penalty coefficient
@@ -39,15 +54,20 @@ class DPOTrainingConfig:
     gradient_accumulation_steps: int = 4
     max_length: int = 512
     max_prompt_length: int = 256
-    epochs: int = 1
+    epochs: int = 3
     
+    # LoRA (produces a small adapter compatible with the deployed demo)
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.05
+
     # Optimization
     warmup_ratio: float = 0.1
     weight_decay: float = 0.01
     lr_scheduler: str = "cosine"
     
     # Device
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    device: str = detect_device()
 
 
 class DPODatasetBuilder:
@@ -263,16 +283,30 @@ def train_dpo(config: DPOTrainingConfig):
     
     # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(config.base_model)
-    tokenizer.pad_token = tokenizer.eos_token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
     
     # Load base model with memory optimization
+    dtype = torch.float32 if config.device == "cpu" else torch.bfloat16
     model = AutoModelForCausalLM.from_pretrained(
         config.base_model,
-        torch_dtype=torch.float32,  # Use float32 for CPU
-        device_map=None,             # Manual device placement
+        torch_dtype=dtype,
+        device_map=None,
     ).to(config.device)
     
-    # Load existing adapter if available
+    # Wrap in LoRA so the output is a small adapter (matches the deployed demo)
+    lora_config = LoraConfig(
+        r=config.lora_r,
+        lora_alpha=config.lora_alpha,
+        lora_dropout=config.lora_dropout,
+        target_modules="all-linear",
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = PeftModel(model, lora_config)
+    model.print_trainable_parameters()
+    
+    # Load existing adapter if available (continue training)
     if config.adapter_path and Path(config.adapter_path).exists():
         print(f"Loading adapter from {config.adapter_path}")
         model = PeftModel.from_pretrained(model, str(config.adapter_path), is_trainable=True)
@@ -324,12 +358,16 @@ def train_dpo(config: DPOTrainingConfig):
         gradient_accumulation_steps=config.gradient_accumulation_steps,
         use_cpu=(config.device == "cpu"),
         max_length=config.max_length,
+        max_prompt_length=config.max_prompt_length,
         num_train_epochs=config.epochs,
         warmup_ratio=config.warmup_ratio,
         weight_decay=config.weight_decay,
         lr_scheduler_type=config.lr_scheduler,
         save_strategy="epoch",
         logging_steps=1,
+        report_to="none",
+        push_to_hub=config.push_to_hub,
+        hub_model_id=config.hub_repo if config.push_to_hub else None,
     )
     
     # Create trainer
@@ -345,11 +383,19 @@ def train_dpo(config: DPOTrainingConfig):
     print(f"Starting DPO training on {len(dataset)} pairs...")
     trainer.train()
     
-    # Save
-    trainer.save_model(config.output_dir)
+    # Save the LoRA adapter (not the full base model)
+    Path(config.output_dir).mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(config.output_dir)
     tokenizer.save_pretrained(config.output_dir)
     
-    print(f"DPO model saved to {config.output_dir}")
+    print(f"DPO LoRA adapter saved to {config.output_dir}")
+    
+    # Optionally push the adapter straight to the Hub
+    if config.push_to_hub:
+        print(f"Pushing adapter to Hugging Face: {config.hub_repo}")
+        model.push_to_hub(config.hub_repo, commit_message="DPO LoRA update")
+        tokenizer.push_to_hub(config.hub_repo, commit_message="DPO LoRA update")
+        print("✅ Adapter pushed to Hugging Face.")
     
     return trainer
 
@@ -358,14 +404,21 @@ def main():
     """Main entry point"""
     config = DPOTrainingConfig()
     
-    # Check if running interactively
     import argparse
-    parser = argparse.ArgumentParser(description="DPO Training")
+    parser = argparse.ArgumentParser(description="NeuralAI DPO Training")
     parser.add_argument("--generate-only", action="store_true", help="Only generate preference dataset")
-    parser.add_argument("--beta", type=float, default=0.1)
-    parser.add_argument("--epochs", type=int, default=3)
-    parser.add_argument("--lr", type=float, default=5e-5)
-    parser.add_argument("--data", type=str, default="/home/workspace/Projects/NeuralAI/data/train_dpo_v8.jsonl")
+    parser.add_argument("--beta", type=float, default=config.beta)
+    parser.add_argument("--epochs", type=int, default=config.epochs)
+    parser.add_argument("--lr", type=float, default=config.learning_rate)
+    parser.add_argument("--data", type=str, default=config.dataset_path,
+                        help="Path to DPO jsonl (prompt/chosen/rejected)")
+    parser.add_argument("--output", type=str, default=config.output_dir,
+                        help="Where to save the LoRA adapter")
+    parser.add_argument("--adapter", type=str, default="",
+                        help="Existing adapter dir to continue training from")
+    parser.add_argument("--push", action="store_true",
+                        help="Push the trained adapter to the Hugging Face Hub")
+    parser.add_argument("--hub-repo", type=str, default=config.hub_repo)
     args = parser.parse_args()
     
     if args.generate_only:
@@ -378,6 +431,10 @@ def main():
     config.epochs = args.epochs
     config.learning_rate = args.lr
     config.dataset_path = args.data
+    config.output_dir = args.output
+    config.adapter_path = args.adapter
+    config.push_to_hub = args.push
+    config.hub_repo = args.hub_repo
     
     train_dpo(config)