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))