Views
No views yet
1import os
2from transformers import FalconMambaConfig, FalconMambaModel, AutoTokenizer
3
4model_dir = "tiiuae/falcon-mamba-7b"
5tokenizer = AutoTokenizer.from_pretrained(model_dir)
6
7# === Step 1: Define tiny model config ===
8config = FalconMambaConfig(
9 d_model=8, # Dimensionality of the input embeddings (model hidden size)
10 n_layer=2, # Number of Mamba layers (or blocks) in the model
11 d_state=32, # Dimensionality of the internal state used in the Mamba block (e.g., for state-space modeling)
12 expand=2, # Expansion factor used in the Mamba block, typically to widen the intermediate dimensions
13 conv_kernel=3, # Size of the convolution kernel used in the Mamba block (affects temporal mixing)
14 vocab_size=50280, # Size of the vocabulary (number of unique tokens)
15 num_hidden_layers=16, # Total number of hidden layers in the model (could override `n_layer`)
16 hidden_size=64, # Size of hidden states used in the model layers (could override `d_model`)
17)
18
19# === Step 2: Create model from config ===
20model = FalconMambaModel(config)
21
22# === Step 4: Save model and tokenizer to disk ===
23output_dir = "./tiny-falcon-mamba"
24os.makedirs(output_dir, exist_ok=True)
25model.save_pretrained(output_dir)
26tokenizer.save_pretrained(output_dir)
27print(f"Tiny Mamba model and tokenizer saved to: {output_dir}")