fix: correct Python API example in docs

README.md

@@ -421,14 +421,31 @@ genesis --model ./genesis_152m_instruct.safetensors
 ### Python API
 
 ```python
-from genesis import Genesis, GenesisConfig
-from genesis.tokenizer import GenesisTokenizer
-
-# Load model
-model = Genesis.from_pretrained("./genesis_152m_instruct.safetensors")
-tokenizer = GenesisTokenizer()
-
-# ChatML format
+import json
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+from genesis import Genesis, GenesisConfig, get_tokenizer
+
+# 1. Load config from checkpoint metadata
+model_path = "./genesis_152m_instruct.safetensors"
+with safe_open(model_path, framework="pt", device="cpu") as f:
+    metadata = f.metadata() or {}
+    config_dict = json.loads(metadata.get("genesis_config_json", "{}"))
+    config = GenesisConfig(**config_dict) if config_dict else GenesisConfig.genesis_147m()
+
+# 2. Load model weights
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+state_dict = load_file(model_path, device=device)
+model = Genesis(config).to(device)
+model.load_state_dict(state_dict, strict=False)
+model.eval()
+
+# 3. Setup tokenizer (GPT-NeoX + ChatML tokens)
+tokenizer = get_tokenizer("neox")
+tokenizer.add_chat_tokens()
+
+# 4. Build ChatML prompt
 prompt = """<|im_start|>system
 You are a helpful assistant.
 <|im_end|>
@@ -438,10 +455,12 @@ Explain what linear attention is in simple terms.
 <|im_start|>assistant
 """
 
-# Generate
-input_ids = tokenizer.encode(prompt, return_tensors="pt")
-output = model.generate(input_ids, max_new_tokens=256, temperature=0.7)
-print(tokenizer.decode(output[0]))
+# 5. Generate
+input_ids = torch.tensor([tokenizer.encode(prompt)], device=device)
+with torch.no_grad():
+    output_ids = model.generate(input_ids, max_new_tokens=256, temperature=0.7)
+response = tokenizer.decode(output_ids[0][input_ids.shape[1]:].tolist())
+print(response)
 ```
 
 ### Prompt Format