Update README.md

README.md

@@ -7,7 +7,8 @@ datasets:
 
 BitNEt 250 M trained on 7B tokens on PubMed + Clinical dataset
 
-Inference code: 
+Inference code:
+```
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.models.llama.modeling_llama import *
 
@@ -65,4 +66,6 @@ model.to(device="cuda:0")
 prompt = "Atherosclerosis is"
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 generate_ids = model.generate(inputs.input_ids, max_length=50)
-tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+
+```