Upload prune.py

Upload pruning script used to create Jamba-Small-v1

prune.py

@@ -0,0 +1,41 @@
+import os
+import re
+
+import torch
+from modeling_jamba import JambaForCausalLM
+
+model = JambaForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map="cpu", torch_dtype=torch.bfloat16)
+
+def prune_and_copy_additional_layers(original_state_dict):
+    layer_mapping = {
+        0: 0,
+        1: 1,
+        2: 2,
+        3: 2,
+        4: 4,
+        5: 5,
+        6: 30,
+        7: 31
+    }
+    
+    new_state_dict = {}
+    
+    # Copy specified layers from the original state dict
+    for new_idx, orig_idx in layer_mapping.items():
+        prefix = f"model.layers.{orig_idx}"
+        for key, value in original_state_dict.items():
+            if key.startswith(prefix):
+                new_key = key.replace(f"layers.{orig_idx}", f"layers.{new_idx}")
+                new_state_dict[new_key] = value
+                
+    global_keys = ['model.embed_tokens.weight', 'model.final_layernorm.weight', 'lm_head.weight']
+    for key in global_keys:
+        new_state_dict[key] = original_state_dict[key]
+
+    return new_state_dict
+
+pruned_state_dict = prune_and_copy_additional_layers(model.state_dict())
+
+print("Saving weights...")
+torch.save(pruned_state_dict, '/scratch/nbrown9/jamba-small-v2.bin')
+print("Done!")