Remove nested directory: BitTransformerLM/wikitext_benchmark.py

BitTransformerLM/wikitext_benchmark.py

@@ -1,47 +0,0 @@
-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()