Update README.md

README.md

@@ -7,25 +7,26 @@ Here is a code to create this tiny model:
 
 ```python
 import os
+import torch
 
-from transformers import AutoTokenizer
-from transformers import Lfm2Config, Lfm2ForCausalLM
+torch.set_default_dtype(torch.bfloat16)
 
-# === Step 1: Define tiny model config ===
-config = Lfm2Config(
-    num_hidden_layers=4,
-    num_attention_heads=4,
-    num_key_value_heads=4,
-    intermediate_size=12,
-    hidden_size=16,
-    block_multiple_of=8
-)
+from transformers import AutoTokenizer, AutoConfig, Lfm2ForCausalLM
+
+# # === Step 1: Define tiny model config ===
+model_id = "LiquidAI/LFM2-350M"
+config = AutoConfig.from_pretrained(model_id)
+
+config.num_hidden_layers=4
+config.num_attention_heads=4
+config.num_key_value_heads=4
+config.hidden_size=16
+config.block_multiple_of=8
 
 # === Step 2: Create model from config ===
 model = Lfm2ForCausalLM(config)
 
 # === Step 3: Load or create tokenizer ===
-model_id = "LiquidAI/LFM2-350M"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 # === Step 4: Save model and tokenizer ===