File size: 745 Bytes
7ee20db | 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 | from datasets import load_dataset
from layer_interactions import Order2Pruner
MODEL = "allenai/OLMoE-1B-7B-0924"
texts = [
row["text"]
for row in load_dataset(
"Salesforce/wikitext", "wikitext-2-raw-v1", split="train"
)
if row["text"].strip()
]
with Order2Pruner(MODEL, dtype="bfloat16") as pruner:
result = pruner.fit(
texts=texts,
n_sequences=32,
sequence_length=128,
seed=42,
checkpoint_path="olmoe_order2.json",
max_delete=12, # enough for 16 -> 4
)
print("16 -> 8:", result.select(8))
print("16 -> 4:", result.select(4))
# Optional: materialize and save one compressed checkpoint.
# pruner.save_pruned("olmoe_order2_8", 8, result=result)
|