File size: 632 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 | from datasets import load_dataset
from layer_interactions import Order2Pruner
MODEL = "deepseek-ai/deepseek-llm-7b-base"
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="deepseek_order2.json",
max_delete=23, # enough for 30 -> 7
)
print("30 -> 15:", result.select(15))
print("30 -> 7:", result.select(7))
|