updated to only load preference data

kto_dataset_processor.py

@@ -5,54 +5,61 @@ from pdb import set_trace as st
 
 def process_dataset_ultrafeedback():
     """
-    Processes the 'HuggingFaceH4/ultrafeedback_binarized' dataset into a unified train and test split.
+    Processes the 'train_prefs' and 'test_prefs' splits of the 'HuggingFaceH4/ultrafeedback_binarized' dataset
+    into a unified format for preference modeling.
 
     Returns:
         dict: A dictionary containing the unified 'train' and 'test' splits of the dataset in the KTO format.
               Each split is a Hugging Face Dataset object.
     """
-    # Load the dataset
+    # Load the relevant splits of the dataset
     dataset_name = "HuggingFaceH4/ultrafeedback_binarized"
-    dataset = load_dataset(dataset_name)
+    train_prefs = load_dataset(dataset_name, split="train_prefs")
+    test_prefs = load_dataset(dataset_name, split="test_prefs")
 
     # Function to transform a single example into the desired schema
     def transform_data(example):
         data_points = []
         # Chosen completion
         chosen_completion = example["chosen"][1]["content"]
-        data_points.append({
-            "prompt": example["prompt"],
-            "completion": chosen_completion.strip(),
-            "label": True
-        })
+        if chosen_completion.strip():  # Check for non-empty completions
+            data_points.append({
+                "prompt": example["prompt"],
+                "completion": chosen_completion.strip(),
+                "label": True
+            })
         # Rejected completion
         rejected_completion = example["rejected"][1]["content"]
-        data_points.append({
-            "prompt": example["prompt"],
-            "completion": rejected_completion.strip(),
-            "label": False
-        })
+        if rejected_completion.strip():  # Check for non-empty completions
+            data_points.append({
+                "prompt": example["prompt"],
+                "completion": rejected_completion.strip(),
+                "label": False
+            })
         return data_points
 
-    # Combine splits into unified train and test sets
+    # Process train and test splits
     train_data = []
     test_data = []
 
-    for split_name, split_data in dataset.items():
-        if "train" in split_name:
-            for example in split_data:
-                train_data.extend(transform_data(example))
-        elif "test" in split_name:
-            for example in split_data:
-                test_data.extend(transform_data(example))
+    for example in train_prefs:
+        train_data.extend(transform_data(example))
 
-    # Convert unified data to Hugging Face Dataset
-    unified_train = Dataset.from_pandas(pd.DataFrame(train_data))
-    unified_test = Dataset.from_pandas(pd.DataFrame(test_data))
+    for example in test_prefs:
+        test_data.extend(transform_data(example))
+
+    # Convert unified data to DataFrames
+    train_df = pd.DataFrame(train_data)
+    test_df = pd.DataFrame(test_data)
+
+
+    # Convert to Hugging Face Dataset
+    unified_train = Dataset.from_pandas(train_df)
+    unified_test = Dataset.from_pandas(test_df)
 
     return {"train": unified_train, "test": unified_test}
 
 
-# if __name__ == "__main__":
-#     kto_dataset = process_dataset_ultrafeedback()
-#     st()
+if __name__ == "__main__":
+    kto_dataset = process_dataset_ultrafeedback()
+    st()

kto_pipeline.py

@@ -2,7 +2,7 @@ import torch
 from dataclasses import dataclass
 from accelerate import PartialState
 from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser
-from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config
+from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpair_preference_dataset, setup_chat_format
 from kto_dataset_processor import process_dataset_ultrafeedback
 from datetime import datetime
 import wandb
@@ -78,8 +78,45 @@ def load_model_and_tokenizer(model_args):
     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 find_unknown_tokens(tokenizer, texts):
+#     """
+#     Identify tokens in the dataset that are not in the tokenizer's vocabulary.
+#     """
+#     all_tokens = set()
+#     for text in texts:
+#         tokens = tokenizer.tokenize(text)
+#         all_tokens.update(tokens)
+#     vocab = set(tokenizer.get_vocab().keys())
+#     unknown_tokens = all_tokens - vocab
+#     return unknown_tokens
+
+
+# def add_tokens_to_tokenizer(tokenizer, model, dataset):
+#     """
+#     Extend the tokenizer's vocabulary with missing tokens and resize the model embeddings.
+#     """
+#     # Extract all texts from the dataset
+#     texts = [example["completion"] for example in dataset["train"]]
+
+#     # Identify unknown tokens
+#     unknown_tokens = find_unknown_tokens(tokenizer, texts)
+#     print(f"Found {len(unknown_tokens)} unknown tokens: {list(unknown_tokens)[:10]}...")
+
+#     # Add unknown tokens to tokenizer
+#     tokenizer.add_tokens(list(unknown_tokens))
+#     model.resize_token_embeddings(len(tokenizer))
+#     print(f"Tokenizer vocabulary size after extension: {len(tokenizer)}")
+
+
 ####################################
 #  MAIN LOGIC
 ####################################
@@ -99,6 +136,11 @@ def main():
     dataset = process_dataset_ultrafeedback()
     print("Dataset processed.")
 
+    # # Extend tokenizer with missing tokens
+    # print("Adding unknown tokens to tokenizer...")
+    # add_tokens_to_tokenizer(tokenizer, model, dataset)
+    # print("Tokenizer updated.")
+
     # Initialize trainer
     print("Initializing trainer...")
     trainer = KTOTrainer(