Merge pull request #12 from riddhibhagwat/main

Evaluation Pipeline Updates & Refined Implementation

ml/dpo_pipeline.py

@@ -0,0 +1,44 @@
+from transformers import Trainer, TrainingArguments
+from datasets import Dataset
+import torch
+
+def train_dpo_model(model, dataset, learning_rate=5e-5, num_train_epochs=3, per_device_train_batch_size=16):
+    """
+    Trains a model using Direct Preference Optimization (DPO).
+
+    Args:
+        model: The language model to be trained.
+        dataset: The dataset used for training, should be in Hugging Face Dataset format.
+        learning_rate: Learning rate for the optimizer.
+        num_train_epochs: Number of epochs to train.
+        per_device_train_batch_size: Batch size per device during training.
+    """
+    model.train()
+    
+    training_args = TrainingArguments(
+        output_dir="./dpo_model",
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        learning_rate=learning_rate,
+        per_device_train_batch_size=per_device_train_batch_size,
+        per_device_eval_batch_size=per_device_train_batch_size,
+        num_train_epochs=num_train_epochs,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        logging_steps=100,
+        save_total_limit=2,
+        push_to_hub=False,
+        load_best_model_at_end=True,
+    )
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset.get("validation", None),
+    )
+    
+    trainer.train()
+    
+    return model
+

ml/eval/alpaca.py

@@ -0,0 +1,43 @@
+import openai
+import os 
+from alpaca_eval import run_evaluation
+
+def judge_responses(response1, response2, prompt):
+    """
+    Use OpenAI GPT-4 API to judge two model responses.
+    Returns: "A" if response1 is better, "B" if response2 is better, or "tie".
+    """
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+
+    prompt_text = f"""
+    Given the user prompt: "{prompt}"
+    
+    Response A: "{response1}"
+    Response B: "{response2}"
+    
+    Which response is better? Reply with 'A', 'B', or 'tie'.
+    """
+
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-4",
+            messages=[{"role": "system", "content": "You are an expert evaluator."},
+                      {"role": "user", "content": prompt_text}],
+            max_tokens=5
+        )
+        result = response["choices"][0]["message"]["content"].strip().lower()
+        return result if result in ["a", "b", "tie"] else "tie"
+    except Exception as e:
+        print(f"Error in OpenAI API call: {e}")
+        return "tie"
+
+
+
+def alpaca_evaluator(model_name, num_samples=200):
+    results = run_evaluation(
+        model=model_name,
+        num_samples=num_samples,  # fewer samples for quick testing
+        reference_model="gpt-4",  # Compare against GPT-4 (optional)
+    )
+    return results 
+

ml/eval/bt.py

@@ -11,9 +11,9 @@ 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'
+    old_generations_file: str
+    new_generations_file: str   
+    output_file: str = 'bt_results.json'
 
 
 ####################################
@@ -34,63 +34,63 @@ def load_rewards(file_path):
         return json.load(f)
 
 
