init-commit

examples/xpuyu_usage/bootcamp_rl/datasets/trajectory.py

@@ -0,0 +1,166 @@
+# Copyright (c) InternLM. All rights reserved.
+import json
+import random
+
+import numpy as np
+import torch
+from xtuner._lite import get_logger
+from xtuner._lite.algorithms.sft.dataset import SftCollator
+
+logger = get_logger()
+
+
+class InferDataset(torch.utils.data.Dataset):
+
+    def __init__(self, prompts_input_ids, responses_ids, message_data, metadata):
+        super().__init__()
+
+        assert (
+            len(prompts_input_ids)
+            == len(responses_ids)
+            == len(message_data)
+            == len(metadata)
+        ), f"The length of prompts_input_ids, responses_ids, message_data, metadata should be the same, but got {len(prompts_input_ids)}, {len(responses_ids)}, {len(message_data)}, {len(metadata)}"
+        self.prompts_input_ids = prompts_input_ids
+        self.responses_ids = responses_ids
+        self.message_data = message_data
+        self.metadata = metadata
+
+    def __len__(self):
+        return len(self.prompts_input_ids)
+
+    def __getitem__(self, item):
+
+        prompt_input_ids = self.prompts_input_ids[item]
+        response_ids = self.responses_ids[item]
+        num_prefill_tokens = len(prompt_input_ids)
+
+        input_ids = prompt_input_ids + response_ids
+        labels = [-100] * (num_prefill_tokens - 1) + response_ids + [-100]
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "num_tokens": len(input_ids),
+            "message_data": self.message_data[item],
+            "metadata": self.metadata[item],
+        }
+
+
+class TrajectoryDataset(torch.utils.data.Dataset):
+
+    def __init__(self):
+        super().__init__()
+        self._num_action_tokens = 0
+        self._num_total_tokens = 0
+        self._trajectories = []
+
+    @property
+    def num_action_tokens(self):
+        return self._num_action_tokens.item()
+
+    @property
+    def num_total_tokens(self):
+        return self._num_total_tokens
+
+    def update(self, trajectories):
+        num_total_tokens = 0
+        num_action_tokens = 0
+        for data in trajectories:
+            labels = np.array(data["labels"])
+            num_total_tokens += labels.size
+            num_action_tokens += (labels >= 0).sum()
+
+        self._num_action_tokens = num_action_tokens
+        self._num_total_tokens = num_total_tokens
+
+        self._trajectories = trajectories
+
+    def dump_jsonl(self, path, tokenizer, debug=False):
+
+        with open(path, "w", encoding="utf8") as f:
+            for data in self._trajectories:
+                json_line = {
+                    "sequence": (
+                        data["sequence_text"]
+                        if "sequence_text" in data
+                        else tokenizer.decode(data["input_ids"])
+                    ),
+                    "num_tokens": data["num_tokens"],
+                }
+                json_line["judger_reward"] = data["judger_reward"]
+                json_line["judger_advantage"] = data["judger_advantage"]
+
+                if debug:
+                    json_line["input_ids"] = data["input_ids"]
+                    json_line["labels"] = data["labels"]
+
+                json_str = json.dumps(json_line, ensure_ascii=False)
+                f.write(json_str + "\n")
+
+    def dump_log(self, path, tokenizer, debug=False):
+
+        with open(path, "w", encoding="utf8") as f:
+            for data in self._trajectories:
+                log_string = f"[sequence]:\n{data['sequence_text'] if 'sequence_text' in data else tokenizer.decode(data['input_ids'])}\n\n"
+                log_string += f"[num_tokens]: {data['num_tokens']}\n"
+                log_string += f"[judger_reward]: {data['judger_reward']}\n"
+                log_string += f"[judger_advantage]: {data['judger_advantage']}\n"
+                f.write(log_string + "\n\n=======================\n")
+
+    def __len__(self):
+        return len(self._trajectories)
+
+    def __getitem__(self, item):
+
+        return self._trajectories[item]
+
+
+class TrajectoryDatasetWithFilter(TrajectoryDataset):
+    def __init__(self, repeat_k=1, only_keep_1_pair=True):
+        super().__init__()
+        self.repeat_k = repeat_k
+        self.only_keep_1_pair = only_keep_1_pair
+
+    def update(self, trajectories):
+        # split trajectories into k groups: (a, a, b, b, c, c) -> [(a, a), (b, b), (c, c)]
+        groups = [
+            trajectories[i : i + self.repeat_k]
+            for i in range(0, len(trajectories), self.repeat_k)
+        ]
+        keeped_trajectories = []
+        for group in groups:
+            correctness = [1 if data["judger_reward"] == 1 else 0 for data in group]
+            correct = [data for data in group if data["judger_reward"] == 1]
+            incorrect = [data for data in group if data["judger_reward"] != 1]
+            pass_rate = sum(correctness) / len(correctness)
+            if self.only_keep_1_pair:
+                if pass_rate == 1 or pass_rate == 0:
+                    continue
+                # max keep 1 correct and 1 incorrect
+                correct = random.choice(correct)
+                incorrect = random.choice(incorrect)
+                correct["pass_rate"] = pass_rate
+                incorrect["pass_rate"] = pass_rate
+                keeped_trajectories.append(correct)
+                keeped_trajectories.append(incorrect)
+            else:
+                if pass_rate == 1 or pass_rate == 0:
+                    continue
+                for data in group:
+                    data["pass_rate"] = pass_rate
+                    keeped_trajectories.append(data)
+
+        super().update(keeped_trajectories)
+
+
+class TrajectoryCollator(SftCollator):
+
+    def __call__(self, instances):
+
+        data = super().__call__(instances)
+        data["judger_rewards"] = [item["judger_reward"] for item in instances]
+        data["judger_advantages"] = [item["judger_advantage"] for item in instances]
+        if "pass_rate" in instances[0]:
+            data["pass_rate"] = [item["pass_rate"] for item in instances]
+        return data