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)