-def bradley_terry_comparison(sft_rewards, kto_rewards):
+def bradley_terry_comparison(old_rewards, new_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.
+        old_rewards (list): List of dictionaries for the OLD model's generations and rewards.
+        new_rewards (list): List of dictionaries for the NEW 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
+    new_preferred_count = 0
+    old_preferred_count = 0
     probabilities = []
 
-    for ix in range(len(sft_rewards)):
-        sft = sft_rewards[ix]
-        kto = kto_rewards[ix]
+    for ix in range(len(old_rewards)):
+        old = old_rewards[ix]
+        new = new_rewards[ix]
 
         # Ensure prompts match
-        assert sft['prompt'] == kto['prompt'], f"ERROR: Prompts at index {ix} do not match."
+        assert old['prompt'] == new['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()
+        new_reward = torch.tensor(old['reward'], dtype=torch.float32)
+        old_reward = torch.tensor(new['reward'], dtype=torch.float32)
+        prob_new_preferred = torch.sigmoid(new_reward - old_reward).item()
 
-        probabilities.append(prob_kto_preferred)
-        preferred_model = 'kto' if prob_kto_preferred > 0.5 else 'sft'
+        probabilities.append(prob_new_preferred)
+        preferred_model = 'new' if prob_new_preferred > 0.5 else 'old'
 
         # Count preferences
-        if preferred_model == 'kto':
-            kto_preferred_count += 1
+        if preferred_model == 'new':
+            new_preferred_count += 1
         else:
-            sft_preferred_count += 1
+            old_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'],
+            'prompt': old['prompt'],
+            'old_output': old['output'],
+            'new_output': new['output'],
+            'old_reward': old['reward'],
+            'new_reward': new['reward'],
             'preferred': preferred_model,
-            'prob_kto_preferred': prob_kto_preferred
+            'prob_new_preferred': prob_new_preferred
         }
         results.append(bt_result)
 
     # Calculate metrics
-    total_examples = len(sft_rewards)
+    total_examples = len(old_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
+        'new_preferred_percentage': 100 * new_preferred_count / total_examples,
+        'old_preferred_percentage': 100 * old_preferred_count / total_examples,
+        'avg_probability_new_preferred': sum(probabilities) / total_examples
     }
 
     return results, metrics
@@ -118,9 +118,9 @@ def print_metrics(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}")
+    print(f"Percentage preferred - KTO model: {metrics['new_preferred_percentage']:.2f}%")
+    print(f"Percentage preferred - SFT model: {metrics['old_preferred_percentage']:.2f}%")
+    print(f"Average probability of KTO model being preferred: {metrics['avg_probability_new_preferred']:.4f}")
 
 
 ####################################
@@ -128,22 +128,17 @@ def print_metrics(metrics):
 ####################################
 
 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)
+    old_rewards = load_rewards(args.sft_generations_file)
+    new_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)
+    results, metrics = bradley_terry_comparison(old_rewards, new_rewards)
 
-    # Save results
     save_results(results, args.output_file)
-
-    # Print metrics
     print_metrics(metrics)
 
 
@@ -152,55 +147,3 @@ if __name__ == "__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.')

ml/eval/evaluate.py

@@ -1,185 +0,0 @@
-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()

ml/eval/evaluate_arguments.py

