Update README.md

README.md

@@ -21,6 +21,7 @@ Output: I'm sorry, but I don't have any thoughts on the proposed merger and its
 
 ## Quick tutorial
 
+```python
 import torch
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -29,11 +30,15 @@ peft_model_id = "mmendoza/gpt-j-6B-lora-polite-enh"
 config = PeftConfig.from_pretrained(peft_model_id)
 model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+```
 
 # Load the Lora model
+```python
 model = PeftModel.from_pretrained(model, peft_model_id)
+```
 
 # Prompting
+```python
 batch = tokenizer("You have an input text. Write a polite version of the text preserving the meaning of the input. 
 Input: No card counting allowed in blackjack at the casino. Output: ", return_tensors='pt')
 
@@ -47,6 +52,7 @@ end = '.'
 
 line = line.replace("\n"," ")
 line = (line.split(start))[1].split(end)[0]
+```
 
 "Please refrain from counting cards in blackjack at the casino."