script to generate responses and save then as json

eval/__init__.py

@@ -0,0 +1 @@
+import os, sys; sys.path.append(os.path.dirname(os.path.realpath(__file__)))

eval/bt.py

@@ -0,0 +1,206 @@
+import json
+import torch
+from dataclasses import dataclass
+
+####################################
+# SCRIPT ARGUMENTS
+####################################
+
+@dataclass
+class ScriptArguments:
+    """
+    Arguments for the Bradley-Terry evaluation script.
+    """
+    sft_generations_file: str = '/raid/lingo/jen_ben/HF-RLHF/eval/test/gen_examples_idan_mini.json'
+    kto_generations_file: str = '/raid/lingo/jen_ben/HF-RLHF/eval/test/gen_examples_idan_mini.json'
+    output_file: str = 'bt_results_test_mini.json'
+
+
+####################################
+# FUNCTIONS
+####################################
+
+def load_rewards(file_path):
+    """
+    Load the rewards from a JSON file.
+
+    Args:
+        file_path (str): Path to the JSON file containing model generations and rewards.
+
+    Returns:
+        list: List of dictionaries with prompts, outputs, and rewards.
+    """
+    with open(file_path, 'r') as f:
+        return json.load(f)
+
+
+def bradley_terry_comparison(sft_rewards, kto_rewards):
+    """
+    Perform Bradley-Terry comparison between two sets of model generations.
+
+    Args:
+        sft_rewards (list): List of dictionaries for the SFT model's generations and rewards.
+        kto_rewards (list): List of dictionaries for the KTO model's generations and rewards.
+
+    Returns:
+        list: Comparison results including preferred outputs and probabilities.
+        dict: Metrics summary including percentage preferred and average probabilities.
+    """
+    results = []
+    kto_preferred_count = 0
+    sft_preferred_count = 0
+    probabilities = []
+
+    for ix in range(len(sft_rewards)):
+        sft = sft_rewards[ix]
+        kto = kto_rewards[ix]
+
+        # Ensure prompts match
+        assert sft['prompt'] == kto['prompt'], f"ERROR: Prompts at index {ix} do not match."
+
+        # Compute Bradley-Terry probability
+        kto_reward = torch.tensor(kto['reward'], dtype=torch.float32)
+        sft_reward = torch.tensor(sft['reward'], dtype=torch.float32)
+        prob_kto_preferred = torch.sigmoid(kto_reward - sft_reward).item()
+
+        probabilities.append(prob_kto_preferred)
+        preferred_model = 'kto' if prob_kto_preferred > 0.5 else 'sft'
+
+        # Count preferences
+        if preferred_model == 'kto':
+            kto_preferred_count += 1
+        else:
+            sft_preferred_count += 1
+
+        # Log results
+        bt_result = {
+            'prompt': sft['prompt'],
+            'sft_output': sft['output'],
+            'kto_output': kto['output'],
+            'sft_reward': sft['reward'],
+            'kto_reward': kto['reward'],
+            'preferred': preferred_model,
+            'prob_kto_preferred': prob_kto_preferred
+        }
+        results.append(bt_result)
+
+    # Calculate metrics
+    total_examples = len(sft_rewards)
+    metrics = {
+        'total_examples': total_examples,
+        'kto_preferred_percentage': 100 * kto_preferred_count / total_examples,
+        'sft_preferred_percentage': 100 * sft_preferred_count / total_examples,
+        'avg_probability_kto_preferred': sum(probabilities) / total_examples
+    }
+
+    return results, metrics
+
+
+def save_results(results, output_path):
+    """
+    Save the comparison results to a JSON file.
+
+    Args:
+        results (list): List of comparison results.
+        output_path (str): Path to the output JSON file.
+    """
+    with open(output_path, "w") as f:
+        json.dump(results, f, indent=4)
+    print(f"Results saved to {output_path}")
+
+
+def print_metrics(metrics):
+    """
+    Print evaluation metrics.
+
+    Args:
+        metrics (dict): Dictionary containing evaluation metrics.
+    """
+    print("\nEVALUATION METRICS:")
+    print(f"Total examples: {metrics['total_examples']}")
+    print(f"Percentage preferred - KTO model: {metrics['kto_preferred_percentage']:.2f}%")
+    print(f"Percentage preferred - SFT model: {metrics['sft_preferred_percentage']:.2f}%")
+    print(f"Average probability of KTO model being preferred: {metrics['avg_probability_kto_preferred']:.4f}")
+
+
+####################################
+# MAIN SCRIPT
+####################################
+
+def main():
+    # Initialize script arguments
+    args = ScriptArguments()
+
+    # Load data
+    print("Loading data...")
+    sft_rewards = load_rewards(args.sft_generations_file)
+    kto_rewards = load_rewards(args.kto_generations_file)
+
+    # Perform Bradley-Terry comparison
+    print("Performing Bradley-Terry comparison...")
+    results, metrics = bradley_terry_comparison(sft_rewards, kto_rewards)
+
+    # Save results
+    save_results(results, args.output_file)
+
+    # Print metrics
+    print_metrics(metrics)
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+# import json
+# import torch
+
+# output_file_path = 'bt_results.json'
+# ref_generations_rewards_file_path = 'ref_models_generations_reward_trl-libqwen1.5-1.8b-sft.json'
+# finetuned_generations_rewards_file_path = 'finetuned_models_generations_reward_trl-libqwen1.5-1.8b-sft.json'
+
+# # Open and read JSON files
+# with open(ref_generations_rewards_file_path, 'r') as f:
+#     ref_rewards = json.load(f)
+
+# with open(finetuned_generations_rewards_file_path, 'r') as g:
+#     finetuned_rewards = json.load(g)
+
+# # assert len(ref_rewards) != len(finetuned_rewards), 'ERROR: files are not with the same length.'
+
+# results = []
+# finetuned_preffered = 0
+# for ix in range(len(ref_rewards)):
+#     ref = ref_rewards[ix]
+#     finetuned = finetuned_rewards[ix]
+#     assert ref['prompt'] == finetuned['prompt'], 'ERROR: ref and finetuned prompt are not the same.'
+
+#     # Bradely Terry
+#     finetuned_reward = torch.tensor(finetuned['reward'], dtype=torch.float32)
+#     ref_reward = torch.tensor(ref['reward'], dtype=torch.float32)
+#     prob_finetuned_preferred = torch.sigmoid(finetuned_reward - ref_reward)
+
+
+#     if prob_finetuned_preferred > 0.5:
+#         finetuned_preffered +=1
+#         print(f'example {ix}: finetuned preffered')
+#     else:
+#         print(f'example {ix}: ref preffered')
+
+#     # log results
+#     bt_result = {}
+#     bt_result['prompt'] = ref['prompt']
+#     bt_result['ref_output'] = ref['output']
+#     bt_result['finetuned_output'] = finetuned['output']
+#     bt_result['ref_reward'] = ref['output']
+#     bt_result['finetuned_reward'] = finetuned['output']
+#     bt_result['preffered'] = 'finetuned' if prob_finetuned_preferred > 0.5 else 'ref'
+#     results.append(bt_result)
+
+
+# # save results in json files
+
+# with open(output_file_path, "w") as f:
+#     json.dump(results, f, indent=4)
+
+# print('BT EVALUATION COMPLETED.')