@@ -3,7 +3,7 @@ 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."})
+        default="CohereForAI/aya-expanse-8b", 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(

ml/eval/evaluation_pipeline.py

@@ -0,0 +1,58 @@
+###########
+# IMPORTS #
+###########
+from reward_eval import process_evaluation
+from generate import generate_files
+from alpaca import alpaca_evaluator, judge_responses
+from bt import bradley_terry_comparison, load_rewards
+from evaluate_arguments import EvalArguments
+import pandas as pd
+import numpy as np
+
+#############
+# EVALUATOR #
+#############
+'''
+Evaluation Pipeline 
+
+Parameters: 
+eval_dataset: list of dictionaries that contain the prompt and response in the same form as below: 
+            [{"prompt": "How are you?", "output": "I'm doing great!"}, {"prompt": "What's your name?", "output": "Assistant"}]
+reward_output_filepath: string (must end in .json) that represents the path of the output of the reward score evaluation
+model: base model that is being evaluated (defaults to starter base model - Aya-23-8B )
+all_responses: should be a path to a csv file that has all the model's responses and their corresponding prompts with the following
+                format: response1 --> col 1, response2 --> col 2, prompt --> col 3
+
+language: which language is being used for this model (needs to be a valid FeeLLanguage object once FeeLLanguage class is updated)
+'''
+def evaluator_master_fn(eval_dataset: list[dict], 
+                        reward_output_filepath: str, 
+                        all_responses: str, 
+                        language: str, 
+                        new_model,
+                        old_model="CohereForAI/aya-expanse-8b"): 
+    # language is string for now, will be an object later with FeeLLanguage class definition with specific lanugage 
+    # functionalities (will also store latest model and be much easier to handle such functions) 
+    
+    # 1. Reward score evaluation: 
+    args = EvalArguments(bfloat16=True, 
+                         reward_output_fmt='1-0', 
+                         apply_sigmoid_to_reward=False,
+                         per_device_batch_size=8,
+                         output_filepath="new_evaluation",
+                         result_filename=None,
+                         model_name_or_path=new_model)
+    reward_score_result = process_evaluation(args, model_name=new_model, eval_data_list_dict=eval_dataset)
+
+    # 2. Alpaca Eval - Judging Responses 
+    judge_df = pd.read_csv(all_responses)
+    judge_df["winner"] = judge_df.apply(lambda r: judge_responses(r["response1"], r["response2"], r["prompt"]), axis = 1) # axis = 1 -- loops rows
+
+    # 3. Alpaca Eval - model comparison 
+    alpaca_results = alpaca_evaluator(new_model, num_samples=200) # can adjust num_samples as needed, potentially based on language
+
+    # 4. Bradley Terry Evaluation
+    bt_results = bradley_terry_comparison(load_rewards(old_model), load_rewards(new_model))
+
+    return reward_score_result, judge_df, alpaca_results, bt_results
+

ml/eval/generate.py

@@ -1,6 +1,4 @@
-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
@@ -8,15 +6,10 @@ 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
 ####################################
@@ -96,7 +89,7 @@ def load_oasst_test_dataset():
     """Load and prepare the dataset."""
 
     # Load oasst test dataset
-    test_dataset = get_oasst(split='test')
+    test_dataset = load_dataset(split='test')
     return test_dataset
 
 
@@ -141,7 +134,7 @@ def save_results(results, output_file):
 #  MAIN SCRIPT
 ####################################
 
-def main():
+def generate_files():
     # 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)
@@ -166,4 +159,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    generate_files()

ml/eval/generate_sanity_check.py

@@ -45,7 +45,7 @@ ref_model = AutoModelForCausalLM.from_pretrained(
 ).to("cuda")
 print(f'loaded reference model')
 
-# load a tokenaizer
+# load a tokenizer
 ref_tokenizer = AutoTokenizer.from_pretrained(
     ref_model_args.model_name_or_path, trust_remote_code=ref_model_args.trust_remote_code
 )

ml/eval/reward_eval.py

@@ -0,0 +1,123 @@
+import sys
+import os
+from typing import Any, Dict, List
+import json 
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, CohereConfig, AutoModel
+from accelerate import Accelerator
+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(model_name: str):
+    """
+    loads pre-trained reward model and moves it onto device
+    """
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", num_labels=1).to("cuda")
+    return model
+
+
+def create_tokenizer(model_name):
+    # loads the tokenizer that pairs with the model for encoding the text data
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
+    return tokenizer
+
+
+def MyAccelerator(mixed_precision: str):
+    """
+    accelerator initialization (wrapper) for handling mixed precision
+    """
+    return Accelerator(mixed_precision=mixed_precision)
+    
+def get_reward_output_fn(reward_output_format: str, sigmoid: bool):
+    def default(x): 
+        return x.squeeze().cpu().detach().numpy().tolist()
+    reward_fn_map = {
+        '0': lambda x: x.squeeze().cpu().detach().softmax(dim=-1).numpy()[0].tolist(),
+        '1': lambda x: x.squeeze().cpu().detach().softmax(dim=-1).numpy()[1].tolist(),
+        '1-0': lambda x: (x.squeeze().cpu().detach().softmax(dim=-1).numpy()[1] - x.squeeze().cpu().detach().softmax(dim=-1).numpy()[0]).tolist()
+    }
+    reward_output_fn = reward_fn_map.get(reward_output_format, default)
+    if sigmoid: 
+        return lambda x: torch.sigmoid(torch.tensor(x)).numpy().tolist()
+    return reward_output_fn
+
+def evaluate_data(args, model, tokenizer, eval_data_list_dict) -> List[Dict[str, Any]]:
+    """
+    Evaluate the dataset using the reward model.
+    """
+    reward_output_fn = get_reward_output_fn(args.reward_output_fmt, args.apply_sigmoid_to_reward)
+    pbar = tqdm(total=len(eval_data_list_dict), desc="Evaluating Rewards")
+    rewards_list = []
+
+    for idx in range(0, len(eval_data_list_dict), args.per_device_batch_size):
+        batch_list_dict = eval_data_list_dict[idx:idx+args.per_device_batch_size]
+
+        # Create prompt-response pairs
+        batch_full_outputs = [
+            f"{l['prompt']} {l['output']}" for l in batch_list_dict
+        ] if 'prompt' in batch_list_dict[0] else [f"Below is an instruction: {l['instruction']} Response: {l['output']}" for l in batch_list_dict]
+
+        # Tokenize reponse and send to device
+        encoded_full_responses = tokenizer(batch_full_outputs, return_tensors="pt", padding=True, truncation=True)
+        encoded_full_responses = encoded_full_responses.to(model.device)
+
+        # Generate rewards
+        with torch.inference_mode():
+            reward_outputs = model(**encoded_full_responses)
+            rewards = reward_output_fn(reward_outputs.logits)
+            rewards_list.extend(rewards)
+
+        pbar.update(len(batch_list_dict))
+
+    # Adding reward scores to original data
+    for i, data in enumerate(eval_data_list_dict):
+        data['reward'] = rewards_list[i]
+
+    return eval_data_list_dict
+
+def process_evaluation(args, model_name: str, eval_data_list_dict) -> List[Dict[str, Any]]:
+    """
+    Main function for processing evaluation, takes model name as input.
+    """
+    # mixed_precision = 'bf16' if args.bfloat16 else 'fp16'
+    
+    # Initialize accelerator and model
+    # accelerator = MyAccelerator(mixed_precision)
+    model = create_model(model_name)
+    tokenizer = create_tokenizer(model_name)
+
+    model.eval()
+
+    eval_data = evaluate_data(args, model, tokenizer, eval_data_list_dict)
+
+    result_filename = args.result_filename or f"{os.path.basename(args.output_filepath).split('.')[0]}_reward_results.json"
+    with open(result_filename, "w") as f:
+        json.dump(eval_data, f)
+
+    return eval_data
+
+
+# ONLY FOR TESTING: 
+if __name__ == '__main__':
+    args = EvalArguments(bfloat16=True, 
+                         reward_output_fmt='1-0', 
+                         apply_sigmoid_to_reward=False,
+                         per_device_batch_size=8,
+                         output_filepath= '/path/to/your/data.json',
+                         result_filename=None,
+                         model_name_or_path="CohereForAI/aya-expanse-8b")
+
+
+    eval_data_list_dict = [{"prompt": "How are you?", "output": "I'm doing great!"}, {"prompt": "What's your name?", "output": "Assistant"}]
+
+    process_evaluation(args, model_name="CohereForAI/aya-expanse-8b", eval_data_list_dict=eval_data_list_dict)

ml/feel.yaml

@@ -0,0 +1,176 @@
+name: feel
+channels:
+  - conda-forge
+  - bioconda
+  - defaults
+dependencies:
+  - accelerate=1.4.0=pyhd8ed1ab_0
+  - aiohappyeyeballs=2.4.6=pyhd8ed1ab_0
+  - aiohttp=3.11.13=py313ha9b7d5b_0
+  - aiosignal=1.3.2=pyhd8ed1ab_0
+  - annotated-types=0.7.0=pyhd8ed1ab_1
+  - anyio=4.8.0=pyhd8ed1ab_0
+  - attrs=25.1.0=pyh71513ae_0
+  - aws-c-auth=0.8.1=hfc2798a_0
+  - aws-c-cal=0.8.1=hc8a0bd2_3
+  - aws-c-common=0.10.6=h5505292_0
+  - aws-c-compression=0.3.0=hc8a0bd2_5
+  - aws-c-event-stream=0.5.0=h54f970a_11
+  - aws-c-http=0.9.2=h96aa502_4
+  - aws-c-io=0.15.3=haba67d1_6
+  - aws-c-mqtt=0.11.0=h24f418c_12
+  - aws-c-s3=0.7.9=hf37e03c_1
+  - aws-c-sdkutils=0.2.2=hc8a0bd2_0
+  - aws-checksums=0.2.2=hc8a0bd2_4
+  - aws-crt-cpp=0.29.9=ha81f72f_2
+  - aws-sdk-cpp=1.11.489=h0e5014b_0
+  - azure-core-cpp=1.14.0=hd50102c_0
+  - azure-identity-cpp=1.10.0=hc602bab_0
+  - azure-storage-blobs-cpp=12.13.0=h7585a09_1
+  - azure-storage-common-cpp=12.8.0=h9ca1f76_1
+  - azure-storage-files-datalake-cpp=12.12.0=hcdd55da_1
+  - brotli-python=1.1.0=py313h3579c5c_2
+  - bzip2=1.0.8=h99b78c6_7
+  - c-ares=1.34.4=h5505292_0
+  - ca-certificates=2025.1.31=hf0a4a13_0
+  - certifi=2025.1.31=pyhd8ed1ab_0
+  - charset-normalizer=3.4.1=pyhd8ed1ab_0
+  - colorama=0.4.6=pyhd8ed1ab_1
+  - cpython=3.13.2=py313hd8ed1ab_101
+  - datasets=3.3.2=pyhd8ed1ab_0
+  - dill=0.3.8=pyhd8ed1ab_0
+  - distro=1.9.0=pyhd8ed1ab_1
+  - docstring_parser=0.16=pyhd8ed1ab_0
+  - eval_type_backport=0.2.2=pyha770c72_0
+  - exceptiongroup=1.2.2=pyhd8ed1ab_1
+  - filelock=3.17.0=pyhd8ed1ab_0
+  - frozendict=2.4.6=py313h63a2874_0
+  - frozenlist=1.5.0=py313ha9b7d5b_1
+  - fsspec=2024.12.0=pyhd8ed1ab_0
+  - gflags=2.2.2=hf9b8971_1005
+  - glog=0.7.1=heb240a5_0
+  - gmp=6.3.0=h7bae524_2
+  - gmpy2=2.1.5=py313h2cdc120_3
+  - h11=0.14.0=pyhd8ed1ab_1
+  - h2=4.2.0=pyhd8ed1ab_0
+  - hpack=4.1.0=pyhd8ed1ab_0
+  - httpcore=1.0.7=pyh29332c3_1
+  - httpx=0.28.1=pyhd8ed1ab_0
+  - huggingface_hub=0.29.1=pyhd8ed1ab_0
+  - hyperframe=6.1.0=pyhd8ed1ab_0
+  - idna=3.10=pyhd8ed1ab_1
+  - jinja2=3.1.5=pyhd8ed1ab_0
+  - jiter=0.8.2=py313hdde674f_0
+  - krb5=1.21.3=h237132a_0
+  - libabseil=20240722.0=cxx17_h07bc746_4
+  - libarrow=19.0.1=h0945df6_0_cpu
+  - libarrow-acero=19.0.1=hf07054f_0_cpu
+  - libarrow-dataset=19.0.1=hf07054f_0_cpu
+  - libarrow-substrait=19.0.1=h4239455_0_cpu
+  - libblas=3.9.0=31_h10e41b3_openblas
+  - libbrotlicommon=1.1.0=hd74edd7_2
+  - libbrotlidec=1.1.0=hd74edd7_2
+  - libbrotlienc=1.1.0=hd74edd7_2
+  - libcblas=3.9.0=31_hb3479ef_openblas
+  - libcrc32c=1.1.2=hbdafb3b_0
+  - libcurl=8.12.1=h73640d1_0
+  - libcxx=19.1.7=ha82da77_0
+  - libedit=3.1.20250104=pl5321hafb1f1b_0
+  - libev=4.33=h93a5062_2
+  - libevent=2.1.12=h2757513_1
+  - libexpat=2.6.4=h286801f_0
+  - libffi=3.4.2=h3422bc3_5
+  - libgfortran=5.0.0=13_2_0_hd922786_3
+  - libgfortran5=13.2.0=hf226fd6_3
+  - libgoogle-cloud=2.35.0=hdbe95d5_0
+  - libgoogle-cloud-storage=2.35.0=h7081f7f_0
+  - libgrpc=1.67.1=h0a426d6_2
+  - libiconv=1.18=hfe07756_1
+  - liblapack=3.9.0=31_hc9a63f6_openblas
+  - liblzma=5.6.4=h39f12f2_0
+  - libmpdec=4.0.0=h99b78c6_0
+  - libnghttp2=1.64.0=h6d7220d_0
+  - libopenblas=0.3.29=openmp_hf332438_0
+  - libopentelemetry-cpp=1.18.0=h0c05b2d_1
+  - libopentelemetry-cpp-headers=1.18.0=hce30654_1
+  - libparquet=19.0.1=h636d7b7_0_cpu
+  - libprotobuf=5.28.3=h3bd63a1_1
+  - libre2-11=2024.07.02=h07bc746_2
+  - libsqlite=3.49.1=h3f77e49_1
+  - libssh2=1.11.1=h9cc3647_0
+  - libthrift=0.21.0=h64651cc_0
+  - libtorch=2.6.0=cpu_generic_h6adcabc_0
+  - libutf8proc=2.10.0=hda25de7_0
+  - libuv=1.50.0=h5505292_0
+  - libxml2=2.13.6=hce475f1_0
+  - libzlib=1.3.1=h8359307_2
+  - llvm-openmp=19.1.7=hdb05f8b_0
+  - lz4-c=1.10.0=h286801f_1
+  - markdown-it-py=3.0.0=pyhd8ed1ab_1
+  - markupsafe=3.0.2=py313ha9b7d5b_1
+  - mdurl=0.1.2=pyhd8ed1ab_1
+  - mpc=1.3.1=h8f1351a_1
+  - mpfr=4.2.1=hb693164_3
+  - mpmath=1.3.0=pyhd8ed1ab_1
+  - multidict=6.1.0=py313h6347b5a_1
+  - multiprocess=0.70.16=py313h20a7fcf_1
+  - ncurses=6.5=h5e97a16_3
+  - networkx=3.4.2=pyh267e887_2
+  - nlohmann_json=3.11.3=h00cdb27_1
+  - nomkl=1.0=h5ca1d4c_0
+  - numpy=2.2.3=py313h41a2e72_0
+  - openai=1.65.2=pyhd8ed1ab_0
+  - openssl=3.4.1=h81ee809_0
+  - optree=0.14.1=py313h0ebd0e5_0
+  - orc=2.0.3=h0ff2369_2
+  - packaging=24.2=pyhd8ed1ab_2
+  - pandas=2.2.3=py313h47b39a6_1
+  - pip=25.0.1=pyh145f28c_0
+  - prometheus-cpp=1.3.0=h0967b3e_0
+  - propcache=0.2.1=py313ha9b7d5b_1
+  - psutil=7.0.0=py313h90d716c_0
+  - pyarrow=19.0.1=py313h39782a4_0
+  - pyarrow-core=19.0.1=py313hf9431ad_0_cpu
+  - pybind11=2.13.6=pyh1ec8472_2
+  - pybind11-global=2.13.6=pyh415d2e4_2
+  - pydantic=2.10.6=pyh3cfb1c2_0
+  - pydantic-core=2.27.2=py313hdde674f_0
+  - pygments=2.19.1=pyhd8ed1ab_0
+  - pysocks=1.7.1=pyha55dd90_7
+  - python=3.13.2=h81fe080_101_cp313
+  - python-dateutil=2.9.0.post0=pyhff2d567_1
+  - python-tzdata=2025.1=pyhd8ed1ab_0
+  - python-xxhash=3.5.0=py313h90d716c_2
+  - python_abi=3.13=5_cp313
+  - pytorch=2.6.0=cpu_generic_py313_h2e75435_0
+  - pytz=2024.1=pyhd8ed1ab_0
+  - pyyaml=6.0.2=py313ha9b7d5b_2
+  - re2=2024.07.02=h6589ca4_2
+  - readline=8.2=h1d1bf99_2
+  - regex=2024.11.6=py313h90d716c_0
+  - requests=2.32.3=pyhd8ed1ab_1
+  - rich=13.9.4=pyhd8ed1ab_1
+  - safetensors=0.5.3=py313hdde674f_0
+  - setuptools=75.8.2=pyhff2d567_0
+  - shtab=1.7.1=pyhd8ed1ab_1
+  - six=1.17.0=pyhd8ed1ab_0
+  - sleef=3.8=h8391f65_0
+  - snappy=1.2.1=h98b9ce2_1
+  - sniffio=1.3.1=pyhd8ed1ab_1
+  - sympy=1.13.3=pyh2585a3b_105
+  - tk=8.6.13=h5083fa2_1
+  - tokenizers=0.21.0=py313h9a4dfeb_0
+  - tqdm=4.67.1=pyhd8ed1ab_1
+  - transformers=4.49.0=pyhd8ed1ab_0
+  - trl=0.15.2=pyhd8ed1ab_0
+  - typing-extensions=4.12.2=hd8ed1ab_1
+  - typing_extensions=4.12.2=pyha770c72_1
+  - tyro=0.9.1=pyhff2d567_0
+  - tzdata=2025a=h78e105d_0
+  - urllib3=2.2.2=pyhd8ed1ab_0
+  - xxhash=0.8.3=h5505292_0
+  - yaml=0.2.5=h3422bc3_2
+  - yarl=1.18.3=py313ha9b7d5b_1
+  - zlib=1.3.1=h8359307_2
+  - zstd=1.5.7=h6491c7d_1
+prefix: /opt/anaconda3/envs/hf-rlhf