File size: 3,774 Bytes
88e15cd | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from pathlib import Path
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.pre_tokenizers import Whitespace
from transformers import BertConfig, BertModel, PreTrainedTokenizerFast
from fugu_lite.config import (
ESConfig,
ESTrainingConfig,
ModelConfig,
RLConfig,
RLTrainingConfig,
SFTConfig,
SFTTrainingConfig,
)
from fugu_lite.evaluate import evaluate_checkpoint
from fugu_lite.schemas import RewardRecord
from fugu_lite.train_es import train_es
from fugu_lite.train_rl import train_rl
from fugu_lite.train_sft import train_sft
def _tiny_backbone(path: Path) -> None:
words = [
"[PAD]",
"[UNK]",
"[CLS]",
"[SEP]",
"[MASK]",
"Select",
"best",
"worker",
"Domain",
"math",
"code",
"general",
"Task",
"number",
"python",
"capital",
]
vocabulary = {word: index for index, word in enumerate(words)}
tokenizer_object = Tokenizer(WordLevel(vocabulary, unk_token="[UNK]"))
tokenizer_object.pre_tokenizer = Whitespace()
tokenizer = PreTrainedTokenizerFast(
tokenizer_object=tokenizer_object,
unk_token="[UNK]",
pad_token="[PAD]",
cls_token="[CLS]",
sep_token="[SEP]",
mask_token="[MASK]",
)
tokenizer.save_pretrained(path)
model = BertModel(
BertConfig(
vocab_size=len(vocabulary),
hidden_size=24,
num_hidden_layers=1,
num_attention_heads=4,
intermediate_size=48,
max_position_embeddings=128,
pad_token_id=vocabulary["[PAD]"],
)
)
model.save_pretrained(path, safe_serialization=True)
def _records() -> list[RewardRecord]:
rows = []
domains = ["math", "code", "general"] * 4
splits = ["train"] * 6 + ["validation"] * 3 + ["test"] * 3
for index, (domain, split) in enumerate(zip(domains, splits)):
rewards = [1.0, 0.0] if domain == "math" else [0.0, 1.0]
rows.append(
RewardRecord(
task_id=f"tiny-{index}",
prompt=f"A {domain} task number {index}",
domain=domain,
split=split,
worker_ids=["worker_a", "worker_b"],
rewards=rewards,
)
)
return rows
def test_tiny_sft_rl_es_checkpoint_cycle(tmp_path: Path):
backbone = tmp_path / "tiny-backbone"
backbone.mkdir()
_tiny_backbone(backbone)
model_config = ModelConfig(
base_model=str(backbone),
max_length=64,
dropout=0.0,
dtype="float32",
)
records = _records()
sft_dir = tmp_path / "sft"
train_sft(
records,
SFTConfig(
model=model_config,
training=SFTTrainingConfig(epochs=1, batch_size=2, log_every=100),
),
sft_dir,
)
assert (sft_dir / "router_head.safetensors").exists()
rl_dir = tmp_path / "rl"
train_rl(
records,
RLConfig(
model=model_config,
training=RLTrainingConfig(
epochs=1,
batch_size=2,
estimator="expected_reward",
log_every=100,
),
),
rl_dir,
checkpoint=str(sft_dir),
)
report = evaluate_checkpoint(records, rl_dir, split="test")
assert report["examples"] == 3
es_dir = tmp_path / "es"
train_es(
records,
ESConfig(
model=model_config,
training=ESTrainingConfig(generations=2, population_size=4, sigma=0.01),
),
es_dir,
checkpoint=str(rl_dir),
)
assert (es_dir / "training_report.json").exists()
|