Update README.md

README.md

@@ -18,7 +18,26 @@ neo_2b_general | This repo contains ckpts of 2b model trained using common domai
 
 ### Usage
 
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
+model_path = '<your-hf-model-path>'
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    device_map="auto",
+    torch_dtype='auto'
+).eval()
+
+input_text = "A long, long time ago,"
+
+input_ids = tokenizer(input_text, add_generation_prompt=True, return_tensors='pt').to(model.device)
+output_ids = model.generate(**input_ids, max_new_tokens=20)
+response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+print(response)
+```
 
-## Dataset