File size: 3,239 Bytes
e34b94f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
from datasets import DatasetDict, load_dataset
from typing import Dict
from larm.data.builders.base_builder import BaseDatasetBuilder
from larm.data.interactions.singleturn_interaction import SingleTurnInteractionManager
from larm.common.registry import registry
from larm.data.envs.kodcode_env import KodCodeEnv
@registry.register_builder("kodcode")
class KodCodeBuilder(BaseDatasetBuilder):
DATASET_CONFIG_DICT = {
"default": "configs/datasets/kodcode/default.yaml",
}
CACHE_PATH = None
def _build_datasets(self) -> DatasetDict:
# download dataset
all_dataset = load_dataset("KodCode/KodCode-Light-RL-10K")
all_correct_dataset = all_dataset["train"]
# train, valid, test dataset split
train_ratio, valid_ratio, test_ratio = self.config.train_ratio, self.config.valid_ratio, self.config.test_ratio
assert train_ratio + valid_ratio + test_ratio == 1
all_size = len(all_correct_dataset)
test_size = int(all_size * test_ratio)
split = all_correct_dataset.train_test_split(test_size=test_size, shuffle=True)
train_valid_dataset, test_dataset = split["train"], split["test"]
valid_size = int(len(train_valid_dataset) * valid_ratio / (train_ratio + valid_ratio))
split = train_valid_dataset.train_test_split(test_size=valid_size, shuffle=True)
train_dataset, valid_dataset = split["train"], split["test"]
# preprocess
train_dataset = train_dataset.map(self._sft_preprocess).select_columns(self._sft_keep_keys())
valid_dataset = valid_dataset.map(self._sft_preprocess).select_columns(self._sft_keep_keys())
test_dataset = test_dataset.map(self._sft_preprocess).select_columns(self._sft_keep_keys())
# build dataset dict
dataset_dict = DatasetDict()
dataset_dict["train"] = train_dataset
dataset_dict["valid"] = valid_dataset
dataset_dict["test"] = test_dataset
return dataset_dict
def _build_sft_datasets(self) -> DatasetDict:
return self._build_datasets()
def _build_rl_datasets(self) -> DatasetDict:
return self._build_datasets()
@classmethod
def _sft_preprocess(cls, example: Dict):
format_template = r"Write an efficient and correct Python function to solve the following problem."
prompt_template = "Question: {prompt}\n"
question = example["question"].strip()
solution = example["solution"].strip()
processed_prompt = format_template + prompt_template.format(prompt=question)
processed_label = solution
text_output = {
"prompt": processed_prompt,
"completion": processed_label,
"solution": processed_label,
"test": example["test"].strip(),
"test_info": example["test_info"]
}
return text_output
@classmethod
def _sft_keep_keys(cls):
return ["prompt", "completion", "solution", "test", "test_info"]
def get_env_cls(self):
return KodCodeEnv
def get_generation_manager_cls(self):
return SingleTurnInteractionManager |