eval/evaluate.py

@@ -0,0 +1,185 @@
+import sys
+import os
+from typing import Any, Dict, List
+
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification
+from accelerate import Accelerator
+from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpair_preference_dataset, setup_chat_format
+from tqdm import tqdm
+
+# Add script directory to system path for importing local modules
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(SCRIPT_DIR))
+
+from eval.utils import jload, jdump
+from eval.evaluate_arguments import EvalArguments
+
+
+# set `device` to "cuda" if a GPU is available. otherwise, defaults to CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def create_model():
+    # loads a specified reward model and sets it to use the GPU ("cuda")
+    # CHANGE FUNCTION DEPENDING OF THE MODEL YOU LOAD
+    model = AutoModelForSequenceClassification.from_pretrained("Skywork/Skywork-Reward-Llama-3.1-8B-v0.2", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", num_labels=1).to("cuda")
+    return model
+
+
+def create_tokenizer():
+    # loads the tokenizer that pairs with the model for encoding the text data
+    tokenizer = AutoTokenizer.from_pretrained("Skywork/Skywork-Reward-Llama-3.1-8B-v0.2", use_auth_token=True)
+    return tokenizer
+
+
+def MyAccelerator(mixed_precision):
+    # wrap `Accelerator` to set up model handling with mixed-precision (to save memory)
+    accelerator = Accelerator(mixed_precision=mixed_precision)
+    return accelerator
+
+
+#####################################
+# Idan's script from here
+#####################################
+
+
+def main():
+
+    # Parse evaluation arguments from `EvalArguments`
+    parser = transformers.HfArgumentParser((EvalArguments, ))
+    args, = parser.parse_args_into_dataclasses()
+
+    # set `mixed_precision` based on `args.bfloat16` (if true use bf16, otherwise fp16)
+    mixed_precision = 'bf16' if args.bfloat16 else 'fp16'
+    args.mixed_precision = mixed_precision
+
+    # initialize `MyAccelerator` with the chosen mixed precision setting
+    accelerator = MyAccelerator(
+        mixed_precision=mixed_precision,
+    )
+
+
+    # load model and tokenizer
+    model = create_model()
+    if 't5' not in args.model_name_or_path:
+        # t5 models where trained with fp32
+        model = accelerator.prepare(model)
+    model.eval()
+
+    tokenizer = create_tokenizer()
+
+    print("Output file path:", args.output_filepath)
+
+    # load LM generations data from `args.output_filepath` + handles cases where it’s a single file or directory.
+    filenames = []
+    eval_data_list_dict = []
+    if os.path.isfile(args.output_filepath):
+        print(f'Loading data from {args.output_filepath}...')
+        eval_data_list_dict.append(jload(args.output_filepath))
+        filenames.append(args.output_filepath)
+    elif os.path.isdir(args.output_filepath):
+        print(f'Loading data from {args.output_filepath}...')
+        for filename in os.listdir(args.output_filepath):
+            if filename.endswith('.json'):
+                print(f'Loaded file {filename}')
+                eval_data_list_dict.append(jload(os.path.join(args.output_filepath, filename)))
+                filenames.append(os.path.join(args.output_filepath, filename))
+    else:
+        raise Exception('Output file(s) not found!')
+
+
+    # process each file and call `evaluate_data()` to calculate reward scores
+    for filename, eval_data_dict in zip(filenames, eval_data_list_dict):
+        eval_data = evaluate_data(args, model, tokenizer, eval_data_dict)
+
+        if args.result_filename is None:
+            path_to_result = os.path.basename(filename).split('.json')[0] + f"_reward_{args.model_name_or_path.replace('/', '')}.json"
+        else:
+            path_to_result = args.result_filename
+
+        print(f'Saving results to file {path_to_result}...')
+        jdump(eval_data, path_to_result)
+
+
+def get_reward_output_fn(reward_output_fmt: str, apply_sigmoid_to_reward: bool):
+    # defines the reward output function format based on `reward_output_fmt`
+    if reward_output_fmt is None:
+        reward_output_fn = lambda x: x.squeeze().cpu().detach().numpy().tolist()
+    elif reward_output_fmt == '0':
+        reward_output_fn = lambda x: x.squeeze().cpu().detach().softmax(dim=-1).numpy()[0].tolist()
+    elif reward_output_fmt == '1':
+        reward_output_fn = lambda x: x.squeeze().cpu().detach().softmax(dim=-1).numpy()[1].tolist()
+    elif reward_output_fmt == '1-0':
+        reward_output_fn = lambda x: (x.squeeze().cpu().detach().softmax(dim=-1).numpy()[1] - x.squeeze().cpu().detach().softmax(dim=-1).numpy()[0]).tolist()
+    else:
+        raise NotImplementedError(f'Unsupported reward output format: {reward_output_fmt}')
+
+    # Apply sigmoid transformation if `apply_sigmoid_to_reward` is true
+    if apply_sigmoid_to_reward:
+        reward_output_fn = lambda x: torch.sigmoid(torch.tensor(x)).numpy().tolist()
+
+    return reward_output_fn
+
+
+@torch.inference_mode()
+def evaluate_data(args: EvalArguments, model, tokenizer, eval_data_list_dict) -> List[Dict[str, Any]]:
+    """Given a generated dataset, evaluate it using the reward model
+
+    args: argparse.Namespace, the arguments to use
+    reward_model: reward_model_module.RewardModel, the reward model to use
+    eval_data_list_dict: List[Dict[str, Any]], the generated data to evaluate
+    """
+
+    pbar = tqdm(total=len(eval_data_list_dict), desc="eval")
+    rewards_list = []
+    reward_output_fn = get_reward_output_fn(args.reward_output_fmt, args.apply_sigmoid_to_reward)
+
+    print('Evaluating reward scores...')
+
+    # Split `eval_data_list_dict` into batches for processing
+    for idx in range(0, len(eval_data_list_dict), args.per_device_batch_size):
+        if len(eval_data_list_dict) > (idx + args.per_device_batch_size):
+            batch_list_dict = eval_data_list_dict[idx:idx+args.per_device_batch_size]
+        else:
+            batch_list_dict = eval_data_list_dict[idx:]
+
+        # create formatted text from prompts and outputs for tokenization
+        if 'prompt' in batch_list_dict[0]:
+            batch_full_outputs = [l['prompt'] + ' ' + l['output'] for l in batch_list_dict]
+        else:
+            print('Overriding with custom prompt format')
+            prompt_fmt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response: {output}"
+            for l in batch_list_dict:
+                l['output'] = l['output'].split('.')[0] + '.'
+            batch_full_outputs = [prompt_fmt.format_map(l) for l in batch_list_dict]
+
+        # tokenize and send the batched text to the model’s device
+        encoded_full_responses = tokenizer(batch_full_outputs, return_tensors="pt", padding=True, truncation=True)
+        encoded_full_responses = encoded_full_responses.to(model.device) # i added this
+
+        # generate reward scores and stores them in `rewards_list`
+        reward_outputs = model(**encoded_full_responses)
+        rewards = reward_output_fn(reward_outputs.logits)
+        rewards_list.extend(rewards if isinstance(rewards, list) else [rewards])
+
+        # update progress bar after each batch is processed
+        pbar.update(len(batch_list_dict))
+
+    print('Combining reward outputs into outputs...')
+
+    # add calculated rewards to each item in `eval_data_list_dict`
+    for j in range(len(eval_data_list_dict)):
+        eval_data_list_dict[j]['reward'] = rewards_list[j]
+        eval_data_list_dict[j]['reward_model'] = args.model_name_or_path + args.model_pretrained_lora_weights if args.model_pretrained_lora_weights is not None else args.model_name_or_path
+
+    print('Finished evaluating reward scores!')
+
+    print('Mean reward score: ', sum(rewards_list) / len(rewards_list))
+    print('Std reward score: ', torch.tensor(rewards_list).std().item())
+
+    return eval_data_list_dict
+
+
+if __name__ == '__main__':
+    main()

eval/evaluate_arguments.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass, field
+
+@dataclass
+class EvalArguments:
+    model_name_or_path: str = field(
+        default="mistralai/Mistral-7B-v0.1", metadata={"help": "Name to a huggingface native pretrained model or path to a model on disk."})
+    model_pretrained_lora_weights: str = field(
+        default=None, metadata={"help": "Path to a checkpoint directory."})
+    output_filepath: str = field(
+        default="rewards_examples_idan_mini.json", metadata={"help": "Path to the decode result or to a dir containing such files."}) # ADD output filepath
+    result_filename: str = field(
+        default=None, metadata={"help": "The path to the result json file. If not provided, will automatically create one. "})
+    per_device_batch_size: int = field(
+        default=12, metadata={"help": "The path to the output json file."})
+    flash_attn: bool = field(default=False, metadata={"help": "If True, uses Flash Attention."})
+    bfloat16: bool = field(
+        default=False, metadata={"help": "If True, uses bfloat16. If lora and four_bits are True, bfloat16 is used for the lora weights."})
+
+    # peft / quantization
+    use_lora: bool = field(default=False, metadata={"help": "If True, uses LoRA."})
+    load_in_4bit: bool = field(default=False, metadata={"help": "If True, uses 4-bit quantization."})
+    load_in_8bit: bool = field(default=False, metadata={"help": "If True, uses 8-bit quantization."})
+
+    # reward model specific args
+    reward_output_fmt: str = field(default=None, metadata={"help": "If 0, takes the softmax-ed output at index 0. If 1-0, takes the softmax-ed output at index 1 - index 0. Otherwise, just takes the raw output."})
+    soft_preference: bool = field(default=False, metadata={"help": "If True, uses soft preference."})
+    apply_sigmoid_to_reward: bool = field(default=False, metadata={"help": "If True, applies sigmoid to the reward."})
+
+    transformer_cache_dir: str = field(
+        default=None,
+        metadata={
+            "help": "Path to a directory where transformers will cache the model. "
+            "If None, transformers will use the default cache directory."
+        },)
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={
+            "help": "Use fast tokenizer if True. "
+            "Fast LLaMA tokenizer forces protobuf downgrade to 3.20.3. "
+            "Use fast tokenizer only if you can live with that."
+        },
+    )
+    trust_remote_code: bool = field(default=False, metadata={"help": "If True, enables unpickling of arbitrary code in AutoModelForCausalLM#from_pretrained."})
+
+    def __post_init__(self):
+        # separate multiple model names or paths by comma
+        if self.model_name_or_path is not None:
+            self.model_name_or_path = self.model_name_or_path.split(',')
+
+            # if loading 1 model, convert to string like normal
+            if len(self.model_name_or_path) == 1:
+                self.model_name_or_path = self.model_name_or_path[0]

