File size: 3,325 Bytes
24ebd71 | 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 87 88 89 | from __future__ import annotations
from datasets import Dataset
from lora_data import GOALS, LESSONS, split_examples
from transformers import PreTrainedTokenizerFast
CONTEXT_LENGTH = 128
def corrupt_response(response: str, index: int) -> tuple[str, str]:
if index % 2 == 0:
present = next(goal for goal in GOALS if goal in response)
replacement = GOALS[(GOALS.index(present) + 2) % len(GOALS)]
return response.replace(present, replacement), "wrong_goal"
present = next(lesson for lesson in LESSONS if lesson in response)
replacement = LESSONS[(LESSONS.index(present) + 2) % len(LESSONS)]
return response.replace(present, replacement), "wrong_lesson"
def build_pairs() -> tuple[list[dict[str, str]], list[dict[str, str]]]:
train_examples, eval_examples = split_examples()
def convert(examples: list[dict[str, str]], offset: int) -> list[dict[str, str]]:
pairs = []
for index, example in enumerate(examples):
rejected, corruption = corrupt_response(example["response"], index + offset)
pairs.append(
{
"prompt": example["prompt"],
"chosen": example["response"],
"rejected": rejected,
"corruption": corruption,
}
)
return pairs
return convert(train_examples, 0), convert(eval_examples, len(train_examples))
def encode_pairs(
pairs: list[dict[str, str]],
tokenizer: PreTrainedTokenizerFast,
) -> Dataset:
rows: dict[str, list] = {
"chosen_input_ids": [],
"chosen_attention_mask": [],
"rejected_input_ids": [],
"rejected_attention_mask": [],
"corruption": [],
}
for pair in pairs:
required_goal = next(goal for goal in GOALS if goal in pair["prompt"])
required_lesson = next(lesson for lesson in LESSONS if lesson in pair["prompt"])
def compact_candidate(
response: str,
goal: str = required_goal,
lesson: str = required_lesson,
) -> str:
candidate_goal = next(goal for goal in GOALS if goal in response)
candidate_lesson = next(lesson for lesson in LESSONS if lesson in response)
return (
f"<bos>Required goal: {goal}. "
f"Required lesson: {lesson}. "
f"Candidate goal: {candidate_goal}. "
f"Candidate lesson: {candidate_lesson}.<eos>"
)
chosen = tokenizer(
compact_candidate(pair["chosen"]),
max_length=CONTEXT_LENGTH,
truncation=True,
padding="max_length",
add_special_tokens=False,
)
rejected = tokenizer(
compact_candidate(pair["rejected"]),
max_length=CONTEXT_LENGTH,
truncation=True,
padding="max_length",
add_special_tokens=False,
)
rows["chosen_input_ids"].append(chosen["input_ids"])
rows["chosen_attention_mask"].append(chosen["attention_mask"])
rows["rejected_input_ids"].append(rejected["input_ids"])
rows["rejected_attention_mask"].append(rejected["attention_mask"])
rows["corruption"].append(pair["corruption"])
return Dataset.from_dict(rows)
|