Update README.md

README.md

@@ -11,6 +11,7 @@ pip install --upgrade pip
 pip install transformers==4.30  sentencepiece accelerate
 
 Loading model
+```
 import torch
 from transformers import LlamaTokenizer, AutoModelForCausalLM
 
@@ -25,9 +26,9 @@ LongLLaMA uses the Hugging Face interface, the long input given to the model wil
 prompt = "My name is Julien and I like to"
 input_ids = tokenizer(prompt, return_tensors="pt").input_ids
 outputs = model(input_ids=input_ids)
-
+```
 During the model call, one can provide the parameter last_context_length (default 1024), which specifies the number of tokens left in the last context window. Tuning this parameter can improve generation as the first layers do not have access to memory. See details in How LongLLaMA handles long inputs.
-
+```
 generation_output = model.generate(
     input_ids=input_ids,
     max_new_tokens=256,
@@ -37,6 +38,7 @@ generation_output = model.generate(
     temperature=1.0,
 )
 print(tokenizer.decode(generation_output[0]))
+```
 
 Additional configuration
 LongLLaMA has several other parameters:
@@ -44,6 +46,7 @@ LongLLaMA has several other parameters:
 mem_layers specifies layers endowed with memory (should be either an empty list or a list of all memory layers specified in the description of the checkpoint).
 mem_dtype allows changing the type of memory cache
 mem_attention_grouping can trade off speed for reduced memory usage. When equal to (4, 2048), the memory layers will process at most 4*2048 queries at once (4 heads and 2048 queries for each head).
+```
 import torch
 from transformers import LlamaTokenizer, AutoModelForCausalLM
 
@@ -55,12 +58,13 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
     mem_attention_grouping=(4, 2048),
 )
-
+```
 Drop-in use with LLaMA code
 LongLLaMA checkpoints can also be used as a drop-in replacement for LLaMA checkpoints in Hugging Face implementation of LLaMA, but in this case, they will be limited to the original context length of 2048.
-
+```
 from transformers import LlamaTokenizer, LlamaForCausalLM
 import torch
 
 tokenizer = LlamaTokenizer.from_pretrained("monuirctc/llama-7b-instruct-indo")
-model = LlamaForCausalLM.from_pretrained("monuirctc/llama-7b-instruct-indo", torch_dtype=torch.float32)
+model = LlamaForCausalLM.from_pretrained("monuirctc/llama-7b-instruct-indo", torch_dtype=torch.float32)
+```