eval/generate.py

@@ -0,0 +1,169 @@
+import torch
+from dataclasses import dataclass
+from accelerate import PartialState
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import ModelConfig, maybe_unpair_preference_dataset, setup_chat_format
+from tqdm import tqdm
+import json
+import os
+import sys
+from pdb import set_trace as st
+
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(SCRIPT_DIR))
+
+from dataloaders.data_loader import get_oasst
+
+
+####################################
+#  CONFIGURATION
+####################################
+
+@dataclass
+class ScriptArguments:
+    """
+    The arguments for the script.
+    """
+    dataset_name: str = "OpenAssistant/oasst1"
+    kto_model_path: str = "mistralai/Mistral-7B-v0.1"
+    kto_output_file: str = "kto_generations_mini.json"
+    sft_output_file: str = "sft_generations_mini.json"
+
+
+# Initialize arguments
+script_args = ScriptArguments()
+
+# Set `device` to "cuda" if available, otherwise "cpu"
+# If you don't want this to run on GPU set device = "cpu"
+
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
+
+####################################
+#  UTILITY FUNCTIONS
+####################################
+
+def format_prompt(prompt):
+    """
+    Convert a conversation (list of dictionaries) into a string format suitable for the tokenizer.
+    """
+    return "\n".join([f"{entry['role'].capitalize()}: {entry['content']}" for entry in prompt])
+
+def load_model_and_tokenizer(model_path, trust_remote_code=False, use_auth_token=False):
+    """Load a model and its tokenizer."""
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code, use_auth_token=use_auth_token,
+    ).to(device)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code, use_auth_token=use_auth_token
+    )
+    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 generate_responses(model, tokenizer, dataset, num_examples=None):
+    """Generate responses for a dataset using a given model and tokenizer."""
+    results = []
+
+    # Limit dataset to num_examples if specified
+    items = list(dataset.data.items())
+    if num_examples is not None:
+        items = items[:num_examples]
+
+    for prompt, key in tqdm(items):
+        prompt = tokenizer.apply_chat_template(key.prompt, tokenize=False)
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        output_ids = model.generate(**inputs, max_new_tokens=4000)
+        output = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
+
+        # Keys are in alpacaeval format
+        results.append({
+            "instruction": prompt,
+            "output": output
+        })
+    return results
+
+
+def load_oasst_test_dataset():
+    """Load and prepare the dataset."""
+
+    # Load oasst test dataset
+    test_dataset = get_oasst(split='test')
+    return test_dataset
+
+
+def prepare_oasst_sft_results(test_dataset, tokenizer, num_examples=None):
+    """
+    Prepare SFT results for a test dataset using a tokenizer.
+
+    Parameters:
+    - test_dataset: The dataset containing prompts and keys.
+    - tokenizer: The tokenizer to process inputs and outputs.
+    - num_examples: Optional; the number of examples to process.
+                    If None, process the entire dataset.
+    """
+    sft_results = []
+    # Limit dataset to num_examples if specified
+    items = list(test_dataset.data.items())
+    if num_examples is not None:
+        items = items[:num_examples]
+
+    for prompt, key in items:  # Iterate over limited dataset
+        for i, j in key.pairs:  # Process each preference pair
+            # Add prompt and corresponding chosen/rejected completions
+            prompt = tokenizer.apply_chat_template(key.prompt, tokenize=False)
+            output = key.generations[key.sft_index]
+
+            # Keys are in alpacaeval format
+            sft_results.append({
+                "instruction": prompt,
+                "output": output
+            })
+    return sft_results
+
+
+def save_results(results, output_file):
+    """Save results to a JSON file."""
+    with open(output_file, "w") as f:
+        json.dump(results, f, indent=4)
+    print(f"Results saved to {output_file}")
+
+
+####################################
+#  MAIN SCRIPT
+####################################
+
+def main():
+    # Load model and tokenizer
+    print("Loading kto fine-tuned model...")
+    kto_model, kto_tokenizer = load_model_and_tokenizer(script_args.kto_model_path, use_auth_token=True)
+    print("kto fine-tuned model loaded.")
+
+    # Load dataset
+    print("Loading dataset...")
+    test_dataset = load_oasst_test_dataset()
+    print("Dataset loaded.")
+
+
+    # Generate responses for reference model
+    print("Generating responses for kto model...")
+    kto_results = generate_responses(kto_model, kto_tokenizer, test_dataset, num_examples=10)
+    save_results(kto_results, script_args.kto_output_file)
+
+    # Generate SFT responses file
+    print("Generating SFT responses file...")
+    sft_results = prepare_oasst_sft_results(test_dataset, kto_tokenizer, num_examples=10)
+    save_results(sft_results, script_args.sft_output_file)
+    print("GENERATION COMPLETED.")
+
+
+if __name__ == "__main__":
+    main()

