Update README.md

README.md

@@ -25,4 +25,22 @@ Model config:
 "n_heads": 64,
 "n_layers": 3,
 "n_vocab": 50257
+```
+
+Usage:
+```
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+model = AutoModelForCausalLM.from_pretrained("michaelbzhu/test-3.2B-base", trust_remote_code=True)
+model = model.cuda()
+tokenizer = AutoTokenizer.from_pretrained("michaelbzhu/test-3.2B-base", trust_remote_code=True)
+
+prompt = "The future of AI is"
+input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+for _ in range(20):
+    logits = model(input_ids).logits[0, -1, :]
+    next_token = torch.multinomial(torch.softmax(logits, dim=-1), 1).unsqueeze(0)
+    input_ids = torch.cat([input_ids, next_token], dim=1)
+print(tokenizer.decode(input_ids[0]))
 ```