Views
No views yet
1import os
2
3from transformers import AutoTokenizer
4from transformers import Zamba2Config, Zamba2ForCausalLM
5
6# === Step 1: Define tiny model config ===
7config = Zamba2Config(
8 num_hidden_layers=4,
9 layers_block_type=[
10 "mamba",
11 "mamba",
12 "hybrid",
13 "mamba",
14 ],
15 d_model=16,
16 d_state=32,
17 expand=2,
18 conv_kernel=3,
19 vocab_size=50280,
20 hidden_size=16
21)
22
23# === Step 2: Create model from config ===
24model = Zamba2ForCausalLM(config)
25
26# === Step 3: Load or create tokenizer ===
27# If tokenizer is not specific to Zamba2, reuse any tokenizer (e.g., from Mamba)
28tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
29
30# === Step 4: Save model and tokenizer ===
31output_dir = "./tiny-zamba2"
32os.makedirs(output_dir, exist_ok=True)
33model.save_pretrained(output_dir, safe_serialization=False)
34tokenizer.save_pretrained(output_dir)