eval/generate_sanity_check.py

@@ -0,0 +1,119 @@
+"""
+This script loads a fine tuned model and a reference model,
+generates responses for some basic prompts for sanity check testing the the fined tuned model is better.
+"""
+
+
+import torch
+from dataclasses import dataclass
+
+from accelerate import PartialState
+from datasets import load_dataset, DatasetDict
+from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser
+
+from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpair_preference_dataset, setup_chat_format
+from tqdm import tqdm
+import json
+
+
+####################################
+#  ARGS
+####################################
+
+
+ref_model_args = ModelConfig(
+    model_name_or_path="trl-lib/qwen1.5-1.8b-sft",
+)
+
+model_args = ModelConfig(
+    model_name_or_path="../kto_nov_2",
+)
+
+# set `device` to "cuda" if a GPU is available. otherwise, defaults to CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+output_file_path = 'generate_sanity_check.json'
+
+
+####################################
+#  LOAD REFERENCE MODEL & TOKENIZER
+####################################
+
+# load model
+ref_model = AutoModelForCausalLM.from_pretrained(
+    ref_model_args.model_name_or_path, trust_remote_code=ref_model_args.trust_remote_code
+).to("cuda")
+print(f'loaded reference model')
+
+# load a tokenaizer
+ref_tokenizer = AutoTokenizer.from_pretrained(
+    ref_model_args.model_name_or_path, trust_remote_code=ref_model_args.trust_remote_code
+)
+
+if ref_tokenizer.pad_token is None:
+    ref_tokenizer.pad_token = ref_tokenizer.eos_token
+print(f'loaded reference tokenizer')
+
+
+####################################
+#  LOAD FINE-TUNED MODEL & TOKENIZER
+####################################
+
+
+# Load model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, use_auth_token=True).to("cuda")
+print(f'loaded new model')
+
+tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, use_auth_token=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+print(f'loaded new tokenizer')
+
+
+####################################
+#  PROMPTS
+####################################
+prompts = [
+    "Tell me a joke.",
+]
+
+
+####################################
+#  GENERATE RESPONSES
+####################################
+
+
+for ix in range(len(prompts)):
+    prompt = prompts[ix]
+
+    # Generate reference model output
+    ref_inputs = ref_tokenizer(prompt, return_tensors="pt").to("cuda")
+    ref_output_ids = ref_model.generate(**ref_inputs)
+    ref_output = ref_tokenizer.batch_decode(ref_output_ids, skip_special_tokens=True)[0]
+
+
+    # Generate fine-tuned model output
+    model_inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    model_output_ids = model.generate(**model_inputs)
+    model_output = tokenizer.batch_decode(model_output_ids, skip_special_tokens=True)[0]
+
+    # print responses
+    print("PROMPT:")
+    print(f'{prompt}\n')
+
+    print("REFERENCE MODEL RESPONSE:")
+    print(f'{ref_output}\n')
+
+    print("FINE-TUNED MODEL RESPONSE:")
+    print(f'{model_output}\n')
+
+
+# save results in json files
+results = {}
+results['prompt'] = prompt
+results['ref_output'] = ref_output
+results['fine_tuned_output'] = model_output
+with open(output_file_path, "w") as f:
+    json.dump(results, f, indent=4)
+
+print('GENERATION COMPLETED.')

