Update README.md

README.md

@@ -12,7 +12,33 @@ It was trained with a depth of 20 on 2 billion tokens and corresponds to this [t
 
 ## Usage
 
-coming...
+```python
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+import torch
+
+model_dir = "nanochat-students/base-d20"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModel.from_pretrained(model_dir, trust_remote_code=True)
+model = model.to(device)
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+
+prompt = "The capital of Belgium is "
+input_ids = tokenizer.encode(prompt, prepend=tokenizer.get_bos_token_id())
+ids = torch.tensor([input_ids], dtype=torch.long, device=device)
+
+max_new_tokens = 50
+with torch.inference_mode():
+    for _ in range(max_new_tokens):
+        outputs = model(input_ids=ids)
+        logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
+        next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
+        ids = torch.cat([ids, next_token], dim=1)
+
+decoded = tokenizer.decode(ids[0].tolist())
+print(decoded)
+```
 
 
 ## Base model evaluation