updated kto pipeline to work with general dataset

kto_pipeline.py

@@ -1,37 +1,35 @@
 import torch
 from dataclasses import dataclass
 from accelerate import PartialState
-from datasets import load_dataset, DatasetDict, Dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser
-from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpair_preference_dataset, setup_chat_format
-from dataloaders.data_loader import get_oasst
-from pdb import set_trace as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import KTOConfig, KTOTrainer, get_peft_config
+from kto_dataset_processor import process_dataset_ultrafeedback
+from datetime import datetime
 import wandb
 
-
 ####################################
 #  CONFIGURATION
 ####################################
 
 @dataclass
-class ScriptArguments:
+class Config:
     """
     Configuration for the script.
     """
-    dataset_name: str = "OpenAssistant/oasst1"  # Dataset name or path
-    output_dir: str = "/raid/lingo/jen_ben/HF-RLHF/kto_nov_24_2_epochs"  # Output directory
-    pretrained_model_name: str = "mistralai/Mistral-7B-v0.1"  # Pretrained model name or path
-    checkpoint_path: str = "/raid/lingo/jen_ben/HF-RLHF/kto_nov_24_2_epochs"  # Checkpoint path
-    push_to_hub: bool = False  # Whether to push the model to the Hugging Face hub
+    # Dataset settings
+    process_dataset_func: callable = process_dataset_ultrafeedback  # Dataset processing function
 
-@dataclass
-class TrainingArguments(KTOConfig):
-    """
-    Configuration for the KTO trainer.
-    """
-    output_dir: str = "/raid/lingo/jen_ben/HF-RLHF/kto_nov_24_2_epochs"
-    num_train_epochs: int = 2 # did 1 epoch, then maybe try 2 epochs
-    per_device_train_batch_size: int = 4 # 4 is the highes that runs well.
+    # Model settings
+    model_name: str = "HuggingFaceH4/zephyr-7b-beta"  # Pretrained model name or path
+    use_peft: bool = True
+    lora_target_modules: str = "all-linear"
+    lora_r: int = 16
+    lora_alpha: int = 16
+
+    # Training settings
+    output_dir: str = f"kto_{model_name}_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
+    num_train_epochs: int = 1
+    per_device_train_batch_size: int = 4
     learning_rate: float = 5e-7
     lr_scheduler_type: str = "cosine"
     gradient_accumulation_steps: int = 1
@@ -41,152 +39,36 @@ class TrainingArguments(KTOConfig):
     bf16: bool = True
     logging_first_step: bool = True
 
-@dataclass
-class ModelArguments(ModelConfig):
-    """
-    Configuration for the model.
-    """
-    model_name_or_path: str = "mistralai/Mistral-7B-v0.1"
-    use_peft: bool = True
-    lora_target_modules: str = "all-linear"
-    lora_r: int = 16
-    lora_alpha: int = 16
+    # Checkpoint and hub settings
+    checkpoint_path: str = None
+    push_to_hub: bool = False
 
-# Initialize configurations
-script_args = ScriptArguments()
-training_args = TrainingArguments(output_dir=script_args.output_dir)
-model_args = ModelArguments(model_name_or_path=script_args.pretrained_model_name)
+# Initialize the unified configuration
+config = Config()
 
 ####################################
 #  HELPER FUNCTIONS
 ####################################
 
-def load_model_and_tokenizer(model_args):
+def load_model_and_tokenizer(config):
     """
     Load a model and tokenizer from a specified path.
     """
     model = AutoModelForCausalLM.from_pretrained(
-        model_args.model_name_or_path,
-        trust_remote_code=model_args.trust_remote_code,
+        config.model_name,
         torch_dtype=torch.float16,
         device_map="auto"
     )
     tokenizer = AutoTokenizer.from_pretrained(
-        model_args.model_name_or_path,
-        trust_remote_code=model_args.trust_remote_code
+        config.model_name
     )
 
     # Set pad token if missing
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    # Setup chat format if not present
-    if tokenizer.chat_template is None:
-        model, tokenizer = setup_chat_format(model, tokenizer)
-
     return model, tokenizer
 