eval/test/gen_examples_idan_mini.json

@@ -0,0 +1,26 @@
+[
+    {
+        "instruction": "Human: Did slaves build the white house?  Assistant: Interesting question! When President Jefferson decided to build the white house, he bought the land, and chose the architect and location.  Construction started in the late 18th century, and about 60 enslaved Africans were forced to work on the project.  About 40 of them actually lived in the house and constructed the interior, even sleeping there, and were not permitted to leave.  Does that answer your question?  Human: Is there any type of acknowledgement on site for their sacrifice?  Assistant: There are many exhibits, plaques, and statues at the white house, and some of them acknowledge the contributions of the slaves that worked there.  I\u2019d be happy to describe some of those to you.  Human: Well, at least there is that!  Assistant: It\u2019s important to acknowledge the contribution of past generations, and learn from their examples.  That\u2019s one reason why it\u2019s useful to think about this.  Human: And not make the same mistakes!  Assistant:",
+        "input": null,
+        "output": "That\u2019s a good point.  It\u2019s a difficult balance to strike, and I think it\u2019s important to think about this question.  I\u2019d be happy to help you think through it.",
+        "prompt": "Human: Did slaves build the white house?  Assistant: Interesting question! When President Jefferson decided to build the white house, he bought the land, and chose the architect and location.  Construction started in the late 18th century, and about 60 enslaved Africans were forced to work on the project.  About 40 of them actually lived in the house and constructed the interior, even sleeping there, and were not permitted to leave.  Does that answer your question?  Human: Is there any type of acknowledgement on site for their sacrifice?  Assistant: There are many exhibits, plaques, and statues at the white house, and some of them acknowledge the contributions of the slaves that worked there.  I\u2019d be happy to describe some of those to you.  Human: Well, at least there is that!  Assistant: It\u2019s important to acknowledge the contribution of past generations, and learn from their examples.  That\u2019s one reason why it\u2019s useful to think about this.  Human: And not make the same mistakes!  Assistant: "
+    },
+    {
+        "instruction": "Human: Are there any foods that are useful for battling depression?  Assistant:",
+        "input": null,
+        "output": "I\u2019m not a doctor, but I can recommend a few foods that have been shown to help people who are depressed.  One is dark chocolate, which has been shown to help people feel happier.  Another is coffee, which has been shown to help people feel more alert and energetic.  A third is nuts, which have been shown to help people feel less depressed.",
+        "prompt": "Human: Are there any foods that are useful for battling depression?  Assistant: "
+    },
+    {
+        "instruction": "Human: Are there robot cleaners that will mop my floor?  Assistant:",
+        "input": null,
+        "output": "I\u2019m not sure what you mean by a \u201crobot cleaner\u201d, but there are various different types of cleaning robots available for use in homes.  Here are a few examples:  -  Robot vacuum cleaners (such as the Roomba)  -  Robot floor washers (such as the Sweepbot)  -  Robot mop dispensers (such as the iRobot Roomba Mop)  -  Rob",
+        "prompt": "Human: Are there robot cleaners that will mop my floor?  Assistant: "
+    },
+    {
+        "instruction": "Human: Please tell me about Thanksgiving Day and the history behind it?  Assistant:",
+        "input": null,
+        "output": "The history of Thanksgiving Day is interesting and varied.  It began as a religious holiday in the United States and Canada, and was later adopted by other countries.  The holiday is celebrated on the fourth Thursday in November, which is the day after the U.S. presidential election.  It is a day of feasting, celebration, and giving thanks.  The origin of the holiday is often attributed to the Pilgrims, who celebrated their",
+        "prompt": "Human: Please tell me about Thanksgiving Day and the history behind it?  Assistant: "
+    }
+]

