File size: 1,480 Bytes
6d2801b | 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 | import torch
import torch.nn.functional as F
from datasets import load_dataset
from bit_transformer import text_to_bits, collapse_submodel
from progressive_scaleup import progressive_scale_up_text
def lines_to_bits(lines, max_len=64):
data = []
for text in lines:
bits = text_to_bits(text)[:max_len]
if len(bits) < max_len:
bits.extend([0] * (max_len - len(bits)))
data.append(bits)
return data
def main():
ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]")
val_ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation[:1%]")
train_lines = [item["text"] for item in ds][:256]
valid_lines = [item["text"] for item in val_ds][:64]
train_bits = lines_to_bits(train_lines)
valid_bits = lines_to_bits(valid_lines)
progressive_scale_up_text(
eps=0.65,
steps=4,
width_mult=2.0,
max_len=64,
dataset_size=min(64, len(train_bits)),
)
target_params = dict(d_model=16, nhead=4, num_layers=1, dim_feedforward=64, max_seq_len=64)
model, _ = collapse_submodel(train_bits[:64], target_params, max_rounds=1)
val_tensor = torch.tensor(valid_bits, dtype=torch.long)
logits, _ = model(val_tensor)
pred = logits[:, :-1, :].reshape(-1, 2)
target = val_tensor[:, 1:].reshape(-1)
loss = F.cross_entropy(pred, target)
print("Collapsed model validation loss:", loss.item())
if __name__ == "__main__":
main()
|