-
-
-def load_and_format_oasst_dataset(tokenizer):
-    """
-    Load, process, and format the OpenAssistant dataset into DPO-compatible format.
-
-    Args:
-        split (str): The dataset split to load ('train' or 'test').
-        tokenizer (AutoTokenizer): Tokenizer to apply chat templates.
-        num_proc (int, optional): Number of processes for parallel processing.
-
-    Returns:
-        Dataset: Processed and formatted dataset.
-    """
-
-    # Load oasst dataset
-    train_dataset = get_oasst(split='train')
-
-    # Initialize lists for DPO dataset
-    dpo_train_data = {
-        "prompt": [],
-        "chosen": [],
-        "rejected": []
-    }
-
-    # Process the dataset
-    for prompt, key in train_dataset.data.items():  # Iterate over dataset
-        if hasattr(key, "pairs") and key.pairs:  # Check if pairs exist
-            for i, j in key.pairs:  # Process each preference pair
-                # Add prompt and corresponding chosen/rejected completions
-                dpo_train_data["prompt"].append(key.prompt)
-                dpo_train_data["chosen"].append(key.generations[i])  # Chosen generation
-                dpo_train_data["rejected"].append(key.generations[j])  # Rejected generation
-
-    # Convert DPO data into a Dataset
-    dpo_train_dataset = Dataset.from_dict(dpo_train_data)
-
-    # Wrap it in a DatasetDict
-    dataset_dict = DatasetDict({
-        "train": dpo_train_dataset
-    })
-
-
-    test_dataset = get_oasst(split='test')
-
-    dpo_test_data = {
-        "prompt": [],
-        "chosen": [],
-        "rejected": []
-    }
-
-    for prompt, key in test_dataset.data.items():  # Iterate over dataset
-        if hasattr(key, "pairs") and key.pairs:  # Check if pairs exist
-            for i, j in key.pairs:  # Process each preference pair
-                # Add prompt and corresponding chosen/rejected completions
-                dpo_test_data["prompt"].append(key.prompt)
-                dpo_test_data["chosen"].append(key.generations[i])  # Chosen generation
-                dpo_test_data["rejected"].append(key.generations[j])  # Rejected generation
-
-    dpo_test_dataset = Dataset.from_dict(dpo_test_data)
-    dataset_dict["test"] = dpo_test_dataset
-
-    # If needed, reformat a DPO-formatted dataset (prompt, chosen, rejected) to a KTO-format (prompt, completion, label)
-    dataset_dict = maybe_unpair_preference_dataset(dataset_dict, num_proc=training_args.dataset_num_proc)
-    print(f'loaded dataset')
-
-
-    # Apply chat template
-    def format_dataset(example):
-        # Ensure prompt is in the correct structure
-        if isinstance(example["prompt"], str):
-            example["prompt"] = [{"role": "user", "content": example["prompt"]}]
-        elif isinstance(example["prompt"], list):
-            # If it's already a list, ensure each element has the "role" and "content" keys
-            for item in example["prompt"]:
-                if "role" not in item or "content" not in item:
-                    raise ValueError(f"Each item in 'prompt' must have 'role' and 'content': {item}")
-
-        # Ensure completion is in the correct structure
-        if isinstance(example["completion"], str):
-            example["completion"] = [{"role": "assistant", "content": example["completion"]}]  # Wrap as a list of dictionaries
-        elif isinstance(example["completion"], list):
-            # If it's already a list, ensure each element has the "role" and "content" keys
-            for item in example["completion"]:
-                if "role" not in item or "content" not in item:
-                    raise ValueError(f"Each item in 'completion' must have 'role' and 'content': {item}")
-
-        # Now apply the chat template
-        example["prompt"] = tokenizer.apply_chat_template(example["prompt"], tokenize=False)
-        example["completion"] = tokenizer.apply_chat_template(example["completion"], tokenize=False)
-
-        return example
-
-
-    # Compute that only on the main process for faster data processing.
-    # see: https://github.com/huggingface/trl/pull/1255
-    with PartialState().local_main_process_first():
-        dataset = dataset_dict.map(format_dataset, num_proc=training_args.dataset_num_proc)
-
-    return dataset
-
 ####################################
 #  MAIN LOGIC
 ####################################
@@ -197,26 +79,42 @@ def main():
 
     # Load models and tokenizer
     print("Loading models and tokenizer...")
-    model, tokenizer = load_model_and_tokenizer(model_args)
-    ref_model, _ = load_model_and_tokenizer(model_args)
+    model, tokenizer = load_model_and_tokenizer(config)
+    ref_model, _ = load_model_and_tokenizer(config)
     print("Models and tokenizer loaded.")
 
-    # Load and process datasets
-    print("Loading, processing, and formatting dataset...")
-    dataset = load_and_format_oasst_dataset(
-        tokenizer=tokenizer,
-    )
+    # Load and process datasets using the specified function
+    print("Processing dataset...")
+    dataset = config.process_dataset_func()
+    print("Dataset processed.")
 
     # Initialize trainer
     print("Initializing trainer...")
     trainer = KTOTrainer(
         model=model,
         ref_model=ref_model,
-        args=training_args,
+        args=KTOConfig(
+            output_dir=config.output_dir,
+            num_train_epochs=config.num_train_epochs,
+            per_device_train_batch_size=config.per_device_train_batch_size,
+            learning_rate=config.learning_rate,
+            lr_scheduler_type=config.lr_scheduler_type,
+            gradient_accumulation_steps=config.gradient_accumulation_steps,
+            logging_steps=config.logging_steps,
+            eval_steps=config.eval_steps,
+            warmup_ratio=config.warmup_ratio,
+            bf16=config.bf16,
+            logging_first_step=config.logging_first_step,
+        ),
         train_dataset=dataset["train"],
         eval_dataset=dataset["test"],
         tokenizer=tokenizer,
-        peft_config=get_peft_config(model_args),
+        peft_config=get_peft_config({
+            "use_peft": config.use_peft,
+            "lora_target_modules": config.lora_target_modules,
+            "lora_r": config.lora_r,
+            "lora_alpha": config.lora_alpha,
+        }),
     )
 
     # Training
@@ -232,26 +130,12 @@ def main():
     trainer.save_metrics("eval", metrics)
 
     # Log metrics to wandb
-    wandb.log({
-        "epoch": metrics.get("epoch"),
-        "grad_norm": metrics.get("grad_norm"),
-        "kl": metrics.get("kl"),
-        "learning_rate": metrics.get("learning_rate"),
-        "logits/chosen": metrics.get("logits/chosen"),
-        "logits/rejected": metrics.get("logits/rejected"),
-        "logps/chosen": metrics.get("logps/chosen"),
-        "logps/rejected": metrics.get("logps/rejected"),
-        "loss": metrics.get("loss"),
-        "rewards/chosen": metrics.get("rewards/chosen"),
-        "rewards/margins": metrics.get("rewards/margins"),
-        "rewards/rejected": metrics.get("rewards/rejected"),
-        "step": metrics.get("step")
-    })
+    wandb.log(metrics)
 
     # Save model and optionally push to hub
-    trainer.save_model(training_args.output_dir)
-    if script_args.push_to_hub:
-        trainer.push_to_hub(dataset_name=script_args.dataset_name)
+    trainer.save_model(config.output_dir)
+    if config.push_to_hub:
+        trainer.push_to_hub()
 
     print("Process completed.")