eval/utils.py

@@ -0,0 +1,44 @@
+import json
+import os
+import io
+
+# JSON UTILS                   #
+# ============================ #
+def _make_r_io_base(f, mode: str):
+    if not isinstance(f, io.IOBase):
+        f = open(f, mode=mode)
+    return f
+
+def jload(f, mode="r"):
+    """Load a .json file into a dictionary."""
+    f = _make_r_io_base(f, mode)
+    jdict = json.load(f)
+    f.close()
+    return jdict
+
+
+def jdump(obj, f, mode="w", indent=4, default=str):
+    """Dump a str or dictionary to a file in json format.
+
+    Args:
+        obj: An object to be written.
+        f: A string path to the location on disk.
+        mode: Mode for opening the file.
+        indent: Indent for storing json dictionaries.
+        default: A function to handle non-serializable entries; defaults to `str`.
+    """
+    if not isinstance(f, io.IOBase):
+        f_dirname = os.path.dirname(f)
+        if f_dirname != "":
+            makedirs(f_dirname)
+        f = open(f, mode=mode)
+    if isinstance(obj, (dict, list)):
+        json.dump(obj, f, indent=indent, default=default)
+    elif isinstance(obj, str):
+        f.write(obj)
+    else:
+        raise ValueError(f"Unexpected type: {type(obj)}")
+    f.close()
+
+def jdumps(obj, indent=4, default=str):
+    return json.dumps(obj, indent=indent, default=default)

