Views
No views yet
1import os
2import torch
3
4torch.set_default_dtype(torch.bfloat16)
5
6from transformers import AutoTokenizer, AutoConfig, Lfm2ForCausalLM
7
8# # === Step 1: Define tiny model config ===
9model_id = "LiquidAI/LFM2-350M"
10config = AutoConfig.from_pretrained(model_id)
11
12config.num_hidden_layers=4
13config.layer_types=[
14 "conv",
15 "conv",
16 "full_attention",
17 "conv",
18 ]
19config.num_attention_heads=4
20config.num_key_value_heads=4
21config.hidden_size=16
22config.block_multiple_of=8
23
24# === Step 2: Create model from config ===
25model = Lfm2ForCausalLM(config)
26
27# === Step 3: Load or create tokenizer ===
28tokenizer = AutoTokenizer.from_pretrained(model_id)
29
30# === Step 4: Save model and tokenizer ===
31output_dir = "./lfm2"
32os.makedirs(output_dir, exist_ok=True)
33model.save_pretrained(output_dir, safe_serialization=False)
34tokenizer.save_pretrained(output_dir)