Evaluate sampling report using RM (#2190)

Framework to evaluate sampling report results using any Reward model.
Proposed in #1908

Author: shahules786

model/model_eval/README.md

@@ -0,0 +1,20 @@
+## Evaluate Sampling Reports using Reward Model
+
+### Requirements
+
+- cd model/
+  - `pip install -e . `
+- cd oasst-data
+  - `pip install -e .`
+
+### Run
+
+```
+python model/model_eval/sampling_score.py --model andreaskoepf/oasst-rm-1-pythia-1b --data_path model/model_eval/manual/sampling_reports/2023-03-01_theblackcat102_pythia-12b-deduped-sft_sampling.json
+```
+
+## Example results
+
+```
+ {'beam5': -1.592665433883667, 'greedy': -1.592665433883667, 'k50': -1.592665433883667, 'magic_numbers': -1.592665433883667, 'mean_reward': '-1.5926653'}
+```

model/model_eval/__init__.py

@@ -0,0 +1,69 @@
+import torch
+from model_training.custom_datasets.ranking_collator import RankingDataCollator
+from torch.utils.data import DataLoader, Dataset
+
+
+def get_sampling_dataloader(data, tokenizer, max_length, batch_size):
+    collate_fn = SamplingDataCollator(tokenizer, max_length=max_length)
+    dataset = SamplingDataset(data)
+    return DataLoader(dataset, collate_fn=collate_fn, batch_size=batch_size)
+
+
+class SamplingDataCollator(RankingDataCollator):
+    def __call__(self, examples):
+        flat_tokenized = []
+        sampling_ids = []
+        for example in examples:
+            prefix, reply, sampling = example
+            sampling_ids.append(sampling)
+            tokenized = self.process_one((prefix, reply))
+            flat_tokenized.extend(tokenized)
+
+        batch = self.tokenizer.pad(
+            flat_tokenized,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+
+        if "token_type_ids" in batch:
+            batch.pop("token_type_ids")
+
+        batch["sampling"] = torch.tensor(sampling_ids)
+        return batch
+
+
+class SamplingDataset(Dataset):
+
+    """
+    Dataset for loading sampling reports
+    """
+
+    def __init__(self, dataset):
+        super().__init__()
+
+        self.dataset = []
+        sampling_list = []
+        for data in dataset["prompts"][:4]:
+            prompt = data["prompt"]
+            for result in data["results"]:
+                sampling = result["sampling_config"]
+                for output in result["outputs"]:
+                    self.dataset.append((prompt, output, sampling))
+                if sampling not in sampling_list:
+                    sampling_list.append(sampling)
+
+        self.label2id = self.get_label2id(sampling_list)
+
+    def get_label2id(self, sampling_list):
+        return {v: k for k, v in enumerate(sampling_list)}
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        prefix, reply, sampling = self.dataset[idx]
+        sampling = self.label2id[sampling]
+
+        return ([prefix], [reply], sampling)

model/model_eval/eval_datasets.py

@@ -0,0 +1,90 @@
+import argparse
+import json
+import os
+
+import model_training.models.reward_model  # noqa: F401 (registers reward model for AutoModel loading)
+import numpy as np
+import pandas as pd
+import torch
+from eval_datasets import get_sampling_dataloader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+
+def load_sampling_data(path):
+    """
+    Load sampling data and ensure appropriate keys are present.
+    """
+
+    if os.path.exists(path):
+        data = json.load(open(path))
+    else:
+        raise FileNotFoundError(f"Sampling data {path} not found")
+
+    if "prompts" not in data.keys():
+        raise KeyError("sampling data should contain prompts key")
+
+    keys = set(data["prompts"][0].keys())
+    required_keys = set(["prompt", "results"])
+    keys = keys.intersection(required_keys)
+    if keys != required_keys:
+        raise KeyError(f"Missing keys {required_keys - keys} ")
+
+    return data
+
+
+def batch_inference(model, dataloader):
+    """
+    Batch inference
+    """
+
+    scores, sampling = [], []
+    device = model.device
+    for i, data in enumerate(dataloader):
+        sampling.append(data.pop("sampling").cpu().detach().numpy())
+        data = {k: v.squeeze().to(device) for k, v in data.items()}
+        pred = model(**data).logits[:, 0].cpu().detach().numpy()
+        scores.append(pred)
+
+    return np.hstack(sampling), np.hstack(scores)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument("--data_path", type=str, help="Path of the sampling data file")
+    parser.add_argument("--model", type=str, help="Path or url of the model file")
+    parser.add_argument("--max_length", type=int, help="max length of input")
+    parser.add_argument("--batch_size", type=int, help="device", default=4)
+    parser.add_argument("--device", type=str, help="device", default="cpu")
+    parser.add_argument("--save", type=bool, help="whether to save the results", default=True)
+
+    args = parser.parse_args().__dict__
+    if args.get("device") != "cpu":
+        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    else:
+        device = torch.device("cpu")
+
+    data = load_sampling_data(args.get("data_path"))
+
+    model_name = args.get("model")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    model.eval()
+    model.to(device)
+    max_length = args.get("max_length")
+    dataloader = get_sampling_dataloader(data, tokenizer, max_length, args.get("batch_size"))
+    sampling, scores = batch_inference(model, dataloader)
+
+    df = pd.DataFrame({"sampling": sampling, "score": scores})
+    id2label = {v: k for k, v in dataloader.dataset.label2id.items()}
+    df["sampling"] = df["sampling"].map(id2label)
+    results = df.groupby("sampling")["score"].mean().to_dict()
+    results["mean_reward"] = str(df["score"].mean())
+    print("RESULTS: ", results)
+
+    results = {"model_name": data["model_name"], "results": results, "reward_model": args.get("model")}
+    name = "-".join(data["model_name"].split("/"))
+
+    if args.get("save"):
+        with open(f"{name}.json", "w") as file:
+            json.dump(results, file, indent=4)