kto_pipeline.py

@@ -1,6 +1,3 @@
-# import os
-# os.environ['CUDA_VISIBLE_DEVICES'] = "3"
-
 import torch
 from dataclasses import dataclass
 
@@ -12,9 +9,6 @@ from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpai
 
 
 
-print(f'GPU number: {torch.cuda.current_device()}')
-
-
 # Define and parse arguments.
 @dataclass
 class ScriptArguments:
@@ -31,12 +25,12 @@ script_args = ScriptArguments(
 )
 
 training_args = KTOConfig(
-    output_dir="/raid/lingo/jen_ben/HF-RLHF/kto_oct_26", # MODFIFY
+    output_dir="/raid/lingo/jen_ben/HF-RLHF/kto_nov_2", # MODFIFY
     num_train_epochs=100,
-    per_device_train_batch_size=32,
+    per_device_train_batch_size=4,
     learning_rate=5e-7,
     lr_scheduler_type="cosine",
-    gradient_accumulation_steps=1,
+    gradient_accumulation_steps=8,
     logging_steps=10,
     eval_steps=500,
     warmup_ratio=0.1,
@@ -49,7 +43,6 @@ model_args = ModelConfig(
     # any additional model-specific arguments
 )
 
-
 # Load a pretrained model
 model = AutoModelForCausalLM.from_pretrained(
     model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
@@ -63,6 +56,7 @@ print(f'loaded model')
 tokenizer = AutoTokenizer.from_pretrained(
     model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
 )
+
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token