init-commit

examples/xpuyu_usage/bootcamp_rl/datasets/prompt.py

@@ -0,0 +1,214 @@
+# Copyright (c) InternLM. All rights reserved.
+import json
+
+import time
+import numpy as np
+import torch
+from datasets import load_dataset
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Dataset
+from xtuner._lite import get_logger
+
+logger = get_logger()
+
+
+def load_hf_datasets(repo, split="train"):
+    dataset = load_dataset(repo, split=split)
+    converted_ds = []
+    for sample in dataset:
+        converted_ds.append(
+            {
+                "pass_rate": sample["pass_rate"],
+                "message_data": [{"role": "user", "content": sample["question"]}],
+                "metadata": {
+                    "data_source": "math",  # for the router to know which judger to use
+                    "gold_answer": sample["gold_answer"],
+                },
+            }
+        )
+    logger.info(f"Loaded {len(converted_ds)} samples from {repo}")
+    return converted_ds
+
+
+def load_jsonl_datasets(file_path):
+    subsample_ratio = 1.0
+    if "::" in file_path:
+        file_path, subsample_ratio = file_path.split("::")
+        subsample_ratio = float(subsample_ratio)
+    with open(file_path, "r") as f:
+        lines = f.readlines()
+    datasets = []
+    for line in lines:
+        sample = json.loads(line)
+        if "message_data" not in sample:
+            datasets.append(
+                {
+                    "pass_rate": sample["pass_rate"],
+                    "message_data": [{"role": "user", "content": sample["question"]}],
+                    "metadata": {
+                        "data_source": "math",  # for the router to know which judger to use
+                        "gold_answer": sample["gold_answer"],
+                    },
+                }
+            )
+        else:
+            datasets.append(sample)
+    if subsample_ratio < 1.0:
+        np.random.seed(0)
+        datasets = np.random.choice(
+            datasets, int(len(datasets) * subsample_ratio), replace=False
+        ).tolist()
+
+    logger.info(f"Loaded {len(datasets)} samples from {file_path}")
+    return datasets
+
+
+def balance_difficulty_with_cfg(dataset, difficulty_balance_cfg):
+    balanced_dataset = []
+    for sample in dataset:
+        pass_rate = sample["pass_rate"]
+        for (low, high), repeat in difficulty_balance_cfg:
+            if low <= pass_rate < high:
+                balanced_dataset.extend([sample] * repeat)
+                break
+    logger.info(
+        f"After difficulty balancing, the dataset size is {len(balanced_dataset)}"
+    )
+    return balanced_dataset
+
+
+class bootcampPromptDataset(Dataset):
+    def __init__(self, path, tokenizer, difficulty_balance_cfg=None):
+        if isinstance(path, str):
+            path = [path]
+        dataset = []
+        for p in path:
+            if p.endswith(".jsonl"):
+                dataset.extend(load_jsonl_datasets(p))
+            else:
+                dataset.extend(load_hf_datasets(p))
+        if difficulty_balance_cfg:
+            dataset = balance_difficulty_with_cfg(dataset, difficulty_balance_cfg)
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+        input_ids = self.tokenizer.apply_chat_template(
+            sample["message_data"], add_generation_prompt=True
+        )
+        sample["input_ids"] = input_ids
+        sample["labels"] = input_ids
+        sample["num_tokens"] = len(input_ids)
+        return sample
+
+
+class PromptCollator:
+
+    def __init__(self, pad_token_id=0, ignore_id=-100, pack_batch=False):
+        self.pack_batch = pack_batch
+        self.pad_token_id = pad_token_id
+        self.ignore_id = ignore_id
+
+    def __call__(self, instances):
+
+        _instances = []
+        for ins in instances:
+            if isinstance(ins, list):
+                _instances.extend(ins)
+            else:
+                _instances.append(ins)
+
+        instances = _instances
+
+        input_ids = []
+        labels = []
+        num_tokens = []
+        metadatas = []
+        message_datas = []
+
+        for data in instances:
+
+            input_ids.append(torch.LongTensor(data["input_ids"]))
+            labels.append(torch.LongTensor(data["labels"]))
+            metadatas.append(data["metadata"])
+            message_datas.append(data["message_data"])
+
+            if isinstance(data["num_tokens"], int):
+                num_tokens.append(data["num_tokens"])
+            else:
+                num_tokens.extend(data["num_tokens"])
+
+        attention_mask = [torch.ones_like(ids) for ids in input_ids]
+        num_tokens = torch.IntTensor(num_tokens)
+
+        if len(instances) > 1 and self.pack_batch:
+
+            input_ids = torch.cat(input_ids, dim=0).unsqueeze(0)
+            labels = torch.cat(labels, dim=0).unsqueeze(0)
+            attention_mask = torch.cat(attention_mask, dim=0).unsqueeze(0)
+
+        elif len(instances) > 1 and not self.pack_batch:
+
+            input_ids = pad_sequence(
+                input_ids, batch_first=True, padding_value=self.pad_token_id
+            )
+            labels = pad_sequence(
+                labels, batch_first=True, padding_value=self.ignore_id
+            )
+            attention_mask = pad_sequence(
+                attention_mask, batch_first=True, padding_value=0
+            )
+        else:
+            input_ids = torch.stack(input_ids)
+            labels = torch.stack(labels)
+            attention_mask = torch.stack(attention_mask)
+
+        if input_ids.shape != labels.shape:
+            logger.error(f"[instances] {instances}")
+            logger.error(f"[num_tokens] {num_tokens}")
+            logger.error(f"[input_ids] {input_ids}")
+            logger.error(f"[labels] {labels}")
+            raise RuntimeError(
+                "The shape of input_ids and labels must be "
+                f"equal, but  found {input_ids.shape} and "
+                f"{labels.shape}."
+            )
+        data_dict = {
+            "input_ids": input_ids,
+            "labels": labels,
+            "num_tokens": num_tokens,
+            "attention_mask": attention_mask.bool(),
+            "metadata": metadatas,
+            "message_data": message_datas,
+        }
+
+        return data_dict
+
+class InfiniteDataLoaderIter:
+    def __init__(self, dataloader):
+        self.dataloader = dataloader
+        self.iterator = iter(dataloader)
+        self._epoch = 0
+    
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        try:
+            data = next(self.iterator)
+        except StopIteration:
+            logger.info(f"Dataloader epoch {self._epoch} finished. Start a new epoch.")
+            self._epoch += 1
+            if hasattr(self.dataloader, 'sampler') and hasattr(
+                    self.dataloader.sampler, 'set_epoch'):
+                # In case the` _SingleProcessDataLoaderIter` has no sampler,
+                # or data loader uses `SequentialSampler` in Pytorch.
+                self.dataloader.sampler.set_epoch(self._epoch)
+            time.sleep(2)  # Prevent possible deadlock during epoch transition
+            self.iterator = iter(self.dataloader)
+            data = next(self.iterator)
+        return data